Improve data entry flow typing (#83901)

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.