Improve entity_platform helper typing (#75464)

* Improve entity_platform helper typing

* Add protocol class

* Apply suggestions from code review

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

.strict-typing

@@ -20,6 +20,7 @@ homeassistant.helpers.deprecation
 homeassistant.helpers.discovery
 homeassistant.helpers.dispatcher
 homeassistant.helpers.entity
+homeassistant.helpers.entity_platform
 homeassistant.helpers.entity_values
 homeassistant.helpers.event
 homeassistant.helpers.reload

homeassistant/helpers/entity_platform.py

@@ -2,11 +2,10 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable, Coroutine, Iterable
+from collections.abc import Awaitable, Callable, Coroutine, Iterable
 from contextvars import ContextVar
 from datetime import datetime, timedelta
 from logging import Logger, getLogger
-from types import ModuleType
 from typing import TYPE_CHECKING, Any, Protocol
 from urllib.parse import urlparse
 
@@ -71,6 +70,36 @@ class AddEntitiesCallback(Protocol):
         """Define add_entities type."""
 
 
+class EntityPlatformModule(Protocol):
+    """Protocol type for entity platform modules."""
+
+    async def async_setup_platform(
+        self,
+        hass: HomeAssistant,
+        config: ConfigType,
+        async_add_entities: AddEntitiesCallback,
+        discovery_info: DiscoveryInfoType | None = None,
+    ) -> None:
+        """Set up an integration platform async."""
+
+    def setup_platform(
+        self,
+        hass: HomeAssistant,
+        config: ConfigType,
+        add_entities: AddEntitiesCallback,
+        discovery_info: DiscoveryInfoType | None = None,
+    ) -> None:
+        """Set up an integration platform."""
+
+    async def async_setup_entry(
+        self,
+        hass: HomeAssistant,
+        entry: config_entries.ConfigEntry,
+        async_add_entities: AddEntitiesCallback,
+    ) -> None:
+        """Set up an integration platform from a config entry."""
+
+
 class EntityPlatform:
     """Manage the entities for a single platform."""
 
@@ -81,7 +110,7 @@ class EntityPlatform:
         logger: Logger,
         domain: str,
         platform_name: str,
-        platform: ModuleType | None,
+        platform: EntityPlatformModule | None,
         scan_interval: timedelta,
         entity_namespace: str | None,
     ) -> None:
@@ -95,7 +124,7 @@ class EntityPlatform:
         self.entity_namespace = entity_namespace
         self.config_entry: config_entries.ConfigEntry | None = None
         self.entities: dict[str, Entity] = {}
-        self._tasks: list[asyncio.Future] = []
+        self._tasks: list[asyncio.Task[None]] = []
         # Stop tracking tasks after setup is completed
         self._setup_complete = False
         # Method to cancel the state change listener
@@ -169,10 +198,12 @@ class EntityPlatform:
             return
 
         @callback
-        def async_create_setup_task() -> Coroutine:
+        def async_create_setup_task() -> Coroutine[
+            Any, Any, None
+        ] | asyncio.Future[None]:
             """Get task to set up platform."""
             if getattr(platform, "async_setup_platform", None):
-                return platform.async_setup_platform(  # type: ignore[no-any-return,union-attr]
+                return platform.async_setup_platform(  # type: ignore[union-attr]
                     hass,
                     platform_config,
                     self._async_schedule_add_entities,
@@ -181,7 +212,7 @@ class EntityPlatform:
 
             # This should not be replaced with hass.async_add_job because
             # we don't want to track this task in case it blocks startup.
-            return hass.loop.run_in_executor(  # type: ignore[return-value]
+            return hass.loop.run_in_executor(
                 None,
                 platform.setup_platform,  # type: ignore[union-attr]
                 hass,
@@ -211,18 +242,18 @@ class EntityPlatform:
         platform = self.platform
 
         @callback
-        def async_create_setup_task() -> Coroutine:
+        def async_create_setup_task() -> Coroutine[Any, Any, None]:
             """Get task to set up platform."""
             config_entries.current_entry.set(config_entry)
 
-            return platform.async_setup_entry(  # type: ignore[no-any-return,union-attr]
+            return platform.async_setup_entry(  # type: ignore[union-attr]
                 self.hass, config_entry, self._async_schedule_add_entities_for_entry
             )
 
         return await self._async_setup_platform(async_create_setup_task)
 
     async def _async_setup_platform(
-        self, async_create_setup_task: Callable[[], Coroutine], tries: int = 0
+        self, async_create_setup_task: Callable[[], Awaitable[None]], tries: int = 0
     ) -> bool:
         """Set up a platform via config file or config entry.
 
@@ -701,7 +732,7 @@ class EntityPlatform:
     def async_register_entity_service(
         self,
         name: str,
-        schema: dict | vol.Schema,
+        schema: dict[str, Any] | vol.Schema,
         func: str | Callable[..., Any],
         required_features: Iterable[int] | None = None,
     ) -> None:
@@ -753,7 +784,7 @@ class EntityPlatform:
             return
 
         async with self._process_updates:
-            tasks = []
+            tasks: list[Coroutine[Any, Any, None]] = []
             for entity in self.entities.values():
                 if not entity.should_poll:
                     continue

mypy.ini

@@ -72,6 +72,9 @@ disallow_any_generics = true
 [mypy-homeassistant.helpers.entity]
 disallow_any_generics = true
 
+[mypy-homeassistant.helpers.entity_platform]
+disallow_any_generics = true
+
 [mypy-homeassistant.helpers.entity_values]
 disallow_any_generics = true