From a2f1b882279e96d05306d7aad0534e2501136433 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Tue, 30 Aug 2022 18:14:06 +0200 Subject: [PATCH] Use generics in litterrobot (#77537) --- .../components/litterrobot/button.py | 3 +- .../components/litterrobot/entity.py | 22 ++++---- .../components/litterrobot/select.py | 4 +- .../components/litterrobot/sensor.py | 51 +++++++++---------- .../components/litterrobot/switch.py | 8 ++- .../components/litterrobot/vacuum.py | 3 +- 6 files changed, 46 insertions(+), 45 deletions(-) diff --git a/homeassistant/components/litterrobot/button.py b/homeassistant/components/litterrobot/button.py index 7c7990edf07..b833500ec4c 100644 --- a/homeassistant/components/litterrobot/button.py +++ b/homeassistant/components/litterrobot/button.py @@ -32,10 +32,9 @@ async def async_setup_entry( ) -class LitterRobotResetWasteDrawerButton(LitterRobotEntity, ButtonEntity): +class LitterRobotResetWasteDrawerButton(LitterRobotEntity[LitterRobot3], ButtonEntity): """Litter-Robot reset waste drawer button.""" - robot: LitterRobot3 _attr_icon = "mdi:delete-variant" _attr_entity_category = EntityCategory.CONFIG diff --git a/homeassistant/components/litterrobot/entity.py b/homeassistant/components/litterrobot/entity.py index 7248cbe9315..8471e007ce9 100644 --- a/homeassistant/components/litterrobot/entity.py +++ b/homeassistant/components/litterrobot/entity.py @@ -4,9 +4,9 @@ from __future__ import annotations from collections.abc import Callable, Coroutine from datetime import time import logging -from typing import Any +from typing import Any, Generic, TypeVar -from pylitterbot import LitterRobot, Robot +from pylitterbot import Robot from pylitterbot.exceptions import InvalidCommandException from typing_extensions import ParamSpec @@ -23,17 +23,20 @@ from .const import DOMAIN from .hub import LitterRobotHub _P = ParamSpec("_P") +_RobotT = TypeVar("_RobotT", bound=Robot) _LOGGER = logging.getLogger(__name__) REFRESH_WAIT_TIME_SECONDS = 8 -class LitterRobotEntity(CoordinatorEntity[DataUpdateCoordinator[bool]]): +class LitterRobotEntity( + CoordinatorEntity[DataUpdateCoordinator[bool]], Generic[_RobotT] +): """Generic Litter-Robot entity representing common data and methods.""" _attr_has_entity_name = True - def __init__(self, robot: Robot, entity_type: str, hub: LitterRobotHub) -> None: + def __init__(self, robot: _RobotT, entity_type: str, hub: LitterRobotHub) -> None: """Pass coordinator to CoordinatorEntity.""" super().__init__(hub.coordinator) self.robot = robot @@ -59,12 +62,10 @@ class LitterRobotEntity(CoordinatorEntity[DataUpdateCoordinator[bool]]): ) -class LitterRobotControlEntity(LitterRobotEntity): +class LitterRobotControlEntity(LitterRobotEntity[_RobotT]): """A Litter-Robot entity that can control the unit.""" - robot: LitterRobot - - def __init__(self, robot: Robot, entity_type: str, hub: LitterRobotHub) -> None: + def __init__(self, robot: _RobotT, entity_type: str, hub: LitterRobotHub) -> None: """Init a Litter-Robot control entity.""" super().__init__(robot=robot, entity_type=entity_type, hub=hub) self._refresh_callback: CALLBACK_TYPE | None = None @@ -128,13 +129,12 @@ class LitterRobotControlEntity(LitterRobotEntity): ) -class LitterRobotConfigEntity(LitterRobotControlEntity): +class LitterRobotConfigEntity(LitterRobotControlEntity[_RobotT]): """A Litter-Robot entity that can control configuration of the unit.""" - robot: LitterRobot _attr_entity_category = EntityCategory.CONFIG - def __init__(self, robot: Robot, entity_type: str, hub: LitterRobotHub) -> None: + def __init__(self, robot: _RobotT, entity_type: str, hub: LitterRobotHub) -> None: """Init a Litter-Robot control entity.""" super().__init__(robot=robot, entity_type=entity_type, hub=hub) self._assumed_state: bool | None = None diff --git a/homeassistant/components/litterrobot/select.py b/homeassistant/components/litterrobot/select.py index 403f7a8c257..2889499f1c4 100644 --- a/homeassistant/components/litterrobot/select.py +++ b/homeassistant/components/litterrobot/select.py @@ -1,6 +1,8 @@ """Support for Litter-Robot selects.""" from __future__ import annotations +from pylitterbot import LitterRobot + from homeassistant.components.select import SelectEntity from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant @@ -29,7 +31,7 @@ async def async_setup_entry( ) -class LitterRobotSelect(LitterRobotConfigEntity, SelectEntity): +class LitterRobotSelect(LitterRobotConfigEntity[LitterRobot], SelectEntity): """Litter-Robot Select.""" _attr_icon = "mdi:timer-outline" diff --git a/homeassistant/components/litterrobot/sensor.py b/homeassistant/components/litterrobot/sensor.py index 5f6579e83c5..386d0e04f3c 100644 --- a/homeassistant/components/litterrobot/sensor.py +++ b/homeassistant/components/litterrobot/sensor.py @@ -1,12 +1,13 @@ """Support for Litter-Robot sensors.""" from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Iterable from dataclasses import dataclass from datetime import datetime -from typing import Any, Union, cast +import itertools +from typing import Any, Generic, Union, cast -from pylitterbot import FeederRobot, LitterRobot, Robot +from pylitterbot import FeederRobot, LitterRobot from homeassistant.components.sensor import ( SensorDeviceClass, @@ -20,7 +21,7 @@ from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback from .const import DOMAIN -from .entity import LitterRobotEntity +from .entity import LitterRobotEntity, _RobotT from .hub import LitterRobotHub @@ -36,30 +37,23 @@ def icon_for_gauge_level(gauge_level: int | None = None, offset: int = 0) -> str @dataclass -class RobotSensorEntityDescription(SensorEntityDescription): +class RobotSensorEntityDescription(SensorEntityDescription, Generic[_RobotT]): """A class that describes robot sensor entities.""" icon_fn: Callable[[Any], str | None] = lambda _: None - should_report: Callable[[Robot], bool] = lambda _: True + should_report: Callable[[_RobotT], bool] = lambda _: True -@dataclass -class LitterRobotSensorEntityDescription(RobotSensorEntityDescription): - """A class that describes Litter-Robot sensor entities.""" - - should_report: Callable[[LitterRobot], bool] = lambda _: True - - -class LitterRobotSensorEntity(LitterRobotEntity, SensorEntity): +class LitterRobotSensorEntity(LitterRobotEntity[_RobotT], SensorEntity): """Litter-Robot sensor entity.""" - entity_description: RobotSensorEntityDescription + entity_description: RobotSensorEntityDescription[_RobotT] def __init__( self, - robot: LitterRobot | FeederRobot, + robot: _RobotT, hub: LitterRobotHub, - description: RobotSensorEntityDescription, + description: RobotSensorEntityDescription[_RobotT], ) -> None: """Initialize a Litter-Robot sensor entity.""" assert description.name @@ -84,31 +78,31 @@ class LitterRobotSensorEntity(LitterRobotEntity, SensorEntity): LITTER_ROBOT_SENSORS = [ - LitterRobotSensorEntityDescription( + RobotSensorEntityDescription[LitterRobot]( name="Waste Drawer", key="waste_drawer_level", native_unit_of_measurement=PERCENTAGE, icon_fn=lambda state: icon_for_gauge_level(state, 10), ), - LitterRobotSensorEntityDescription( + RobotSensorEntityDescription[LitterRobot]( name="Sleep Mode Start Time", key="sleep_mode_start_time", device_class=SensorDeviceClass.TIMESTAMP, should_report=lambda robot: robot.sleep_mode_enabled, ), - LitterRobotSensorEntityDescription( + RobotSensorEntityDescription[LitterRobot]( name="Sleep Mode End Time", key="sleep_mode_end_time", device_class=SensorDeviceClass.TIMESTAMP, should_report=lambda robot: robot.sleep_mode_enabled, ), - LitterRobotSensorEntityDescription( + RobotSensorEntityDescription[LitterRobot]( name="Last Seen", key="last_seen", device_class=SensorDeviceClass.TIMESTAMP, entity_category=EntityCategory.DIAGNOSTIC, ), - LitterRobotSensorEntityDescription( + RobotSensorEntityDescription[LitterRobot]( name="Status Code", key="status_code", device_class="litterrobot__status_code", @@ -116,7 +110,7 @@ LITTER_ROBOT_SENSORS = [ ), ] -FEEDER_ROBOT_SENSOR = RobotSensorEntityDescription( +FEEDER_ROBOT_SENSOR = RobotSensorEntityDescription[FeederRobot]( name="Food Level", key="food_level", native_unit_of_measurement=PERCENTAGE, @@ -131,16 +125,17 @@ async def async_setup_entry( ) -> None: """Set up Litter-Robot sensors using config entry.""" hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id] - async_add_entities( - [ + entities: Iterable[LitterRobotSensorEntity] = itertools.chain( + ( LitterRobotSensorEntity(robot=robot, hub=hub, description=description) for description in LITTER_ROBOT_SENSORS for robot in hub.litter_robots() - ] - + [ + ), + ( LitterRobotSensorEntity( robot=robot, hub=hub, description=FEEDER_ROBOT_SENSOR ) for robot in hub.feeder_robots() - ] + ), ) + async_add_entities(entities) diff --git a/homeassistant/components/litterrobot/switch.py b/homeassistant/components/litterrobot/switch.py index 69050057050..de401576a78 100644 --- a/homeassistant/components/litterrobot/switch.py +++ b/homeassistant/components/litterrobot/switch.py @@ -3,6 +3,8 @@ from __future__ import annotations from typing import Any +from pylitterbot import LitterRobot + from homeassistant.components.switch import SwitchEntity from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant @@ -13,7 +15,9 @@ from .entity import LitterRobotConfigEntity from .hub import LitterRobotHub -class LitterRobotNightLightModeSwitch(LitterRobotConfigEntity, SwitchEntity): +class LitterRobotNightLightModeSwitch( + LitterRobotConfigEntity[LitterRobot], SwitchEntity +): """Litter-Robot Night Light Mode Switch.""" @property @@ -37,7 +41,7 @@ class LitterRobotNightLightModeSwitch(LitterRobotConfigEntity, SwitchEntity): await self.perform_action_and_assume_state(self.robot.set_night_light, False) -class LitterRobotPanelLockoutSwitch(LitterRobotConfigEntity, SwitchEntity): +class LitterRobotPanelLockoutSwitch(LitterRobotConfigEntity[LitterRobot], SwitchEntity): """Litter-Robot Panel Lockout Switch.""" @property diff --git a/homeassistant/components/litterrobot/vacuum.py b/homeassistant/components/litterrobot/vacuum.py index 11a437e893c..9a4b825045f 100644 --- a/homeassistant/components/litterrobot/vacuum.py +++ b/homeassistant/components/litterrobot/vacuum.py @@ -4,6 +4,7 @@ from __future__ import annotations import logging from typing import Any +from pylitterbot import LitterRobot from pylitterbot.enums import LitterBoxStatus import voluptuous as vol @@ -68,7 +69,7 @@ async def async_setup_entry( ) -class LitterRobotCleaner(LitterRobotControlEntity, StateVacuumEntity): +class LitterRobotCleaner(LitterRobotControlEntity[LitterRobot], StateVacuumEntity): """Litter-Robot "Vacuum" Cleaner.""" _attr_supported_features = (