From 259befa65fab68f8376950728613a3bc6a34bd09 Mon Sep 17 00:00:00 2001 From: Christopher Bailey Date: Mon, 17 Jan 2022 15:51:55 -0500 Subject: [PATCH] Cleans up various asserts/type ignores for UniFi Protect (#63824) Co-authored-by: epenet --- .../components/unifiprotect/models.py | 24 +++---- .../components/unifiprotect/number.py | 14 ++--- .../components/unifiprotect/select.py | 62 ++++++++----------- .../components/unifiprotect/sensor.py | 44 ++++++------- .../components/unifiprotect/switch.py | 16 ++--- 5 files changed, 74 insertions(+), 86 deletions(-) diff --git a/homeassistant/components/unifiprotect/models.py b/homeassistant/components/unifiprotect/models.py index 49e6c61907fe..df11c46a5d9c 100644 --- a/homeassistant/components/unifiprotect/models.py +++ b/homeassistant/components/unifiprotect/models.py @@ -4,9 +4,9 @@ from __future__ import annotations from collections.abc import Callable, Coroutine from dataclasses import dataclass import logging -from typing import Any +from typing import Any, Generic, TypeVar -from pyunifiprotect.data import NVR, ProtectAdoptableDeviceModel +from pyunifiprotect.data import ProtectDeviceModel from homeassistant.helpers.entity import EntityDescription @@ -14,17 +14,19 @@ from .utils import get_nested_attr _LOGGER = logging.getLogger(__name__) +T = TypeVar("T", bound=ProtectDeviceModel) + @dataclass -class ProtectRequiredKeysMixin: +class ProtectRequiredKeysMixin(Generic[T]): """Mixin for required keys.""" ufp_required_field: str | None = None ufp_value: str | None = None - ufp_value_fn: Callable[[ProtectAdoptableDeviceModel | NVR], Any] | None = None + ufp_value_fn: Callable[[T], Any] | None = None ufp_enabled: str | None = None - def get_ufp_value(self, obj: ProtectAdoptableDeviceModel | NVR) -> Any: + def get_ufp_value(self, obj: T) -> Any: """Return value from UniFi Protect device.""" if self.ufp_value is not None: return get_nested_attr(obj, self.ufp_value) @@ -36,7 +38,7 @@ class ProtectRequiredKeysMixin: "`ufp_value` or `ufp_value_fn` is required" ) - def get_ufp_enabled(self, obj: ProtectAdoptableDeviceModel | NVR) -> bool: + def get_ufp_enabled(self, obj: T) -> bool: """Return value from UniFi Protect device.""" if self.ufp_enabled is not None: return bool(get_nested_attr(obj, self.ufp_enabled)) @@ -44,15 +46,13 @@ class ProtectRequiredKeysMixin: @dataclass -class ProtectSetableKeysMixin(ProtectRequiredKeysMixin): - """Mixin to for settable values.""" +class ProtectSetableKeysMixin(ProtectRequiredKeysMixin, Generic[T]): + """Mixin for settable values.""" ufp_set_method: str | None = None - ufp_set_method_fn: Callable[ - [ProtectAdoptableDeviceModel, Any], Coroutine[Any, Any, None] - ] | None = None + ufp_set_method_fn: Callable[[T, Any], Coroutine[Any, Any, None]] | None = None - async def ufp_set(self, obj: ProtectAdoptableDeviceModel, value: Any) -> None: + async def ufp_set(self, obj: T, value: Any) -> None: """Set value for UniFi Protect device.""" assert isinstance(self, EntityDescription) _LOGGER.debug("Setting %s to %s for %s", self.name, value, obj.name) diff --git a/homeassistant/components/unifiprotect/number.py b/homeassistant/components/unifiprotect/number.py index 9a5cd565090e..7f47dfba7fb7 100644 --- a/homeassistant/components/unifiprotect/number.py +++ b/homeassistant/components/unifiprotect/number.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass from datetime import timedelta -from typing import Any +from typing import Generic from pyunifiprotect.data.devices import Camera, Light @@ -16,7 +16,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from .const import DOMAIN from .data import ProtectData from .entity import ProtectDeviceEntity, async_all_device_entities -from .models import ProtectSetableKeysMixin +from .models import ProtectSetableKeysMixin, T @dataclass @@ -30,18 +30,16 @@ class NumberKeysMixin: @dataclass class ProtectNumberEntityDescription( - ProtectSetableKeysMixin, NumberEntityDescription, NumberKeysMixin + ProtectSetableKeysMixin, NumberEntityDescription, NumberKeysMixin, Generic[T] ): """Describes UniFi Protect Number entity.""" -def _get_pir_duration(obj: Any) -> int: - assert isinstance(obj, Light) +def _get_pir_duration(obj: Light) -> int: return int(obj.light_device_settings.pir_duration.total_seconds()) -async def _set_pir_duration(obj: Any, value: float) -> None: - assert isinstance(obj, Light) +async def _set_pir_duration(obj: Light, value: float) -> None: await obj.set_duration(timedelta(seconds=value)) @@ -97,7 +95,7 @@ LIGHT_NUMBERS: tuple[ProtectNumberEntityDescription, ...] = ( ufp_value="light_device_settings.pir_sensitivity", ufp_set_method="set_sensitivity", ), - ProtectNumberEntityDescription( + ProtectNumberEntityDescription[Light]( key="duration", name="Auto-shutoff Duration", icon="mdi:camera-timer", diff --git a/homeassistant/components/unifiprotect/select.py b/homeassistant/components/unifiprotect/select.py index b2995736eeb7..83033a19ff46 100644 --- a/homeassistant/components/unifiprotect/select.py +++ b/homeassistant/components/unifiprotect/select.py @@ -6,7 +6,7 @@ from dataclasses import dataclass from datetime import timedelta from enum import Enum import logging -from typing import Any, Final +from typing import Any, Final, Generic from pyunifiprotect.api import ProtectApiClient from pyunifiprotect.data import ( @@ -19,9 +19,7 @@ from pyunifiprotect.data import ( RecordingMode, Viewer, ) -from pyunifiprotect.data.base import ProtectAdoptableDeviceModel from pyunifiprotect.data.devices import Sensor -from pyunifiprotect.data.nvr import NVR from pyunifiprotect.data.types import ChimeType, MountType import voluptuous as vol @@ -37,7 +35,7 @@ from homeassistant.util.dt import utcnow from .const import ATTR_DURATION, ATTR_MESSAGE, DOMAIN, TYPE_EMPTY_VALUE from .data import ProtectData from .entity import ProtectDeviceEntity, async_all_device_entities -from .models import ProtectSetableKeysMixin +from .models import ProtectSetableKeysMixin, T _LOGGER = logging.getLogger(__name__) _KEY_LIGHT_MOTION = "light_motion" @@ -104,15 +102,14 @@ SET_DOORBELL_LCD_MESSAGE_SCHEMA = vol.Schema( @dataclass -class ProtectSelectEntityDescription(ProtectSetableKeysMixin, SelectEntityDescription): +class ProtectSelectEntityDescription( + ProtectSetableKeysMixin, SelectEntityDescription, Generic[T] +): """Describes UniFi Protect Select entity.""" ufp_options: list[dict[str, Any]] | None = None - ufp_options_callable: Callable[ - [ProtectApiClient], list[dict[str, Any]] - ] | None = None + ufp_options_fn: Callable[[ProtectApiClient], list[dict[str, Any]]] | None = None ufp_enum_type: type[Enum] | None = None - ufp_set_method: str | None = None def _get_viewer_options(api: ProtectApiClient) -> list[dict[str, Any]]: @@ -140,13 +137,11 @@ def _get_paired_camera_options(api: ProtectApiClient) -> list[dict[str, Any]]: return options -def _get_viewer_current(obj: Any) -> str: - assert isinstance(obj, Viewer) +def _get_viewer_current(obj: Viewer) -> str: return obj.liveview_id -def _get_light_motion_current(obj: Any) -> str: - assert isinstance(obj, Light) +def _get_light_motion_current(obj: Light) -> str: # a bit of extra to allow On Motion Always/Dark if ( obj.light_mode_settings.mode == LightModeType.MOTION @@ -156,15 +151,13 @@ def _get_light_motion_current(obj: Any) -> str: return obj.light_mode_settings.mode.value -def _get_doorbell_current(obj: Any) -> str | None: - assert isinstance(obj, Camera) +def _get_doorbell_current(obj: Camera) -> str | None: if obj.lcd_message is None: return None return obj.lcd_message.text -async def _set_light_mode(obj: Any, mode: str) -> None: - assert isinstance(obj, Light) +async def _set_light_mode(obj: Light, mode: str) -> None: lightmode, timing = LIGHT_MODE_TO_SETTINGS[mode] await obj.set_light_settings( LightModeType(lightmode), @@ -172,10 +165,7 @@ async def _set_light_mode(obj: Any, mode: str) -> None: ) -async def _set_paired_camera( - obj: ProtectAdoptableDeviceModel | NVR, camera_id: str -) -> None: - assert isinstance(obj, (Sensor, Light)) +async def _set_paired_camera(obj: Light | Sensor, camera_id: str) -> None: if camera_id == TYPE_EMPTY_VALUE: camera: Camera | None = None else: @@ -183,8 +173,7 @@ async def _set_paired_camera( await obj.set_paired_camera(camera) -async def _set_doorbell_message(obj: Any, message: str) -> None: - assert isinstance(obj, Camera) +async def _set_doorbell_message(obj: Camera, message: str) -> None: if message.startswith(DoorbellMessageType.CUSTOM_MESSAGE.value): await obj.set_lcd_text(DoorbellMessageType.CUSTOM_MESSAGE, text=message) elif message == TYPE_EMPTY_VALUE: @@ -193,8 +182,7 @@ async def _set_doorbell_message(obj: Any, message: str) -> None: await obj.set_lcd_text(DoorbellMessageType(message)) -async def _set_liveview(obj: Any, liveview_id: str) -> None: - assert isinstance(obj, Viewer) +async def _set_liveview(obj: Viewer, liveview_id: str) -> None: liveview = obj.api.bootstrap.liveviews[liveview_id] await obj.set_liveview(liveview) @@ -221,7 +209,7 @@ CAMERA_SELECTS: tuple[ProtectSelectEntityDescription, ...] = ( ufp_value="isp_settings.ir_led_mode", ufp_set_method="set_ir_led_model", ), - ProtectSelectEntityDescription( + ProtectSelectEntityDescription[Camera]( key="doorbell_text", name="Doorbell Text", icon="mdi:card-text", @@ -229,7 +217,7 @@ CAMERA_SELECTS: tuple[ProtectSelectEntityDescription, ...] = ( device_class=DEVICE_CLASS_LCD_MESSAGE, ufp_required_field="feature_flags.has_lcd_screen", ufp_value_fn=_get_doorbell_current, - ufp_options_callable=_get_doorbell_options, + ufp_options_fn=_get_doorbell_options, ufp_set_method_fn=_set_doorbell_message, ), ProtectSelectEntityDescription( @@ -246,7 +234,7 @@ CAMERA_SELECTS: tuple[ProtectSelectEntityDescription, ...] = ( ) LIGHT_SELECTS: tuple[ProtectSelectEntityDescription, ...] = ( - ProtectSelectEntityDescription( + ProtectSelectEntityDescription[Light]( key=_KEY_LIGHT_MOTION, name="Light Mode", icon="mdi:spotlight", @@ -255,13 +243,13 @@ LIGHT_SELECTS: tuple[ProtectSelectEntityDescription, ...] = ( ufp_value_fn=_get_light_motion_current, ufp_set_method_fn=_set_light_mode, ), - ProtectSelectEntityDescription( + ProtectSelectEntityDescription[Light]( key="paired_camera", name="Paired Camera", icon="mdi:cctv", entity_category=EntityCategory.CONFIG, ufp_value="camera_id", - ufp_options_callable=_get_paired_camera_options, + ufp_options_fn=_get_paired_camera_options, ufp_set_method_fn=_set_paired_camera, ), ) @@ -277,24 +265,24 @@ SENSE_SELECTS: tuple[ProtectSelectEntityDescription, ...] = ( ufp_value="mount_type", ufp_set_method="set_mount_type", ), - ProtectSelectEntityDescription( + ProtectSelectEntityDescription[Sensor]( key="paired_camera", name="Paired Camera", icon="mdi:cctv", entity_category=EntityCategory.CONFIG, ufp_value="camera_id", - ufp_options_callable=_get_paired_camera_options, + ufp_options_fn=_get_paired_camera_options, ufp_set_method_fn=_set_paired_camera, ), ) VIEWER_SELECTS: tuple[ProtectSelectEntityDescription, ...] = ( - ProtectSelectEntityDescription( + ProtectSelectEntityDescription[Viewer]( key="viewer", name="Liveview", icon="mdi:view-dashboard", entity_category=None, - ufp_options_callable=_get_viewer_options, + ufp_options_fn=_get_viewer_options, ufp_value_fn=_get_viewer_current, ufp_set_method_fn=_set_liveview, ), @@ -350,7 +338,7 @@ class ProtectSelects(ProtectDeviceEntity, SelectEntity): # entities with categories are not exposed for voice and safe to update dynamically if ( self.entity_description.entity_category is not None - and self.entity_description.ufp_options_callable is not None + and self.entity_description.ufp_options_fn is not None ): _LOGGER.debug( "Updating dynamic select options for %s", self.entity_description.name @@ -364,8 +352,8 @@ class ProtectSelects(ProtectDeviceEntity, SelectEntity): if self.entity_description.ufp_options is not None: options = self.entity_description.ufp_options else: - assert self.entity_description.ufp_options_callable is not None - options = self.entity_description.ufp_options_callable(self.data.api) + assert self.entity_description.ufp_options_fn is not None + options = self.entity_description.ufp_options_fn(self.data.api) self._attr_options = [item["name"] for item in options] self._hass_to_unifi_options = {item["name"]: item["id"] for item in options} diff --git a/homeassistant/components/unifiprotect/sensor.py b/homeassistant/components/unifiprotect/sensor.py index b51e2c0475b1..29be7ef96d6a 100644 --- a/homeassistant/components/unifiprotect/sensor.py +++ b/homeassistant/components/unifiprotect/sensor.py @@ -4,11 +4,16 @@ from __future__ import annotations from dataclasses import dataclass from datetime import datetime import logging -from typing import Any +from typing import Any, Generic -from pyunifiprotect.data import NVR, Camera, Event -from pyunifiprotect.data.base import ProtectAdoptableDeviceModel -from pyunifiprotect.data.devices import Sensor +from pyunifiprotect.data import ( + NVR, + Camera, + Event, + ProtectAdoptableDeviceModel, + ProtectDeviceModel, + Sensor, +) from homeassistant.components.sensor import ( SensorDeviceClass, @@ -40,7 +45,7 @@ from .entity import ( ProtectNVREntity, async_all_device_entities, ) -from .models import ProtectRequiredKeysMixin +from .models import ProtectRequiredKeysMixin, T _LOGGER = logging.getLogger(__name__) OBJECT_TYPE_NONE = "none" @@ -48,12 +53,14 @@ DEVICE_CLASS_DETECTION = "unifiprotect__detection" @dataclass -class ProtectSensorEntityDescription(ProtectRequiredKeysMixin, SensorEntityDescription): +class ProtectSensorEntityDescription( + ProtectRequiredKeysMixin, SensorEntityDescription, Generic[T] +): """Describes UniFi Protect Sensor entity.""" precision: int | None = None - def get_ufp_value(self, obj: ProtectAdoptableDeviceModel | NVR) -> Any: + def get_ufp_value(self, obj: ProtectDeviceModel) -> Any: """Return value from UniFi Protect device.""" value = super().get_ufp_value(obj) @@ -62,7 +69,7 @@ class ProtectSensorEntityDescription(ProtectRequiredKeysMixin, SensorEntityDescr return value -def _get_uptime(obj: ProtectAdoptableDeviceModel | NVR) -> datetime | None: +def _get_uptime(obj: ProtectDeviceModel) -> datetime | None: if obj.up_since is None: return None @@ -71,26 +78,21 @@ def _get_uptime(obj: ProtectAdoptableDeviceModel | NVR) -> datetime | None: return obj.up_since.replace(second=0, microsecond=0) -def _get_nvr_recording_capacity(obj: Any) -> int: - assert isinstance(obj, NVR) - +def _get_nvr_recording_capacity(obj: NVR) -> int: if obj.storage_stats.capacity is None: return 0 return int(obj.storage_stats.capacity.total_seconds()) -def _get_nvr_memory(obj: Any) -> float | None: - assert isinstance(obj, NVR) - +def _get_nvr_memory(obj: NVR) -> float | None: memory = obj.system_info.memory if memory.available is None or memory.total is None: return None return (1 - memory.available / memory.total) * 100 -def _get_alarm_sound(obj: ProtectAdoptableDeviceModel | NVR) -> str: - assert isinstance(obj, Sensor) +def _get_alarm_sound(obj: Sensor) -> str: alarm_type = OBJECT_TYPE_NONE if ( @@ -103,7 +105,7 @@ def _get_alarm_sound(obj: ProtectAdoptableDeviceModel | NVR) -> str: ALL_DEVICES_SENSORS: tuple[ProtectSensorEntityDescription, ...] = ( - ProtectSensorEntityDescription( + ProtectSensorEntityDescription[ProtectDeviceModel]( key="uptime", name="Uptime", icon="mdi:clock", @@ -244,7 +246,7 @@ SENSE_SENSORS: tuple[ProtectSensorEntityDescription, ...] = ( ufp_value="stats.temperature.value", ufp_enabled="is_temperature_sensor_enabled", ), - ProtectSensorEntityDescription( + ProtectSensorEntityDescription[Sensor]( key="alarm_sound", name="Alarm Sound Detected", ufp_value_fn=_get_alarm_sound, @@ -253,7 +255,7 @@ SENSE_SENSORS: tuple[ProtectSensorEntityDescription, ...] = ( ) NVR_SENSORS: tuple[ProtectSensorEntityDescription, ...] = ( - ProtectSensorEntityDescription( + ProtectSensorEntityDescription[ProtectDeviceModel]( key="uptime", name="Uptime", icon="mdi:clock", @@ -331,7 +333,7 @@ NVR_SENSORS: tuple[ProtectSensorEntityDescription, ...] = ( ufp_value="storage_stats.storage_distribution.free.percentage", precision=2, ), - ProtectSensorEntityDescription( + ProtectSensorEntityDescription[NVR]( key="record_capacity", name="Recording Capacity", native_unit_of_measurement=TIME_SECONDS, @@ -363,7 +365,7 @@ NVR_DISABLED_SENSORS: tuple[ProtectSensorEntityDescription, ...] = ( state_class=SensorStateClass.MEASUREMENT, ufp_value="system_info.cpu.temperature", ), - ProtectSensorEntityDescription( + ProtectSensorEntityDescription[NVR]( key="memory_utilization", name="Memory Utilization", native_unit_of_measurement=PERCENTAGE, diff --git a/homeassistant/components/unifiprotect/switch.py b/homeassistant/components/unifiprotect/switch.py index d2a2267f4885..7922dfbc19f2 100644 --- a/homeassistant/components/unifiprotect/switch.py +++ b/homeassistant/components/unifiprotect/switch.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass import logging -from typing import Any +from typing import Any, Generic from pyunifiprotect.data import Camera, RecordingMode, VideoMode from pyunifiprotect.data.base import ProtectAdoptableDeviceModel @@ -17,26 +17,26 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from .const import DOMAIN from .data import ProtectData from .entity import ProtectDeviceEntity, async_all_device_entities -from .models import ProtectSetableKeysMixin +from .models import ProtectSetableKeysMixin, T _LOGGER = logging.getLogger(__name__) @dataclass -class ProtectSwitchEntityDescription(ProtectSetableKeysMixin, SwitchEntityDescription): +class ProtectSwitchEntityDescription( + ProtectSetableKeysMixin, SwitchEntityDescription, Generic[T] +): """Describes UniFi Protect Switch entity.""" _KEY_PRIVACY_MODE = "privacy_mode" -def _get_is_highfps(obj: Any) -> bool: - assert isinstance(obj, Camera) +def _get_is_highfps(obj: Camera) -> bool: return bool(obj.video_mode == VideoMode.HIGH_FPS) -async def _set_highfps(obj: Any, value: bool) -> None: - assert isinstance(obj, Camera) +async def _set_highfps(obj: Camera, value: bool) -> None: if value: await obj.set_video_mode(VideoMode.HIGH_FPS) else: @@ -74,7 +74,7 @@ CAMERA_SWITCHES: tuple[ProtectSwitchEntityDescription, ...] = ( ufp_value="hdr_mode", ufp_set_method="set_hdr", ), - ProtectSwitchEntityDescription( + ProtectSwitchEntityDescription[Camera]( key="high_fps", name="High FPS", icon="mdi:video-high-definition",