From 033a16b67e7803b3d12b3d79ad80b3476bed0d32 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Thu, 15 Dec 2022 08:45:54 +0100
Subject: [PATCH] Improve data entry flow typing (#83901)

---
 homeassistant/data_entry_flow.py | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/homeassistant/data_entry_flow.py b/homeassistant/data_entry_flow.py
index d4c8a0db4e3..89000d8e0d5 100644
--- a/homeassistant/data_entry_flow.py
+++ b/homeassistant/data_entry_flow.py
@@ -112,16 +112,19 @@ def _async_flow_handler_to_flow_result(
     flows: Iterable[FlowHandler], include_uninitialized: bool
 ) -> list[FlowResult]:
     """Convert a list of FlowHandler to a partial FlowResult that can be serialized."""
-    return [
-        FlowResult(
+    results = []
+    for flow in flows:
+        if not include_uninitialized and flow.cur_step is None:
+            continue
+        result = FlowResult(
             flow_id=flow.flow_id,
             handler=flow.handler,
             context=flow.context,
-            step_id=flow.cur_step["step_id"] if flow.cur_step else None,
         )
-        for flow in flows
-        if include_uninitialized or flow.cur_step is not None
-    ]
+        if flow.cur_step:
+            result["step_id"] = flow.cur_step["step_id"]
+        results.append(result)
+    return results
 
 
 class FlowManager(abc.ABC):
@@ -269,8 +272,10 @@ class FlowManager(abc.ABC):
         cur_step = flow.cur_step
         assert cur_step is not None
 
-        if cur_step.get("data_schema") is not None and user_input is not None:
-            user_input = cur_step["data_schema"](user_input)
+        if (
+            data_schema := cur_step.get("data_schema")
+        ) is not None and user_input is not None:
+            user_input = data_schema(user_input)
 
         # Handle a menu navigation choice
         if cur_step["type"] == FlowResultType.MENU and user_input:
@@ -348,7 +353,7 @@ class FlowManager(abc.ABC):
 
     async def _async_handle_step(
         self,
-        flow: Any,
+        flow: FlowHandler,
         step_id: str,
         user_input: dict | BaseServiceInfo | None,
         step_done: asyncio.Future | None = None,
@@ -415,7 +420,7 @@ class FlowHandler:
     """Handle the configuration flow of a component."""
 
     # Set by flow manager
-    cur_step: dict[str, Any] | None = None
+    cur_step: FlowResult | None = None
 
     # While not purely typed, it makes typehinting more useful for us
     # and removes the need for constant None checks or asserts.