Improve DiscoveryFlowHandler typing (#66511)

homeassistant/components/rpi_power/config_flow.py

@@ -1,6 +1,7 @@
 """Config flow for Raspberry Pi Power Supply Checker."""
 from __future__ import annotations
 
+from collections.abc import Awaitable
 from typing import Any
 
 from rpi_bad_power import new_under_voltage
@@ -18,7 +19,7 @@ async def _async_supported(hass: HomeAssistant) -> bool:
     return under_voltage is not None
 
 
-class RPiPowerFlow(DiscoveryFlowHandler, domain=DOMAIN):
+class RPiPowerFlow(DiscoveryFlowHandler[Awaitable[bool]], domain=DOMAIN):
     """Discovery flow handler."""
 
     VERSION = 1
@@ -35,7 +36,7 @@ class RPiPowerFlow(DiscoveryFlowHandler, domain=DOMAIN):
         self, data: dict[str, Any] | None = None
     ) -> FlowResult:
         """Handle a flow initialized by onboarding."""
-        has_devices = await self._discovery_function(self.hass)  # type: ignore
+        has_devices = await self._discovery_function(self.hass)
 
         if not has_devices:
             return self.async_abort(reason="no_devices_found")

homeassistant/components/sonos/config_flow.py

@@ -1,4 +1,5 @@
 """Config flow for SONOS."""
+from collections.abc import Awaitable
 import dataclasses
 
 from homeassistant import config_entries
@@ -16,7 +17,7 @@ async def _async_has_devices(hass: HomeAssistant) -> bool:
     return bool(await ssdp.async_get_discovery_info_by_st(hass, UPNP_ST))
 
 
-class SonosDiscoveryFlowHandler(DiscoveryFlowHandler):
+class SonosDiscoveryFlowHandler(DiscoveryFlowHandler[Awaitable[bool]]):
     """Sonos discovery flow that callsback zeroconf updates."""
 
     def __init__(self) -> None:

homeassistant/helpers/config_entry_flow.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 from collections.abc import Awaitable, Callable
 import logging
-from typing import TYPE_CHECKING, Any, Union, cast
+from typing import TYPE_CHECKING, Any, Generic, TypeVar, Union, cast
 
 from homeassistant import config_entries
 from homeassistant.components import dhcp, mqtt, ssdp, zeroconf
@@ -15,12 +15,13 @@ from .typing import UNDEFINED, DiscoveryInfoType, UndefinedType
 if TYPE_CHECKING:
     import asyncio
 
-DiscoveryFunctionType = Callable[[HomeAssistant], Union[Awaitable[bool], bool]]
+_R = TypeVar("_R", bound="Awaitable[bool] | bool")
+DiscoveryFunctionType = Callable[[HomeAssistant], _R]
 
 _LOGGER = logging.getLogger(__name__)
 
 
-class DiscoveryFlowHandler(config_entries.ConfigFlow):
+class DiscoveryFlowHandler(config_entries.ConfigFlow, Generic[_R]):
     """Handle a discovery config flow."""
 
     VERSION = 1
@@ -29,7 +30,7 @@ class DiscoveryFlowHandler(config_entries.ConfigFlow):
         self,
         domain: str,
         title: str,
-        discovery_function: DiscoveryFunctionType,
+        discovery_function: DiscoveryFunctionType[_R],
     ) -> None:
         """Initialize the discovery config flow."""
         self._domain = domain
@@ -153,7 +154,7 @@ class DiscoveryFlowHandler(config_entries.ConfigFlow):
 def register_discovery_flow(
     domain: str,
     title: str,
-    discovery_function: DiscoveryFunctionType,
+    discovery_function: DiscoveryFunctionType[Awaitable[bool] | bool],
     connection_class: str | UndefinedType = UNDEFINED,
 ) -> None:
     """Register flow for discovered integrations that not require auth."""
@@ -172,7 +173,7 @@ def register_discovery_flow(
             domain,
         )
 
-    class DiscoveryFlow(DiscoveryFlowHandler):
+    class DiscoveryFlow(DiscoveryFlowHandler[Union[Awaitable[bool], bool]]):
         """Discovery flow handler."""
 
         def __init__(self) -> None: