Improve DiscoveryFlowHandler typing (#66511)

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: