Refactor and unify device fetching for UniFi Protect (#77341)

This commit is contained in:
Christopher Bailey 2022-08-26 07:46:11 -04:00 committed by GitHub
parent dfc3e7d80f
commit 1fb8fbf5de
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 50 additions and 44 deletions

View file

@ -10,6 +10,7 @@ from pyunifiprotect.data import (
Camera,
Event,
Light,
ModelType,
MountType,
ProtectAdoptableDeviceModel,
ProtectModelWithId,
@ -409,12 +410,9 @@ def _async_motion_entities(
) -> list[ProtectDeviceEntity]:
entities: list[ProtectDeviceEntity] = []
devices = (
data.api.bootstrap.cameras.values() if ufp_device is None else [ufp_device]
data.get_by_types({ModelType.CAMERA}) if ufp_device is None else [ufp_device]
)
for device in devices:
if not device.is_adopted:
continue
for description in MOTION_SENSORS:
entities.append(ProtectEventBinarySensor(data, device, description))
_LOGGER.debug(

View file

@ -100,10 +100,8 @@ def _async_remove_adopt_button(
) -> None:
entity_registry = er.async_get(hass)
if device.is_adopted_by_us and (
entity_id := entity_registry.async_get_entity_id(
Platform.BUTTON, DOMAIN, f"{device.mac}_adopt"
)
if entity_id := entity_registry.async_get_entity_id(
Platform.BUTTON, DOMAIN, f"{device.mac}_adopt"
):
entity_registry.async_remove(entity_id)

View file

@ -3,10 +3,12 @@ from __future__ import annotations
from collections.abc import Generator
import logging
from typing import cast
from pyunifiprotect.data import (
Camera as UFPCamera,
CameraChannel,
ModelType,
ProtectAdoptableDeviceModel,
ProtectModelWithId,
StateType,
@ -42,12 +44,10 @@ def get_camera_channels(
"""Get all the camera channels."""
devices = (
data.api.bootstrap.cameras.values() if ufp_device is None else [ufp_device]
data.get_by_types({ModelType.CAMERA}) if ufp_device is None else [ufp_device]
)
for camera in devices:
if not camera.is_adopted_by_us:
continue
camera = cast(UFPCamera, camera)
if not camera.channels:
if ufp_device is None:
# only warn on startup

View file

@ -4,12 +4,13 @@ from __future__ import annotations
from collections.abc import Callable, Generator, Iterable
from datetime import timedelta
import logging
from typing import Any, Union
from typing import Any, Union, cast
from pyunifiprotect import ProtectApiClient
from pyunifiprotect.data import (
NVR,
Bootstrap,
Camera,
Event,
EventType,
Liveview,
@ -35,11 +36,7 @@ from .const import (
DISPATCH_CHANNELS,
DOMAIN,
)
from .utils import (
async_dispatch_id as _ufpd,
async_get_devices,
async_get_devices_by_type,
)
from .utils import async_dispatch_id as _ufpd, async_get_devices_by_type
_LOGGER = logging.getLogger(__name__)
ProtectDeviceType = Union[ProtectAdoptableDeviceModel, NVR]
@ -92,13 +89,17 @@ class ProtectData:
return self._entry.options.get(CONF_MAX_MEDIA, DEFAULT_MAX_MEDIA)
def get_by_types(
self, device_types: Iterable[ModelType]
self, device_types: Iterable[ModelType], ignore_unadopted: bool = True
) -> Generator[ProtectAdoptableDeviceModel, None, None]:
"""Get all devices matching types."""
for device_type in device_types:
yield from async_get_devices_by_type(
devices = async_get_devices_by_type(
self.api.bootstrap, device_type
).values()
for device in devices:
if ignore_unadopted and not device.is_adopted_by_us:
continue
yield device
async def async_setup(self) -> None:
"""Subscribe and do the refresh."""
@ -202,7 +203,8 @@ class ProtectData:
"Doorbell messages updated. Updating devices with LCD screens"
)
self.api.bootstrap.nvr.update_all_messages()
for camera in self.api.bootstrap.cameras.values():
for camera in self.get_by_types({ModelType.CAMERA}):
camera = cast(Camera, camera)
if camera.feature_flags.has_lcd_screen:
self._async_signal_device_update(camera)
@ -250,7 +252,7 @@ class ProtectData:
return
self._async_signal_device_update(self.api.bootstrap.nvr)
for device in async_get_devices(self.api.bootstrap, DEVICES_THAT_ADOPT):
for device in self.get_by_types(DEVICES_THAT_ADOPT):
self._async_signal_device_update(device)
@callback

View file

@ -46,7 +46,9 @@ def _async_device_entities(
entities: list[ProtectDeviceEntity] = []
devices = (
[ufp_device] if ufp_device is not None else data.get_by_types({model_type})
[ufp_device]
if ufp_device is not None
else data.get_by_types({model_type}, ignore_unadopted=False)
)
for device in devices:
assert isinstance(device, (Camera, Light, Sensor, Viewer, Doorlock, Chime))

View file

@ -44,10 +44,7 @@ async def async_setup_entry(
)
entities = []
for device in data.api.bootstrap.lights.values():
if not device.is_adopted_by_us:
continue
for device in data.get_by_types({ModelType.LIGHT}):
if device.can_write(data.api.bootstrap.auth_user):
entities.append(ProtectLight(data, device))

View file

@ -2,11 +2,12 @@
from __future__ import annotations
import logging
from typing import Any
from typing import Any, cast
from pyunifiprotect.data import (
Doorlock,
LockStatusType,
ModelType,
ProtectAdoptableDeviceModel,
ProtectModelWithId,
)
@ -42,10 +43,8 @@ async def async_setup_entry(
)
entities = []
for device in data.api.bootstrap.doorlocks.values():
if not device.is_adopted_by_us:
continue
for device in data.get_by_types({ModelType.DOORLOCK}):
device = cast(Doorlock, device)
entities.append(ProtectLock(data, device))
async_add_entities(entities)

View file

@ -2,9 +2,14 @@
from __future__ import annotations
import logging
from typing import Any
from typing import Any, cast
from pyunifiprotect.data import Camera, ProtectAdoptableDeviceModel, ProtectModelWithId
from pyunifiprotect.data import (
Camera,
ModelType,
ProtectAdoptableDeviceModel,
ProtectModelWithId,
)
from pyunifiprotect.exceptions import StreamError
from homeassistant.components import media_source
@ -51,9 +56,8 @@ async def async_setup_entry(
)
entities = []
for device in data.api.bootstrap.cameras.values():
if not device.is_adopted_by_us:
continue
for device in data.get_by_types({ModelType.CAMERA}):
device = cast(Camera, device)
if device.feature_flags.has_speaker:
entities.append(ProtectMediaPlayer(data, device))

View file

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

View file

@ -4,13 +4,14 @@ from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime
import logging
from typing import Any
from typing import Any, cast
from pyunifiprotect.data import (
NVR,
Camera,
Event,
Light,
ModelType,
ProtectAdoptableDeviceModel,
ProtectDeviceModel,
ProtectModelWithId,
@ -649,12 +650,10 @@ def _async_motion_entities(
) -> list[ProtectDeviceEntity]:
entities: list[ProtectDeviceEntity] = []
devices = (
data.api.bootstrap.cameras.values() if ufp_device is None else [ufp_device]
data.get_by_types({ModelType.CAMERA}) if ufp_device is None else [ufp_device]
)
for device in devices:
if not device.is_adopted_by_us:
continue
device = cast(Camera, device)
for description in MOTION_TRIP_SENSORS:
entities.append(ProtectDeviceSensor(data, device, description))
_LOGGER.debug(