Reduce duplicate code in unifiprotect (#119624)

homeassistant/components/unifiprotect/binary_sensor.py

@@ -26,10 +26,8 @@ from homeassistant.components.binary_sensor import (
 )
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT
 from .data import ProtectData, UFPConfigEntry
 from .entity import (
     BaseProtectEntity,
@@ -39,7 +37,6 @@ from .entity import (
     async_all_device_entities,
 )
 from .models import PermRequired, ProtectEventMixin, ProtectRequiredKeysMixin
-from .utils import async_dispatch_id as _ufpd
 
 _LOGGER = logging.getLogger(__name__)
 _KEY_DOOR = "door"
@@ -641,9 +638,7 @@ async def async_setup_entry(
             entities += _async_event_entities(data, ufp_device=device)
         async_add_entities(entities)
 
-    entry.async_on_unload(
-        async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
-    )
+    data.async_subscribe_adopt(_add_new_device)
 
     entities = async_all_device_entities(
         data, ProtectDeviceBinarySensor, model_descriptions=_MODEL_DESCRIPTIONS
@@ -660,9 +655,7 @@ def _async_event_entities(
     ufp_device: ProtectAdoptableDeviceModel | None = None,
 ) -> list[ProtectDeviceEntity]:
     entities: list[ProtectDeviceEntity] = []
-    devices = (
-        data.get_by_types({ModelType.CAMERA}) if ufp_device is None else [ufp_device]
-    )
+    devices = data.get_cameras() if ufp_device is None else [ufp_device]
     for device in devices:
         for description in EVENT_SENSORS:
             if not description.has_required(device):

homeassistant/components/unifiprotect/button.py

@@ -20,7 +20,7 @@ from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DEVICES_THAT_ADOPT, DISPATCH_ADD, DISPATCH_ADOPT, DOMAIN
+from .const import DEVICES_THAT_ADOPT, DISPATCH_ADD, DOMAIN
 from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity, async_all_device_entities
 from .models import PermRequired, ProtectRequiredKeysMixin, ProtectSetableKeysMixin, T
@@ -147,9 +147,7 @@ async def async_setup_entry(
             )
         )
 
-    entry.async_on_unload(
-        async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
-    )
+    data.async_subscribe_adopt(_add_new_device)
     entry.async_on_unload(
         async_dispatcher_connect(
             hass, _ufpd(entry, DISPATCH_ADD), _async_add_unadopted_device

homeassistant/components/unifiprotect/camera.py

@@ -3,13 +3,12 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, cast
+from typing import Any
 
 from typing_extensions import Generator
 from uiprotect.data import (
     Camera as UFPCamera,
     CameraChannel,
-    ModelType,
     ProtectAdoptableDeviceModel,
     ProtectModelWithId,
     StateType,
@@ -28,7 +27,6 @@ from .const import (
     ATTR_FPS,
     ATTR_HEIGHT,
     ATTR_WIDTH,
-    DISPATCH_ADOPT,
     DISPATCH_CHANNELS,
     DOMAIN,
 )
@@ -73,11 +71,8 @@ def _get_camera_channels(
 ) -> Generator[tuple[UFPCamera, CameraChannel, bool]]:
     """Get all the camera channels."""
 
-    devices = (
-        data.get_by_types({ModelType.CAMERA}) if ufp_device is None else [ufp_device]
-    )
-    for camera in devices:
-        camera = cast(UFPCamera, camera)
+    cameras = data.get_cameras() if ufp_device is None else [ufp_device]
+    for camera in cameras:
         if not camera.channels:
             if ufp_device is None:
                 # only warn on startup
@@ -157,9 +152,7 @@ async def async_setup_entry(
             return
         async_add_entities(_async_camera_entities(hass, entry, data, ufp_device=device))
 
-    entry.async_on_unload(
-        async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
-    )
+    data.async_subscribe_adopt(_add_new_device)
     entry.async_on_unload(
         async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_CHANNELS), _add_new_device)
     )

homeassistant/components/unifiprotect/data.py

@@ -27,7 +27,10 @@ from uiprotect.utils import log_event
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr
-from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.dispatcher import (
+    async_dispatcher_connect,
+    async_dispatcher_send,
+)
 from homeassistant.helpers.event import async_track_time_interval
 
 from .const import (
@@ -81,6 +84,7 @@ class ProtectData:
 
         self.last_update_success = False
         self.api = protect
+        self._adopt_signal = _ufpd(self._entry, DISPATCH_ADOPT)
 
     @property
     def disable_stream(self) -> bool:
@@ -92,6 +96,15 @@ class ProtectData:
         """Max number of events to load at once."""
         return self._entry.options.get(CONF_MAX_MEDIA, DEFAULT_MAX_MEDIA)
 
+    @callback
+    def async_subscribe_adopt(
+        self, add_callback: Callable[[ProtectAdoptableDeviceModel], None]
+    ) -> None:
+        """Add an callback for on device adopt."""
+        self._entry.async_on_unload(
+            async_dispatcher_connect(self._hass, self._adopt_signal, add_callback)
+        )
+
     def get_by_types(
         self, device_types: Iterable[ModelType], ignore_unadopted: bool = True
     ) -> Generator[ProtectAdoptableDeviceModel]:
@@ -105,6 +118,12 @@ class ProtectData:
                     continue
                 yield device
 
+    def get_cameras(self, ignore_unadopted: bool = True) -> Generator[Camera]:
+        """Get all cameras."""
+        return cast(
+            Generator[Camera], self.get_by_types({ModelType.CAMERA}, ignore_unadopted)
+        )
+
     async def async_setup(self) -> None:
         """Subscribe and do the refresh."""
         self._unsub_websocket = self.api.subscribe_websocket(
@@ -206,8 +225,7 @@ class ProtectData:
                 "Doorbell messages updated. Updating devices with LCD screens"
             )
             self.api.bootstrap.nvr.update_all_messages()
-            for camera in self.get_by_types({ModelType.CAMERA}):
-                camera = cast(Camera, camera)
+            for camera in self.get_cameras():
                 if camera.feature_flags.has_lcd_screen:
                     self._async_signal_device_update(camera)
 

homeassistant/components/unifiprotect/light.py

@@ -14,13 +14,10 @@ from uiprotect.data import (
 
 from homeassistant.components.light import ATTR_BRIGHTNESS, ColorMode, LightEntity
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT
 from .data import UFPConfigEntry
 from .entity import ProtectDeviceEntity
-from .utils import async_dispatch_id as _ufpd
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -40,10 +37,7 @@ async def async_setup_entry(
         ):
             async_add_entities([ProtectLight(data, device)])
 
-    entry.async_on_unload(
-        async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
-    )
-
+    data.async_subscribe_adopt(_add_new_device)
     async_add_entities(
         ProtectLight(data, device)
         for device in data.get_by_types({ModelType.LIGHT})

homeassistant/components/unifiprotect/lock.py

@@ -15,13 +15,10 @@ from uiprotect.data import (
 
 from homeassistant.components.lock import LockEntity, LockEntityDescription
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT
 from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity
-from .utils import async_dispatch_id as _ufpd
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -39,9 +36,7 @@ async def async_setup_entry(
         if isinstance(device, Doorlock):
             async_add_entities([ProtectLock(data, device)])
 
-    entry.async_on_unload(
-        async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
-    )
+    data.async_subscribe_adopt(_add_new_device)
 
     async_add_entities(
         ProtectLock(data, cast(Doorlock, device))

homeassistant/components/unifiprotect/media_player.py

@@ -3,11 +3,10 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, cast
+from typing import Any
 
 from uiprotect.data import (
     Camera,
-    ModelType,
     ProtectAdoptableDeviceModel,
     ProtectModelWithId,
     StateType,
@@ -27,13 +26,10 @@ from homeassistant.components.media_player import (
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT
 from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity
-from .utils import async_dispatch_id as _ufpd
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -53,18 +49,13 @@ async def async_setup_entry(
         ):
             async_add_entities([ProtectMediaPlayer(data, device)])
 
-    entry.async_on_unload(
-        async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
+    data.async_subscribe_adopt(_add_new_device)
+    async_add_entities(
+        ProtectMediaPlayer(data, device)
+        for device in data.get_cameras()
+        if device.has_speaker or device.has_removable_speaker
     )
 
-    entities = []
-    for device in data.get_by_types({ModelType.CAMERA}):
-        device = cast(Camera, device)
-        if device.has_speaker or device.has_removable_speaker:
-            entities.append(ProtectMediaPlayer(data, device))
-
-    async_add_entities(entities)
-
 
 class ProtectMediaPlayer(ProtectDeviceEntity, MediaPlayerEntity):
     """A Ubiquiti UniFi Protect Speaker."""

homeassistant/components/unifiprotect/media_source.py

@@ -7,7 +7,7 @@ from datetime import date, datetime, timedelta
 from enum import Enum
 from typing import Any, NoReturn, cast
 
-from uiprotect.data import Camera, Event, EventType, ModelType, SmartDetectObjectType
+from uiprotect.data import Camera, Event, EventType, SmartDetectObjectType
 from uiprotect.exceptions import NvrError
 from uiprotect.utils import from_js_time
 from yarl import URL
@@ -848,8 +848,7 @@ class ProtectMediaSource(MediaSource):
 
         cameras: list[BrowseMediaSource] = [await self._build_camera(data, "all")]
 
-        for camera in data.get_by_types({ModelType.CAMERA}):
-            camera = cast(Camera, camera)
+        for camera in data.get_cameras():
             if not camera.can_read_media(data.api.bootstrap.auth_user):
                 continue
             cameras.append(await self._build_camera(data, camera.id))

homeassistant/components/unifiprotect/number.py

@@ -20,14 +20,11 @@ from uiprotect.data import (
 from homeassistant.components.number import NumberEntity, NumberEntityDescription
 from homeassistant.const import PERCENTAGE, EntityCategory, UnitOfTime
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT
 from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity, async_all_device_entities
 from .models import PermRequired, ProtectRequiredKeysMixin, ProtectSetableKeysMixin, T
-from .utils import async_dispatch_id as _ufpd
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -245,10 +242,7 @@ async def async_setup_entry(
             )
         )
 
-    entry.async_on_unload(
-        async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
-    )
-
+    data.async_subscribe_adopt(_add_new_device)
     async_add_entities(
         async_all_device_entities(
             data,

homeassistant/components/unifiprotect/select.py

@@ -30,14 +30,13 @@ from uiprotect.data import (
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT, TYPE_EMPTY_VALUE
+from .const import TYPE_EMPTY_VALUE
 from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity, async_all_device_entities
 from .models import PermRequired, ProtectRequiredKeysMixin, ProtectSetableKeysMixin, T
-from .utils import async_dispatch_id as _ufpd, async_get_light_motion_current
+from .utils import async_get_light_motion_current
 
 _LOGGER = logging.getLogger(__name__)
 _KEY_LIGHT_MOTION = "light_motion"
@@ -346,9 +345,7 @@ async def async_setup_entry(
             )
         )
 
-    entry.async_on_unload(
-        async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
-    )
+    data.async_subscribe_adopt(_add_new_device)
     async_add_entities(
         async_all_device_entities(
             data, ProtectSelects, model_descriptions=_MODEL_DESCRIPTIONS

homeassistant/components/unifiprotect/sensor.py

@@ -6,7 +6,7 @@ from collections.abc import Sequence
 from dataclasses import dataclass
 from datetime import datetime
 import logging
-from typing import Any, cast
+from typing import Any
 
 from uiprotect.data import (
     NVR,
@@ -37,10 +37,8 @@ from homeassistant.const import (
     UnitOfTime,
 )
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT
 from .data import ProtectData, UFPConfigEntry
 from .entity import (
     BaseProtectEntity,
@@ -50,7 +48,7 @@ from .entity import (
     async_all_device_entities,
 )
 from .models import PermRequired, ProtectEventMixin, ProtectRequiredKeysMixin, T
-from .utils import async_dispatch_id as _ufpd, async_get_light_motion_current
+from .utils import async_get_light_motion_current
 
 _LOGGER = logging.getLogger(__name__)
 OBJECT_TYPE_NONE = "none"
@@ -641,10 +639,7 @@ async def async_setup_entry(
             entities += _async_event_entities(data, ufp_device=device)
         async_add_entities(entities)
 
-    entry.async_on_unload(
-        async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
-    )
-
+    data.async_subscribe_adopt(_add_new_device)
     entities = async_all_device_entities(
         data,
         ProtectDeviceSensor,
@@ -663,31 +658,28 @@ def _async_event_entities(
     ufp_device: Camera | None = None,
 ) -> list[ProtectDeviceEntity]:
     entities: list[ProtectDeviceEntity] = []
-    devices = (
-        data.get_by_types({ModelType.CAMERA}) if ufp_device is None else [ufp_device]
-    )
-    for device in devices:
-        device = cast(Camera, device)
+    cameras = data.get_cameras() if ufp_device is None else [ufp_device]
+    for camera in cameras:
         for description in MOTION_TRIP_SENSORS:
-            entities.append(ProtectDeviceSensor(data, device, description))
+            entities.append(ProtectDeviceSensor(data, camera, description))
             _LOGGER.debug(
                 "Adding trip sensor entity %s for %s",
                 description.name,
-                device.display_name,
+                camera.display_name,
             )
 
-        if not device.feature_flags.has_smart_detect:
+        if not camera.feature_flags.has_smart_detect:
             continue
 
         for event_desc in LICENSE_PLATE_EVENT_SENSORS:
-            if not event_desc.has_required(device):
+            if not event_desc.has_required(camera):
                 continue
 
-            entities.append(ProtectLicensePlateEventSensor(data, device, event_desc))
+            entities.append(ProtectLicensePlateEventSensor(data, camera, event_desc))
             _LOGGER.debug(
                 "Adding sensor entity %s for %s",
                 description.name,
-                device.display_name,
+                camera.display_name,
             )
 
     return entities

homeassistant/components/unifiprotect/switch.py

@@ -20,15 +20,12 @@ from uiprotect.data import (
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.restore_state import RestoreEntity
 
-from .const import DISPATCH_ADOPT
 from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity, ProtectNVREntity, async_all_device_entities
 from .models import PermRequired, ProtectRequiredKeysMixin, ProtectSetableKeysMixin, T
-from .utils import async_dispatch_id as _ufpd
 
 _LOGGER = logging.getLogger(__name__)
 ATTR_PREV_MIC = "prev_mic_level"
@@ -494,10 +491,7 @@ async def async_setup_entry(
         )
         async_add_entities(entities)
 
-    entry.async_on_unload(
-        async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
-    )
-
+    data.async_subscribe_adopt(_add_new_device)
     entities = async_all_device_entities(
         data,
         ProtectSwitch,

homeassistant/components/unifiprotect/text.py

@@ -17,14 +17,11 @@ from uiprotect.data import (
 from homeassistant.components.text import TextEntity, TextEntityDescription
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT
 from .data import ProtectData, UFPConfigEntry
 from .entity import ProtectDeviceEntity, async_all_device_entities
 from .models import PermRequired, ProtectRequiredKeysMixin, ProtectSetableKeysMixin, T
-from .utils import async_dispatch_id as _ufpd
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -78,10 +75,7 @@ async def async_setup_entry(
             )
         )
 
-    entry.async_on_unload(
-        async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
-    )
-
+    data.async_subscribe_adopt(_add_new_device)
     async_add_entities(
         async_all_device_entities(
             data, ProtectDeviceText, model_descriptions=_MODEL_DESCRIPTIONS