From 69900ed8cb759a9e3cad867d25a593070686c8e6 Mon Sep 17 00:00:00 2001 From: Nathan Spencer Date: Tue, 21 Jan 2025 14:12:15 -0700 Subject: [PATCH] Cleanup litterrobot switch entity (#136199) --- .../components/litterrobot/switch.py | 53 +++++++++---------- 1 file changed, 25 insertions(+), 28 deletions(-) diff --git a/homeassistant/components/litterrobot/switch.py b/homeassistant/components/litterrobot/switch.py index 133fd897cc6..a73449b01a1 100644 --- a/homeassistant/components/litterrobot/switch.py +++ b/homeassistant/components/litterrobot/switch.py @@ -17,18 +17,13 @@ from . import LitterRobotConfigEntry from .entity import LitterRobotEntity, _RobotT -@dataclass(frozen=True) -class RequiredKeysMixin(Generic[_RobotT]): - """A class that describes robot switch entity required keys.""" - - set_fn: Callable[[_RobotT, bool], Coroutine[Any, Any, bool]] - - -@dataclass(frozen=True) -class RobotSwitchEntityDescription(SwitchEntityDescription, RequiredKeysMixin[_RobotT]): +@dataclass(frozen=True, kw_only=True) +class RobotSwitchEntityDescription(SwitchEntityDescription, Generic[_RobotT]): """A class that describes robot switch entities.""" entity_category: EntityCategory = EntityCategory.CONFIG + set_fn: Callable[[_RobotT, bool], Coroutine[Any, Any, bool]] + value_fn: Callable[[_RobotT], bool] ROBOT_SWITCHES = [ @@ -36,34 +31,17 @@ ROBOT_SWITCHES = [ key="night_light_mode_enabled", translation_key="night_light_mode", set_fn=lambda robot, value: robot.set_night_light(value), + value_fn=lambda robot: robot.night_light_mode_enabled, ), RobotSwitchEntityDescription[LitterRobot | FeederRobot]( key="panel_lock_enabled", translation_key="panel_lockout", set_fn=lambda robot, value: robot.set_panel_lockout(value), + value_fn=lambda robot: robot.panel_lock_enabled, ), ] -class RobotSwitchEntity(LitterRobotEntity[_RobotT], SwitchEntity): - """Litter-Robot switch entity.""" - - entity_description: RobotSwitchEntityDescription[_RobotT] - - @property - def is_on(self) -> bool | None: - """Return true if switch is on.""" - return bool(getattr(self.robot, self.entity_description.key)) - - async def async_turn_on(self, **kwargs: Any) -> None: - """Turn the switch on.""" - await self.entity_description.set_fn(self.robot, True) - - async def async_turn_off(self, **kwargs: Any) -> None: - """Turn the switch off.""" - await self.entity_description.set_fn(self.robot, False) - - async def async_setup_entry( hass: HomeAssistant, entry: LitterRobotConfigEntry, @@ -78,3 +56,22 @@ async def async_setup_entry( if isinstance(robot, (LitterRobot, FeederRobot)) ] async_add_entities(entities) + + +class RobotSwitchEntity(LitterRobotEntity[_RobotT], SwitchEntity): + """Litter-Robot switch entity.""" + + entity_description: RobotSwitchEntityDescription[_RobotT] + + @property + def is_on(self) -> bool | None: + """Return true if switch is on.""" + return self.entity_description.value_fn(self.robot) + + async def async_turn_on(self, **kwargs: Any) -> None: + """Turn the switch on.""" + await self.entity_description.set_fn(self.robot, True) + + async def async_turn_off(self, **kwargs: Any) -> None: + """Turn the switch off.""" + await self.entity_description.set_fn(self.robot, False)