Use generics in litterrobot (#77537)

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
 

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

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"

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)

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

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 = (