From 5bc2f37bf8f2d6e21ae03d69756916fd68dd6f73 Mon Sep 17 00:00:00 2001
From: Nathan Spencer <natekspencer@gmail.com>
Date: Wed, 31 Aug 2022 05:23:51 -0600
Subject: [PATCH] Add support for Feeder-Robot select (#77512)

* Add support for Feeder-Robot select

* Use lambda to get current selected option

* Use generics and required keys mixin

* Code improvements

* Even more generics

* Fix missing type hint

* Apply suggestions from code review

Co-authored-by: epenet <6771947+epenet@users.noreply.github.com>

Co-authored-by: epenet <6771947+epenet@users.noreply.github.com>
---
 .../components/litterrobot/select.py          | 93 +++++++++++++++----
 1 file changed, 76 insertions(+), 17 deletions(-)

diff --git a/homeassistant/components/litterrobot/select.py b/homeassistant/components/litterrobot/select.py
index 2889499f1c4..a18cd3b46b5 100644
--- a/homeassistant/components/litterrobot/select.py
+++ b/homeassistant/components/litterrobot/select.py
@@ -1,18 +1,61 @@
 """Support for Litter-Robot selects."""
 from __future__ import annotations
 
-from pylitterbot import LitterRobot
+from collections.abc import Callable, Coroutine
+from dataclasses import dataclass
+import itertools
+from typing import Any, Generic, TypeVar
 
-from homeassistant.components.select import SelectEntity
+from pylitterbot import FeederRobot, LitterRobot
+
+from homeassistant.components.select import SelectEntity, SelectEntityDescription
 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
 
-TYPE_CLEAN_CYCLE_WAIT_TIME_MINUTES = "Clean Cycle Wait Time Minutes"
+_CastTypeT = TypeVar("_CastTypeT", int, float)
+
+
+@dataclass
+class RequiredKeysMixin(Generic[_RobotT, _CastTypeT]):
+    """A class that describes robot select entity required keys."""
+
+    current_fn: Callable[[_RobotT], _CastTypeT]
+    options_fn: Callable[[_RobotT], list[_CastTypeT]]
+    select_fn: Callable[
+        [_RobotT, str],
+        tuple[Callable[[_CastTypeT], Coroutine[Any, Any, bool]], _CastTypeT],
+    ]
+
+
+@dataclass
+class RobotSelectEntityDescription(
+    SelectEntityDescription, RequiredKeysMixin[_RobotT, _CastTypeT]
+):
+    """A class that describes robot select entities."""
+
+
+LITTER_ROBOT_SELECT = RobotSelectEntityDescription[LitterRobot, int](
+    key="clean_cycle_wait_time_minutes",
+    name="Clean Cycle Wait Time Minutes",
+    icon="mdi:timer-outline",
+    current_fn=lambda robot: robot.clean_cycle_wait_time_minutes,
+    options_fn=lambda robot: robot.VALID_WAIT_TIMES,
+    select_fn=lambda robot, option: (robot.set_wait_time, int(option)),
+)
+FEEDER_ROBOT_SELECT = RobotSelectEntityDescription[FeederRobot, float](
+    key="meal_insert_size",
+    name="Meal insert size",
+    icon="mdi:scale",
+    unit_of_measurement="cups",
+    current_fn=lambda robot: robot.meal_insert_size,
+    options_fn=lambda robot: robot.VALID_MEAL_INSERT_SIZES,
+    select_fn=lambda robot, option: (robot.set_meal_insert_size, float(option)),
+)
 
 
 async def async_setup_entry(
@@ -22,30 +65,46 @@ async def async_setup_entry(
 ) -> None:
     """Set up Litter-Robot selects using config entry."""
     hub: LitterRobotHub = hass.data[DOMAIN][config_entry.entry_id]
-
     async_add_entities(
-        LitterRobotSelect(
-            robot=robot, entity_type=TYPE_CLEAN_CYCLE_WAIT_TIME_MINUTES, hub=hub
+        itertools.chain(
+            (
+                LitterRobotSelect(robot=robot, hub=hub, description=LITTER_ROBOT_SELECT)
+                for robot in hub.litter_robots()
+            ),
+            (
+                LitterRobotSelect(robot=robot, hub=hub, description=FEEDER_ROBOT_SELECT)
+                for robot in hub.feeder_robots()
+            ),
         )
-        for robot in hub.litter_robots()
     )
 
 
-class LitterRobotSelect(LitterRobotConfigEntity[LitterRobot], SelectEntity):
+class LitterRobotSelect(
+    LitterRobotConfigEntity[_RobotT], SelectEntity, Generic[_RobotT, _CastTypeT]
+):
     """Litter-Robot Select."""
 
-    _attr_icon = "mdi:timer-outline"
+    entity_description: RobotSelectEntityDescription[_RobotT, _CastTypeT]
+
+    def __init__(
+        self,
+        robot: _RobotT,
+        hub: LitterRobotHub,
+        description: RobotSelectEntityDescription[_RobotT, _CastTypeT],
+    ) -> None:
+        """Initialize a Litter-Robot select entity."""
+        assert description.name
+        super().__init__(robot, description.name, hub)
+        self.entity_description = description
+        options = self.entity_description.options_fn(self.robot)
+        self._attr_options = list(map(str, options))
 
     @property
     def current_option(self) -> str | None:
         """Return the selected entity option to represent the entity state."""
-        return str(self.robot.clean_cycle_wait_time_minutes)
-
-    @property
-    def options(self) -> list[str]:
-        """Return a set of selectable options."""
-        return [str(minute) for minute in self.robot.VALID_WAIT_TIMES]
+        return str(self.entity_description.current_fn(self.robot))
 
     async def async_select_option(self, option: str) -> None:
         """Change the selected option."""
-        await self.perform_action_and_refresh(self.robot.set_wait_time, int(option))
+        action, adjusted_option = self.entity_description.select_fn(self.robot, option)
+        await self.perform_action_and_refresh(action, adjusted_option)