Improve WLED typing (#77200)

homeassistant/components/wled/coordinator.py

@@ -2,13 +2,12 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable
 
 from wled import WLED, Device as WLEDDevice, WLEDConnectionClosed, WLEDError
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, EVENT_HOMEASSISTANT_STOP
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import CALLBACK_TYPE, Event, HomeAssistant, callback
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
@@ -38,7 +37,7 @@ class WLEDDataUpdateCoordinator(DataUpdateCoordinator[WLEDDevice]):
             CONF_KEEP_MASTER_LIGHT, DEFAULT_KEEP_MASTER_LIGHT
         )
         self.wled = WLED(entry.data[CONF_HOST], session=async_get_clientsession(hass))
-        self.unsub: Callable | None = None
+        self.unsub: CALLBACK_TYPE | None = None
 
         super().__init__(
             hass,
@@ -85,7 +84,7 @@ class WLEDDataUpdateCoordinator(DataUpdateCoordinator[WLEDDevice]):
                 self.unsub()
                 self.unsub = None
 
-        async def close_websocket(_) -> None:
+        async def close_websocket(_: Event) -> None:
             """Close WebSocket connection."""
             self.unsub = None
             await self.wled.disconnect()

homeassistant/components/wled/diagnostics.py

@@ -17,7 +17,7 @@ async def async_get_config_entry_diagnostics(
     """Return diagnostics for a config entry."""
     coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
 
-    data = {
+    data: dict[str, Any] = {
         "info": async_redact_data(coordinator.data.info.__dict__, "wifi"),
         "state": coordinator.data.state.__dict__,
         "effects": {

homeassistant/components/wled/helpers.py

@@ -1,18 +1,30 @@
 """Helpers for WLED."""
+from __future__ import annotations
 
+from collections.abc import Callable, Coroutine
+from typing import Any, TypeVar
+
+from typing_extensions import Concatenate, ParamSpec
 from wled import WLEDConnectionError, WLEDError
 
 from homeassistant.exceptions import HomeAssistantError
 
+from .models import WLEDEntity
 
-def wled_exception_handler(func):
+_WLEDEntityT = TypeVar("_WLEDEntityT", bound=WLEDEntity)
+_P = ParamSpec("_P")
+
+
+def wled_exception_handler(
+    func: Callable[Concatenate[_WLEDEntityT, _P], Coroutine[Any, Any, Any]]
+) -> Callable[Concatenate[_WLEDEntityT, _P], Coroutine[Any, Any, None]]:
     """Decorate WLED calls to handle WLED exceptions.
 
     A decorator that wraps the passed in function, catches WLED errors,
     and handles the availability of the device in the data coordinator.
     """
 
-    async def handler(self, *args, **kwargs):
+    async def handler(self: _WLEDEntityT, *args: _P.args, **kwargs: _P.kwargs) -> None:
         try:
             await func(self, *args, **kwargs)
             self.coordinator.async_update_listeners()

homeassistant/components/wled/light.py

@@ -253,7 +253,7 @@ class WLEDSegmentLight(WLEDEntity, LightEntity):
 def async_update_segments(
     coordinator: WLEDDataUpdateCoordinator,
     current_ids: set[int],
-    async_add_entities,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Update segments."""
     segment_ids = {light.segment_id for light in coordinator.data.state.segments}

homeassistant/components/wled/number.py

@@ -115,12 +115,12 @@ class WLEDNumber(WLEDEntity, NumberEntity):
 def async_update_segments(
     coordinator: WLEDDataUpdateCoordinator,
     current_ids: set[int],
-    async_add_entities,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Update segments."""
     segment_ids = {segment.segment_id for segment in coordinator.data.state.segments}
 
-    new_entities = []
+    new_entities: list[WLEDNumber] = []
 
     # Process new segments, add them to Home Assistant
     for segment_id in segment_ids - current_ids:

homeassistant/components/wled/select.py

@@ -183,12 +183,12 @@ class WLEDPaletteSelect(WLEDEntity, SelectEntity):
 def async_update_segments(
     coordinator: WLEDDataUpdateCoordinator,
     current_ids: set[int],
-    async_add_entities,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Update segments."""
     segment_ids = {segment.segment_id for segment in coordinator.data.state.segments}
 
-    new_entities = []
+    new_entities: list[WLEDPaletteSelect] = []
 
     # Process new segments, add them to Home Assistant
     for segment_id in segment_ids - current_ids:

homeassistant/components/wled/switch.py

@@ -203,12 +203,12 @@ class WLEDReverseSwitch(WLEDEntity, SwitchEntity):
 def async_update_segments(
     coordinator: WLEDDataUpdateCoordinator,
     current_ids: set[int],
-    async_add_entities,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Update segments."""
     segment_ids = {segment.segment_id for segment in coordinator.data.state.segments}
 
-    new_entities = []
+    new_entities: list[WLEDReverseSwitch] = []
 
     # Process new segments, add them to Home Assistant
     for segment_id in segment_ids - current_ids: