proper generics

Author: dmulcahey

zha/application/platforms/__init__.py

@@ -5,6 +5,7 @@
 from abc import abstractmethod
 import asyncio
 from contextlib import suppress
+import functools
 from functools import cached_property
 import logging
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, final
@@ -29,7 +30,7 @@
 
 if TYPE_CHECKING:
     from zha.zigbee.cluster_handlers import ClusterHandler
-    from zha.zigbee.device import Device
+    from zha.zigbee.device import Device, WebSocketClientDevice
     from zha.zigbee.endpoint import Endpoint
     from zha.zigbee.group import Group
 
@@ -483,10 +484,13 @@ async def async_update(self, _: Any | None = None) -> None:
 class WebSocketClientEntity(BaseEntity, Generic[BaseEntityInfoType]):
     """Entity repsentation for the websocket client."""
 
-    def __init__(self, entity_info: BaseEntityInfoType) -> None:
+    def __init__(
+        self, entity_info: BaseEntityInfoType, device: WebSocketClientDevice
+    ) -> None:
         """Initialize the websocket client entity."""
         super().__init__(entity_info.unique_id)
         self.PLATFORM = entity_info.platform
+        self._device: WebSocketClientDevice = device
         self._entity_info: BaseEntityInfoType = entity_info
         self._attr_enabled = self._entity_info.enabled
         self._attr_fallback_name = self._entity_info.fallback_name
@@ -498,6 +502,11 @@ def __init__(self, entity_info: BaseEntityInfoType) -> None:
         self._attr_device_class = self._entity_info.device_class
         self._attr_state_class = self._entity_info.state_class
 
+    @functools.cached_property
+    def info_object(self) -> BaseEntityInfoType:
+        """Return a representation of the alarm control panel."""
+        return self._entity_info
+
     @property
     def state(self) -> dict[str, Any]:
         """Return the arguments to use in the command."""
@@ -507,3 +516,8 @@ def state(self) -> dict[str, Any]:
     def state(self, value: dict[str, Any]) -> None:
         """Set the state of the entity."""
         self._entity_info.state = value
+
+    async def async_update(self) -> None:
+        """Retrieve latest state."""
+        self.debug("polling current state")
+        await self._device.gateway.entities.refresh_state(self._entity_info)

zha/application/platforms/alarm_control_panel/__init__.py

@@ -183,7 +183,7 @@ async def async_alarm_trigger(self, code: str | None = None, **kwargs) -> None:
 
 
 class WebSocketClientAlarmControlPanel(
-    WebSocketClientEntity, AlarmControlPanelEntityInterface
+    WebSocketClientEntity[AlarmControlPanelEntityInfo], AlarmControlPanelEntityInterface
 ):
     """Alarm control panel entity for the WebSocket API."""
 
@@ -194,13 +194,7 @@ def __init__(
         self, entity_info: AlarmControlPanelEntityInfo, device: WebSocketClientDevice
     ) -> None:
         """Initialize the ZHA alarm control device."""
-        super().__init__(entity_info)
-        self._device: WebSocketClientDevice = device
-
-    @functools.cached_property
-    def info_object(self) -> AlarmControlPanelEntityInfo:
-        """Return a representation of the alarm control panel."""
-        return self._entity_info
+        super().__init__(entity_info, device)
 
     @property
     def code_arm_required(self) -> bool:

zha/application/platforms/binary_sensor/__init__.py

@@ -404,7 +404,9 @@ class DanfossPreheatStatus(BinarySensor):
     _attr_entity_category = EntityCategory.DIAGNOSTIC
 
 
-class WebSocketClientBinarySensor(WebSocketClientEntity, BinarySensorEntityInterface):
+class WebSocketClientBinarySensor(
+    WebSocketClientEntity[BinarySensorEntityInfo], BinarySensorEntityInterface
+):
     """Base class for binary sensors that are updated via a websocket client."""
 
     PLATFORM: Platform = Platform.BINARY_SENSOR
