Type UniFi Protect entities against the public device model (#182617)

homeassistant/components/unifiprotect/camera.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Iterable
 import logging
-from typing import cast, override
+from typing import override
 
 from uiprotect.data import (
     Camera as UFPCamera,
@@ -266,7 +266,7 @@ class ProtectCamera(ProtectDeviceEntity, Camera):
         self._last_image: bytes | None = None
         # The base tracks the private device in hybrid (unchanged behaviour) and
         # the public device in public-only, so it always has a mac to key on.
-        super().__init__(data, cast(ProtectDeviceType, private or public))
+        super().__init__(data, private or public)
         self._attr_unique_id = f"{self.device.mac}_{self._channel_id}"
         self._attr_name = get_camera_base_name(quality)
         # only the default (first active) quality channel is enabled by default
@@ -392,11 +392,7 @@ class ProtectCamera(ProtectDeviceEntity, Camera):
             self._public_missing = False
         else:
             self._public_missing = True
-        device = (
-            self._private
-            if self._private is not None
-            else cast(ProtectDeviceType, self._public)
-        )
+        device = self._private if self._private is not None else self._public
         self._async_updated_event(device)
 
     @override

homeassistant/components/unifiprotect/data.py

@@ -53,7 +53,7 @@ from .const import (
 from .utils import async_get_devices_by_type
 
 _LOGGER = logging.getLogger(__name__)
-type ProtectDeviceType = ProtectAdoptableDeviceModel | NVR
+type ProtectDeviceType = ProtectAdoptableDeviceModel | NVR | PublicDeviceModel
 type UFPConfigEntry = ConfigEntry[ProtectData]
 
 
@@ -836,7 +836,7 @@ class ProtectData:
 
     @callback
     def async_get_public_device(
-        self, device: ProtectDeviceType | PublicDeviceModel
+        self, device: ProtectDeviceType
     ) -> PublicDeviceModel | None:
         """Return the public-API object matching a device, if available."""
         api = self.api

homeassistant/components/unifiprotect/entity.py

@@ -52,7 +52,7 @@ from .data import ProtectData, ProtectDeviceType
 
 _LOGGER = logging.getLogger(__name__)
 
-T = TypeVar("T", bound=ProtectAdoptableDeviceModel | NVR)
+T = TypeVar("T", bound=ProtectDeviceType)
 
 
 class PermRequired(int, Enum):
@@ -144,7 +144,7 @@ def _async_public_only_entities(
         entities.append(
             klass(
                 data,
-                device=cast(ProtectDeviceType, public),
+                device=public,
                 description=description,
             )
         )
@@ -332,7 +332,7 @@ class BaseProtectEntity(Entity):
     def __init__(
         self,
         data: ProtectData,
-        device: ProtectDeviceType | PublicDeviceModel,
+        device: ProtectDeviceType,
         description: EntityDescription | None = None,
     ) -> None:
         """Initialize the entity."""
@@ -342,7 +342,7 @@ class BaseProtectEntity(Entity):
             self._ufp_has_private = False
             self._ufp_public_obj = device
         # The base keys on the mac, which both model trees carry.
-        self.device = cast(ProtectDeviceType, device)
+        self.device = device
 
         if description is None:
             self._attr_unique_id = self.device.mac
@@ -417,7 +417,7 @@ class BaseProtectEntity(Entity):
             self._attr_available = available
 
     @callback
-    def _ufp_set_target(self) -> ProtectDeviceType | PublicDeviceModel:
+    def _ufp_set_target(self) -> ProtectDeviceType:
         """Return the object a description's setter is called on.
 
         A migrated description writes through the public object it reads from,
@@ -505,9 +505,7 @@ class ProtectIsOnEntity(BaseProtectEntity):
     entity_description: ProtectEntityDescription
 
     @override
-    def _async_update_device_from_protect(
-        self, device: ProtectAdoptableDeviceModel | NVR
-    ) -> None:
+    def _async_update_device_from_protect(self, device: ProtectDeviceType) -> None:
         super()._async_update_device_from_protect(device)
         was_on = self._attr_is_on
         value = self.entity_description.get_value(device, self._ufp_public_obj)
@@ -521,30 +519,27 @@ class ProtectDeviceEntity(BaseProtectEntity):
     @callback
     @override
     def _async_set_device_info(self) -> None:
-        if not self._ufp_has_private:
+        if isinstance(device := self.device, PublicDeviceModel):
             # market_name/firmware/URL are private-only; the NVR link uses the
             # device id registered at setup.
-            public = self._ufp_public_obj
-            if TYPE_CHECKING:
-                assert public is not None
             self._attr_device_info = DeviceInfo(
-                name=public.display_name,
-                model=public.type,
-                model_id=public.type,
+                name=device.display_name,
+                model=device.type,
+                model_id=device.type,
                 manufacturer=DEFAULT_BRAND,
-                connections={(dr.CONNECTION_NETWORK_MAC, public.mac)},
+                connections={(dr.CONNECTION_NETWORK_MAC, device.mac)},
                 via_device_id=self.data.nvr_device_id,
             )
             return
         self._attr_device_info = DeviceInfo(
-            name=self.device.display_name,
+            name=device.display_name,
             manufacturer=DEFAULT_BRAND,
-            model=self.device.market_name or self.device.type,
-            model_id=self.device.type,
+            model=device.market_name or device.type,
+            model_id=device.type,
             via_device_id=self.data.nvr_device_id,
-            sw_version=self.device.firmware_version,
-            connections={(dr.CONNECTION_NETWORK_MAC, self.device.mac)},
-            configuration_url=self.device.protect_url,
+            sw_version=device.firmware_version,
+            connections={(dr.CONNECTION_NETWORK_MAC, device.mac)},
+            configuration_url=device.protect_url,
         )
 
 

homeassistant/components/unifiprotect/light.py

@@ -1,7 +1,7 @@
 """Component providing Lights for UniFi Protect."""
 
 import logging
-from typing import Any, cast, override
+from typing import TYPE_CHECKING, Any, cast, override
 
 from uiprotect.data import (
     Light,
@@ -88,7 +88,7 @@ def hass_to_unifi_brightness(value: int) -> int:
 class ProtectLight(ProtectDeviceEntity, LightEntity):
     """A Ubiquiti UniFi Protect Light Entity."""
 
-    device: Light
+    device: Light | PublicLight
 
     _attr_icon = "mdi:spotlight-beam"
     _attr_color_mode = ColorMode.BRIGHTNESS
@@ -108,7 +108,11 @@ class ProtectLight(ProtectDeviceEntity, LightEntity):
         self._ufp_public_obj = public
         # unique_id and device info derive from the base device, so hybrid must
         # keep the private one to leave existing entities unchanged.
-        super().__init__(data, cast(ProtectDeviceType, private or public))
+        device = private or public
+        if TYPE_CHECKING:
+            # The platform only builds a light when at least one side exists.
+            assert device is not None
+        super().__init__(data, device)
 
     @callback
     @override