From e192c99d2f993c12b3605aafe94e4d05febb7373 Mon Sep 17 00:00:00 2001 From: Nathan Spencer Date: Tue, 30 Aug 2022 22:03:02 -0600 Subject: [PATCH] Add support for Feeder-Robot switches (#77503) --- .../components/litterrobot/switch.py | 128 ++++++++++-------- 1 file changed, 71 insertions(+), 57 deletions(-) diff --git a/homeassistant/components/litterrobot/switch.py b/homeassistant/components/litterrobot/switch.py index de401576a78..2f54ede38b8 100644 --- a/homeassistant/components/litterrobot/switch.py +++ b/homeassistant/components/litterrobot/switch.py @@ -1,78 +1,91 @@ """Support for Litter-Robot switches.""" from __future__ import annotations -from typing import Any +from collections.abc import Callable, Coroutine +from dataclasses import dataclass +from typing import Any, Generic, Union -from pylitterbot import LitterRobot +from pylitterbot import FeederRobot, LitterRobot -from homeassistant.components.switch import SwitchEntity +from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from .const import DOMAIN -from .entity import LitterRobotConfigEntity +from .entity import LitterRobotConfigEntity, _RobotT from .hub import LitterRobotHub -class LitterRobotNightLightModeSwitch( - LitterRobotConfigEntity[LitterRobot], SwitchEntity -): - """Litter-Robot Night Light Mode Switch.""" +@dataclass +class RequiredKeysMixin(Generic[_RobotT]): + """A class that describes robot switch entity required keys.""" - @property - def is_on(self) -> bool | None: - """Return true if switch is on.""" - if self._refresh_callback is not None: - return self._assumed_state - return self.robot.night_light_mode_enabled - - @property - def icon(self) -> str: - """Return the icon.""" - return "mdi:lightbulb-on" if self.is_on else "mdi:lightbulb-off" - - async def async_turn_on(self, **kwargs: Any) -> None: - """Turn the switch on.""" - await self.perform_action_and_assume_state(self.robot.set_night_light, True) - - async def async_turn_off(self, **kwargs: Any) -> None: - """Turn the switch off.""" - await self.perform_action_and_assume_state(self.robot.set_night_light, False) + icons: tuple[str, str] + set_fn: Callable[[_RobotT], Callable[[bool], Coroutine[Any, Any, bool]]] -class LitterRobotPanelLockoutSwitch(LitterRobotConfigEntity[LitterRobot], SwitchEntity): - """Litter-Robot Panel Lockout Switch.""" - - @property - def is_on(self) -> bool | None: - """Return true if switch is on.""" - if self._refresh_callback is not None: - return self._assumed_state - return self.robot.panel_lock_enabled - - @property - def icon(self) -> str: - """Return the icon.""" - return "mdi:lock" if self.is_on else "mdi:lock-open" - - async def async_turn_on(self, **kwargs: Any) -> None: - """Turn the switch on.""" - await self.perform_action_and_assume_state(self.robot.set_panel_lockout, True) - - async def async_turn_off(self, **kwargs: Any) -> None: - """Turn the switch off.""" - await self.perform_action_and_assume_state(self.robot.set_panel_lockout, False) +@dataclass +class RobotSwitchEntityDescription(SwitchEntityDescription, RequiredKeysMixin[_RobotT]): + """A class that describes robot switch entities.""" -ROBOT_SWITCHES: list[ - tuple[type[LitterRobotNightLightModeSwitch | LitterRobotPanelLockoutSwitch], str] -] = [ - (LitterRobotNightLightModeSwitch, "Night Light Mode"), - (LitterRobotPanelLockoutSwitch, "Panel Lockout"), +ROBOT_SWITCHES = [ + RobotSwitchEntityDescription[Union[LitterRobot, FeederRobot]]( + key="night_light_mode_enabled", + name="Night Light Mode", + icons=("mdi:lightbulb-on", "mdi:lightbulb-off"), + set_fn=lambda robot: robot.set_night_light, + ), + RobotSwitchEntityDescription[Union[LitterRobot, FeederRobot]]( + key="panel_lock_enabled", + name="Panel Lockout", + icons=("mdi:lock", "mdi:lock-open"), + set_fn=lambda robot: robot.set_panel_lockout, + ), ] +class RobotSwitchEntity(LitterRobotConfigEntity[_RobotT], SwitchEntity): + """Litter-Robot switch entity.""" + + entity_description: RobotSwitchEntityDescription[_RobotT] + + def __init__( + self, + robot: _RobotT, + hub: LitterRobotHub, + description: RobotSwitchEntityDescription[_RobotT], + ) -> None: + """Initialize a Litter-Robot switch entity.""" + assert description.name + super().__init__(robot, description.name, hub) + self.entity_description = description + + @property + def is_on(self) -> bool | None: + """Return true if switch is on.""" + if self._refresh_callback is not None: + return self._assumed_state + return bool(getattr(self.robot, self.entity_description.key)) + + @property + def icon(self) -> str: + """Return the icon.""" + icon_on, icon_off = self.entity_description.icons + return icon_on if self.is_on else icon_off + + async def async_turn_on(self, **kwargs: Any) -> None: + """Turn the switch on.""" + set_fn = self.entity_description.set_fn + await self.perform_action_and_assume_state(set_fn(self.robot), True) + + async def async_turn_off(self, **kwargs: Any) -> None: + """Turn the switch off.""" + set_fn = self.entity_description.set_fn + await self.perform_action_and_assume_state(set_fn(self.robot), False) + + async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, @@ -81,7 +94,8 @@ async def async_setup_entry( """Set up Litter-Robot switches using config entry.""" hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id] async_add_entities( - switch_class(robot=robot, entity_type=switch_type, hub=hub) - for switch_class, switch_type in ROBOT_SWITCHES - for robot in hub.litter_robots() + RobotSwitchEntity(robot=robot, hub=hub, description=description) + for description in ROBOT_SWITCHES + for robot in hub.account.robots + if isinstance(robot, (LitterRobot, FeederRobot)) )