@@ -413,20 +415,9 @@ def __init__(
         self, entity_info: BinarySensorEntityInfo, device: WebSocketClientDevice
     ) -> None:
         """Initialize the ZHA alarm control device."""
-        super().__init__(entity_info)
-        self._device: WebSocketClientDevice = device
-
-    @functools.cached_property
-    def info_object(self) -> BinarySensorEntityInfo:
-        """Return a representation of the binary sensor."""
-        return self._entity_info
+        super().__init__(entity_info, device)
 
     @property
     def is_on(self) -> bool:
         """Return True if the switch is on based on the state machine."""
         return self.info_object.state.state
-
-    async def async_update(self) -> None:
-        """Retrieve latest state."""
-        self.debug("polling current state")
-        await self._device.gateway.entities.refresh_state(self._entity_info)

zha/application/platforms/button/__init__.py

@@ -234,7 +234,9 @@ class AqaraSelfTestButton(WriteAttributeButton):
     _attr_translation_key = "self_test"
 
 
-class WebSocketClientButtonEntity(WebSocketClientEntity, ButtonEntityInterface):
+class WebSocketClientButtonEntity(
+    WebSocketClientEntity[ButtonEntityInfo], ButtonEntityInterface
+):
     """Defines a ZHA button that is controlled via a websocket."""
 
     PLATFORM = Platform.BUTTON
@@ -243,13 +245,7 @@ def __init__(
         self, entity_info: ButtonEntityInfo, device: WebSocketClientDevice
     ) -> None:
         """Initialize the ZHA alarm control device."""
-        super().__init__(entity_info)
-        self._device: WebSocketClientDevice = device
-
-    @functools.cached_property
-    def info_object(self) -> ButtonEntityInfo:
-        """Return a representation of the button."""
-        return self._entity_info
+        super().__init__(entity_info, device)
 
     @functools.cached_property
     def args(self) -> list[Any]:

zha/application/platforms/climate/__init__.py

@@ -962,7 +962,9 @@ async def async_preset_handler(self, preset: str, enable: bool = False) -> None:
             )
 
 
-class WebSocketClientThermostatEntity(WebSocketClientEntity, ClimateEntityInterface):
+class WebSocketClientThermostatEntity(
+    WebSocketClientEntity[ThermostatEntityInfo], ClimateEntityInterface
+):
     """Representation of a ZHA Thermostat device."""
 
     PLATFORM: Platform = Platform.CLIMATE
@@ -971,13 +973,7 @@ def __init__(
         self, entity_info: ThermostatEntityInfo, device: WebSocketClientDevice
     ) -> None:
         """Initialize the ZHA climate entity."""
-        super().__init__(entity_info)
-        self._device: WebSocketClientDevice = device
-
-    @property
-    def info_object(self) -> ThermostatEntityInfo:
-        """Return a representation of the thermostat."""
-        return self._entity_info
+        super().__init__(entity_info, device)
 
     @property
     def current_temperature(self) -> float | None:

zha/application/platforms/cover/__init__.py

@@ -596,7 +596,9 @@ async def async_open_cover(self, **kwargs: Any) -> None:
         self.maybe_emit_state_changed_event()
 
 
-class WebSocketClientCoverEntity(WebSocketClientEntity, CoverEntityInterface):
+class WebSocketClientCoverEntity(
+    WebSocketClientEntity[CoverEntityInfo], CoverEntityInterface
+):
     """Representation of a ZHA cover."""
 
     PLATFORM: Platform = Platform.COVER
@@ -605,13 +607,7 @@ def __init__(
         self, entity_info: CoverEntityInfo, device: WebSocketClientDevice
     ) -> None:
         """Initialize the ZHA fan entity."""
-        super().__init__(entity_info)
-        self._device: WebSocketClientDevice = device
-
-    @property
-    def info_object(self) -> CoverEntityInfo:
-        """Return the info object for this entity."""
-        return self._entity_info
+        super().__init__(entity_info, device)
 
     @property
     def supported_features(self) -> CoverEntityFeature:

zha/application/platforms/device_tracker/__init__.py

