diff --git a/homeassistant/components/litterrobot/button.py b/homeassistant/components/litterrobot/button.py index 6e6cc563c8e..984b28cc96e 100644 --- a/homeassistant/components/litterrobot/button.py +++ b/homeassistant/components/litterrobot/button.py @@ -4,10 +4,9 @@ from __future__ import annotations from collections.abc import Callable, Coroutine from dataclasses import dataclass -import itertools from typing import Any, Generic -from pylitterbot import FeederRobot, LitterRobot3 +from pylitterbot import FeederRobot, LitterRobot3, LitterRobot4, Robot from homeassistant.components.button import ButtonEntity, ButtonEntityDescription from homeassistant.const import EntityCategory @@ -18,6 +17,34 @@ from . import LitterRobotConfigEntry from .entity import LitterRobotEntity, _RobotT +@dataclass(frozen=True, kw_only=True) +class RobotButtonEntityDescription(ButtonEntityDescription, Generic[_RobotT]): + """A class that describes robot button entities.""" + + press_fn: Callable[[_RobotT], Coroutine[Any, Any, bool]] + + +ROBOT_BUTTON_MAP: dict[type[Robot], RobotButtonEntityDescription] = { + LitterRobot3: RobotButtonEntityDescription[LitterRobot3]( + key="reset_waste_drawer", + translation_key="reset_waste_drawer", + entity_category=EntityCategory.CONFIG, + press_fn=lambda robot: robot.reset_waste_drawer(), + ), + LitterRobot4: RobotButtonEntityDescription[LitterRobot4]( + key="reset", + translation_key="reset", + entity_category=EntityCategory.CONFIG, + press_fn=lambda robot: robot.reset(), + ), + FeederRobot: RobotButtonEntityDescription[FeederRobot]( + key="give_snack", + translation_key="give_snack", + press_fn=lambda robot: robot.give_snack(), + ), +} + + async def async_setup_entry( hass: HomeAssistant, entry: LitterRobotConfigEntry, @@ -25,51 +52,15 @@ async def async_setup_entry( ) -> None: """Set up Litter-Robot cleaner using config entry.""" hub = entry.runtime_data - entities: list[LitterRobotButtonEntity] = list( - itertools.chain( - ( - LitterRobotButtonEntity( - robot=robot, hub=hub, description=LITTER_ROBOT_BUTTON - ) - for robot in hub.litter_robots() - if isinstance(robot, LitterRobot3) - ), - ( - LitterRobotButtonEntity( - robot=robot, hub=hub, description=FEEDER_ROBOT_BUTTON - ) - for robot in hub.feeder_robots() - ), - ) - ) + entities = [ + LitterRobotButtonEntity(robot=robot, hub=hub, description=description) + for robot in hub.account.robots + for robot_type, description in ROBOT_BUTTON_MAP.items() + if isinstance(robot, robot_type) + ] async_add_entities(entities) -@dataclass(frozen=True) -class RequiredKeysMixin(Generic[_RobotT]): - """A class that describes robot button entity required keys.""" - - press_fn: Callable[[_RobotT], Coroutine[Any, Any, bool]] - - -@dataclass(frozen=True) -class RobotButtonEntityDescription(ButtonEntityDescription, RequiredKeysMixin[_RobotT]): - """A class that describes robot button entities.""" - - -LITTER_ROBOT_BUTTON = RobotButtonEntityDescription[LitterRobot3]( - key="reset_waste_drawer", - translation_key="reset_waste_drawer", - entity_category=EntityCategory.CONFIG, - press_fn=lambda robot: robot.reset_waste_drawer(), -) -FEEDER_ROBOT_BUTTON = RobotButtonEntityDescription[FeederRobot]( - key="give_snack", - translation_key="give_snack", - press_fn=lambda robot: robot.give_snack(), -) - - class LitterRobotButtonEntity(LitterRobotEntity[_RobotT], ButtonEntity): """Litter-Robot button entity.""" diff --git a/homeassistant/components/litterrobot/strings.json b/homeassistant/components/litterrobot/strings.json index 7acfad69735..3b6e2f01ef9 100644 --- a/homeassistant/components/litterrobot/strings.json +++ b/homeassistant/components/litterrobot/strings.json @@ -38,6 +38,9 @@ } }, "button": { + "reset": { + "name": "Reset" + }, "reset_waste_drawer": { "name": "Reset waste drawer" },