Add mopping abilities to Roborock (#91766)

Co-authored-by: Franck Nijhof <frenck@frenck.nl>

homeassistant/components/roborock/select.py

@@ -0,0 +1,116 @@
+"""Support for Roborock select."""
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from roborock.code_mappings import RoborockMopIntensityCode, RoborockMopModeCode
+from roborock.containers import Status
+from roborock.exceptions import RoborockException
+from roborock.typing import RoborockCommand
+
+from homeassistant.components.select import SelectEntity, SelectEntityDescription
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.util import slugify
+
+from .const import DOMAIN
+from .coordinator import RoborockDataUpdateCoordinator
+from .device import RoborockCoordinatedEntity
+from .models import RoborockHassDeviceInfo
+
+
+@dataclass
+class RoborockSelectDescriptionMixin:
+    """Define an entity description mixin for select entities."""
+
+    api_command: RoborockCommand
+    value_fn: Callable[[Status], str]
+    options_lambda: Callable[[str], list[int]]
+
+
+@dataclass
+class RoborockSelectDescription(
+    SelectEntityDescription, RoborockSelectDescriptionMixin
+):
+    """Class to describe an Roborock select entity."""
+
+
+SELECT_DESCRIPTIONS: list[RoborockSelectDescription] = [
+    RoborockSelectDescription(
+        key="water_box_mode",
+        translation_key="mop_intensity",
+        options=RoborockMopIntensityCode.values(),
+        api_command=RoborockCommand.SET_WATER_BOX_CUSTOM_MODE,
+        value_fn=lambda data: data.water_box_mode,
+        options_lambda=lambda data: [
+            k for k, v in RoborockMopIntensityCode.items() if v == data
+        ],
+    ),
+    RoborockSelectDescription(
+        key="mop_mode",
+        translation_key="mop_mode",
+        options=RoborockMopModeCode.values(),
+        api_command=RoborockCommand.SET_MOP_MODE,
+        value_fn=lambda data: data.mop_mode,
+        options_lambda=lambda data: [
+            k for k, v in RoborockMopModeCode.items() if v == data
+        ],
+    ),
+]
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up Roborock select platform."""
+
+    coordinator: RoborockDataUpdateCoordinator = hass.data[DOMAIN][
+        config_entry.entry_id
+    ]
+    async_add_entities(
+        RoborockSelectEntity(
+            f"{description.key}_{slugify(device_id)}",
+            device_info,
+            coordinator,
+            description,
+        )
+        for device_id, device_info in coordinator.devices_info.items()
+        for description in SELECT_DESCRIPTIONS
+    )
+
+
+class RoborockSelectEntity(RoborockCoordinatedEntity, SelectEntity):
+    """A class to let you set options on a Roborock vacuum where the potential options are fixed."""
+
+    entity_description: RoborockSelectDescription
+
+    def __init__(
+        self,
+        unique_id: str,
+        device_info: RoborockHassDeviceInfo,
+        coordinator: RoborockDataUpdateCoordinator,
+        entity_description: RoborockSelectDescription,
+    ) -> None:
+        """Create a select entity."""
+        self.entity_description = entity_description
+        super().__init__(unique_id, device_info, coordinator)
+
+    async def async_select_option(self, option: str) -> None:
+        """Set the mop intensity."""
+        try:
+            await self.send(
+                self.entity_description.api_command,
+                self.entity_description.options_lambda(option),
+            )
+        except RoborockException as err:
+            raise HomeAssistantError(
+                f"Error while setting {self.entity_description.key} to {option}"
+            ) from err
+
+    @property
+    def current_option(self) -> str | None:
+        """Get the current status of the select entity from device_status."""
+        return self.entity_description.value_fn(self._device_status)