@@ -160,7 +160,7 @@ def handle_cluster_handler_attribute_updated(
 
 
 class WebSocketClientDeviceTrackerEntity(
-    WebSocketClientEntity, DeviceTrackerEntityInterface
+    WebSocketClientEntity[DeviceTrackerEntityInfo], DeviceTrackerEntityInterface
 ):
     """Device tracker entity for the WebSocket API."""
 
@@ -170,13 +170,7 @@ def __init__(
         self, entity_info: DeviceTrackerEntityInfo, device: WebSocketClientDevice
     ) -> None:
         """Initialize the ZHA device tracker."""
-        super().__init__(entity_info)
-        self._device: WebSocketClientDevice = device
-
-    @property
-    def info_object(self) -> DeviceTrackerEntityInfo:
-        """Return a representation of the device tracker."""
-        return self._entity_info
+        super().__init__(entity_info, device)
 
     @property
     def is_connected(self) -> bool:

zha/application/platforms/fan/__init__.py

@@ -538,7 +538,9 @@ def preset_modes_to_name(self) -> dict[int, str]:
         return {6: PRESET_MODE_SMART}
 
 
-class WebSocketClientFanEntity(WebSocketClientEntity, FanEntityInterface):
+class WebSocketClientFanEntity(
+    WebSocketClientEntity[FanEntityInfo], FanEntityInterface
+):
     """Representation of a ZHA fan over WebSocket."""
 
     PLATFORM: Platform = Platform.FAN
@@ -547,13 +549,7 @@ def __init__(
         self, entity_info: FanEntityInfo, device: WebSocketClientDevice
     ) -> None:
         """Initialize the ZHA fan entity."""
-        super().__init__(entity_info)
-        self._device: WebSocketClientDevice = device
-
-    @property
-    def info_object(self) -> FanEntityInfo:
-        """Return the fan entity info."""
-        return self._entity_info
+        super().__init__(entity_info, device)
 
     @property
     def preset_modes(self) -> list[str]:

zha/application/platforms/light/__init__.py

@@ -1330,7 +1330,9 @@ def restore_external_state_attributes(
             self._off_brightness = off_brightness
 
 
-class WebSocketClientLightEntity(WebSocketClientEntity, LightEntityInterface):
+class WebSocketClientLightEntity(
+    WebSocketClientEntity[LightEntityInfo], LightEntityInterface
+):
     """Light entity that sends commands to a websocket client."""
 
     PLATFORM: Platform = Platform.LIGHT
@@ -1339,13 +1341,7 @@ def __init__(
         self, entity_info: LightEntityInfo, device: WebSocketClientDevice
     ) -> None:
         """Initialize the ZHA lock entity."""
-        super().__init__(entity_info)
-        self._device: WebSocketClientDevice = device
-
-    @property
-    def info_object(self) -> LightEntityInfo:
-        """Return a representation of the select."""
-        return self._entity_info
+        super().__init__(entity_info, device)
 
     @property
     def xy_color(self) -> tuple[float, float] | None:

zha/application/platforms/lock/__init__.py

@@ -171,7 +171,9 @@ def restore_external_state_attributes(
         self._state = state
 
 
-class WebSocketClientLockEntity(WebSocketClientEntity, LockEntityInterface):
+class WebSocketClientLockEntity(
+    WebSocketClientEntity[LockEntityInfo], LockEntityInterface
+):
     """Representation of a ZHA lock on the client side."""
 
     PLATFORM: Platform = Platform.LOCK
@@ -180,13 +182,7 @@ def __init__(
         self, entity_info: LockEntityInfo, device: WebSocketClientDevice
     ) -> None:
         """Initialize the ZHA lock entity."""
-        super().__init__(entity_info)
-        self._device: WebSocketClientDevice = device
-
-    @property
-    def info_object(self) -> LockEntityInfo:
-        """Return a representation of the lock."""
-        return self._entity_info
+        super().__init__(entity_info, device)
 
     @property
     def is_locked(self) -> bool: