Add button to reset Litter-Robot 4 (#136191)

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."""
 

homeassistant/components/litterrobot/strings.json

@@ -38,6 +38,9 @@
       }
     },
     "button": {
+      "reset": {
+        "name": "Reset"
+      },
       "reset_waste_drawer": {
         "name": "Reset waste drawer"
       },