Use Capabilities (#447)

custom_components/deebot/binary_sensor.py

@@ -1,6 +1,9 @@
 """Binary sensor module."""
-import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Generic
 
+from deebot_client.capabilities import CapabilityEvent
 from deebot_client.events.water_info import WaterInfoEvent
 from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
@@ -13,50 +16,63 @@
 
 from .const import DOMAIN
 from .controller import DeebotController
-from .entity import DeebotEntity
+from .entity import DeebotEntity, DeebotEntityDescription, EventT
 
-_LOGGER = logging.getLogger(__name__)
 
+@dataclass
+class DeebotBinarySensorEntityMixin(Generic[EventT]):
+    """Deebot binary sensor entity mixin."""
 
-async def async_setup_entry(
-    hass: HomeAssistant,
-    config_entry: ConfigEntry,
-    async_add_entities: AddEntitiesCallback,
-) -> None:
-    """Add entities for passed config_entry in HA."""
-    controller: DeebotController = hass.data[DOMAIN][config_entry.entry_id]
-
-    new_devices = []
-    for vacbot in controller.vacuum_bots:
-        new_devices.append(DeebotMopAttachedBinarySensor(vacbot))
+    value_fn: Callable[[EventT], bool | None]
+    icon_fn: Callable[[bool | None], str | None]
 
-    if new_devices:
-        async_add_entities(new_devices)
 
+@dataclass
+class DeebotBinarySensorEntityDescription(
+    BinarySensorEntityDescription,  # type: ignore
+    DeebotEntityDescription,
+    DeebotBinarySensorEntityMixin[EventT],
+):
+    """Class describing Deebot binary sensor entity."""
 
-class DeebotMopAttachedBinarySensor(DeebotEntity, BinarySensorEntity):  # type: ignore
-    """Deebot mop attached binary sensor."""
 
-    entity_description = BinarySensorEntityDescription(
+ENTITY_DESCRIPTIONS: tuple[DeebotBinarySensorEntityDescription, ...] = (
+    DeebotBinarySensorEntityDescription[WaterInfoEvent](
+        capability_fn=lambda caps: caps.water,
+        value_fn=lambda e: e.mop_attached,
+        icon_fn=lambda is_on: "mdi:water" if is_on else "mdi:water-off",
         key="mop_attached",
         translation_key="mop_attached",
         entity_registry_enabled_default=False,
         entity_category=EntityCategory.DIAGNOSTIC,
+    ),
+)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Add entities for passed config_entry in HA."""
+    controller: DeebotController = hass.data[DOMAIN][config_entry.entry_id]
+    controller.register_platform_add_entities(
+        DeebotBinarySensor, ENTITY_DESCRIPTIONS, async_add_entities
     )
 
-    @property
-    def icon(self) -> str | None:
-        """Return the icon to use in the frontend, if any."""
-        return "mdi:water" if self.is_on else "mdi:water-off"
+
+class DeebotBinarySensor(DeebotEntity[CapabilityEvent[EventT], DeebotBinarySensorEntityDescription], BinarySensorEntity):  # type: ignore
+    """Deebot binary sensor."""
 
     async def async_added_to_hass(self) -> None:
         """Set up the event listeners now that hass is ready."""
         await super().async_added_to_hass()
 
-        async def on_event(event: WaterInfoEvent) -> None:
-            self._attr_is_on = event.mop_attached
+        async def on_event(event: EventT) -> None:
+            self._attr_is_on = self.entity_description.value_fn(event)
+            self._attr_icon = self.entity_description.icon_fn(self._attr_is_on)
             self.async_write_ha_state()
 
         self.async_on_remove(
-            self._vacuum_bot.events.subscribe(WaterInfoEvent, on_event)
+            self._vacuum_bot.events.subscribe(self._capability.event, on_event)
         )

custom_components/deebot/button.py

@@ -1,7 +1,8 @@
 """Binary sensor module."""
-import logging
+from collections.abc import Sequence
+from dataclasses import dataclass
 
-from deebot_client.commands import ResetLifeSpan, SetRelocationState
+from deebot_client.capabilities import CapabilityExecute
 from deebot_client.events import LifeSpan
 from deebot_client.vacuum_bot import VacuumBot
 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
@@ -12,9 +13,27 @@
 
 from .const import DOMAIN
 from .controller import DeebotController
-from .entity import DeebotEntity
+from .entity import DeebotEntity, DeebotEntityDescription
 
-_LOGGER = logging.getLogger(__name__)
+
+@dataclass
+class DeebotButtonEntityDescription(
+    ButtonEntityDescription,  # type: ignore
+    DeebotEntityDescription,
+):
+    """Class describing debbot button entity."""
+
+
+ENTITY_DESCRIPTIONS: tuple[DeebotButtonEntityDescription, ...] = (
+    DeebotButtonEntityDescription(
+        capability_fn=lambda caps: caps.map.relocation if caps.map else None,
+        key="relocate",
+        translation_key="relocate",
+        icon="mdi:map-marker-question",
+        entity_registry_enabled_default=True,  # Can be enabled as they don't poll data
+        entity_category=EntityCategory.DIAGNOSTIC,
+    ),
+)
 
 
 async def async_setup_entry(
@@ -24,18 +43,27 @@ async def async_setup_entry(
 ) -> None:
     """Add entities for passed config_entry in HA."""
     controller: DeebotController = hass.data[DOMAIN][config_entry.entry_id]
+    controller.register_platform_add_entities(
+        DeebotButtonEntity, ENTITY_DESCRIPTIONS, async_add_entities
+    )
 
-    new_devices = []
-    for vacbot in controller.vacuum_bots:
-        for component in LifeSpan:
-            new_devices.append(DeebotResetLifeSpanButtonEntity(vacbot, component))
-        new_devices.append(DeebotRelocateButtonEntity(vacbot))
+    def generate_reset_life_span(
+        device: VacuumBot,
+    ) -> Sequence[DeebotResetLifeSpanButtonEntity]:
+        return [
+            DeebotResetLifeSpanButtonEntity(device, component)
+            for component in device.capabilities.life_span.types
+        ]
 
-    if new_devices:
-        async_add_entities(new_devices)
+    controller.register_platform_add_entities_generator(
+        async_add_entities, generate_reset_life_span
+    )
 
 
-class DeebotResetLifeSpanButtonEntity(DeebotEntity, ButtonEntity):  # type: ignore
+class DeebotResetLifeSpanButtonEntity(
+    DeebotEntity[None, ButtonEntityDescription],
+    ButtonEntity,  # type: ignore
+):
     """Deebot reset life span button entity."""
 
     def __init__(self, vacuum_bot: VacuumBot, component: LifeSpan):
@@ -47,25 +75,20 @@ def __init__(self, vacuum_bot: VacuumBot, component: LifeSpan):
             entity_registry_enabled_default=True,  # Can be enabled as they don't poll data
             entity_category=EntityCategory.CONFIG,
         )
-        super().__init__(vacuum_bot, entity_description)
-        self._component = component
+        super().__init__(vacuum_bot, None, entity_description)
+        self._command = vacuum_bot.capabilities.life_span.reset(component)
 
     async def async_press(self) -> None:
         """Press the button."""
-        await self._vacuum_bot.execute_command(ResetLifeSpan(self._component))
-
+        await self._vacuum_bot.execute_command(self._command)
 
-class DeebotRelocateButtonEntity(DeebotEntity, ButtonEntity):  # type: ignore
-    """Deebot relocate button entity."""
 
-    entity_description = ButtonEntityDescription(
-        key="relocate",
-        translation_key="relocate",
-        icon="mdi:map-marker-question",
-        entity_registry_enabled_default=True,  # Can be enabled as they don't poll data
-        entity_category=EntityCategory.DIAGNOSTIC,
-    )
+class DeebotButtonEntity(
+    DeebotEntity[CapabilityExecute, DeebotButtonEntityDescription],
+    ButtonEntity,  # type: ignore
+):
+    """Deebot button entity."""
 
     async def async_press(self) -> None:
         """Press the button."""
-        await self._vacuum_bot.execute_command(SetRelocationState())
+        await self._vacuum_bot.execute_command(self._capability.execute())

custom_components/deebot/config_flow.py

@@ -218,11 +218,12 @@ def _get_options_schema(
     select_options = []
 
     for entry in devices:
-        label = entry.get("nick", entry["name"])
+        api_info = entry.api_device_info
+        label = api_info.get("nick", api_info["name"])
         if not label:
-            label = entry["name"]
+            label = api_info["name"]
         select_options.append(
-            selector.SelectOptionDict(value=entry["name"], label=label)
+            selector.SelectOptionDict(value=api_info["name"], label=label)
         )
 
     return vol.Schema(

custom_components/deebot/controller.py

@@ -2,13 +2,13 @@
 import logging
 import random
 import string
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping, Sequence
 from typing import Any
 
 from deebot_client.api_client import ApiClient
 from deebot_client.authentication import Authenticator
 from deebot_client.exceptions import InvalidAuthenticationError
-from deebot_client.models import Configuration
+from deebot_client.models import ApiDeviceInfo, Configuration
 from deebot_client.mqtt_client import MqttClient, MqttConfiguration
 from deebot_client.util import md5
 from deebot_client.vacuum_bot import VacuumBot
@@ -21,6 +21,11 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers import aiohttp_client
+from homeassistant.helpers.device_registry import DeviceEntry
+from homeassistant.helpers.entity import EntityDescription
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from custom_components.deebot.entity import DeebotEntity, DeebotEntityDescription
 
 from .const import CONF_CLIENT_DEVICE_ID, CONF_CONTINENT, CONF_COUNTRY
 
@@ -33,7 +38,7 @@ class DeebotController:
     def __init__(self, hass: HomeAssistant, config: Mapping[str, Any]):
         self._hass_config: Mapping[str, Any] = config
         self._hass: HomeAssistant = hass
-        self.vacuum_bots: list[VacuumBot] = []
+        self._devices: list[VacuumBot] = []
         verify_ssl = config.get(CONF_VERIFY_SSL, True)
         device_id = config.get(CONF_CLIENT_DEVICE_ID)
 
@@ -71,11 +76,15 @@ async def initialize(self) -> None:
             await self._mqtt.connect()
 
             for device in devices:
-                if device["name"] in self._hass_config.get(CONF_DEVICES, []):
+                if device.api_device_info["name"] in self._hass_config.get(
+                    CONF_DEVICES, []
+                ):
                     bot = VacuumBot(device, self._authenticator)
-                    _LOGGER.debug("New vacbot found: %s", device["name"])
+                    _LOGGER.debug(
+                        "New vacbot found: %s", device.api_device_info["name"]
+                    )
                     await bot.initialize(self._mqtt)
-                    self.vacuum_bots.append(bot)
+                    self._devices.append(bot)
 
             _LOGGER.debug("Controller initialize complete")
         except InvalidAuthenticationError as ex:
@@ -85,9 +94,50 @@ async def initialize(self) -> None:
             _LOGGER.error(msg, exc_info=True)
             raise ConfigEntryNotReady(msg) from ex
 
+    def register_platform_add_entities(
+        self,
+        entity_class: type[DeebotEntity],
+        descriptions: tuple[DeebotEntityDescription, ...],
+        async_add_entities: AddEntitiesCallback,
+    ) -> None:
+        """Create entities from descriptions and add them."""
+        new_entites: list[DeebotEntity] = []
+
+        for device in self._devices:
+            for description in descriptions:
+                if capability := description.capability_fn(device.capabilities):
+                    new_entites.append(entity_class(device, capability, description))
+
+        if new_entites:
+            async_add_entities(new_entites)
+
+    def register_platform_add_entities_generator(
+        self,
+        async_add_entities: AddEntitiesCallback,
+        func: Callable[[VacuumBot], Sequence[DeebotEntity[Any, EntityDescription]]],
+    ) -> None:
+        """Add entities generated through the provided function."""
+        new_entites: list[DeebotEntity[Any, EntityDescription]] = []
+
+        for device in self._devices:
+            new_entites.extend(func(device))
+
+        if new_entites:
+            async_add_entities(new_entites)
+
+    def get_device_info(self, device: DeviceEntry) -> ApiDeviceInfo | dict[str, str]:
+        """Get the device info for the given entry."""
+        for bot in self._devices:
+            for identifier in device.identifiers:
+                if bot.device_info.did == identifier[1]:
+                    return bot.device_info.api_device_info
+
+        _LOGGER.error("Could not find the device with entry: %s", device.json_repr)
+        return {"error": "Could not find the device"}
+
     async def teardown(self) -> None:
         """Disconnect controller."""
-        for bot in self.vacuum_bots:
+        for bot in self._devices:
             await bot.teardown()
         await self._mqtt.disconnect()
         await self._authenticator.teardown()

custom_components/deebot/diagnostics.py

@@ -25,10 +25,8 @@ async def async_get_device_diagnostics(
         "config": async_redact_data(config_entry.as_dict(), REDACT_CONFIG)
     }
 
-    for bot in controller.vacuum_bots:
-        for identifier in device.identifiers:
-            if bot.device_info.did == identifier[1]:
-                diag["device"] = async_redact_data(bot.device_info, REDACT_DEVICE)
-                break
+    diag["device"] = async_redact_data(
+        controller.get_device_info(device), REDACT_DEVICE
+    )
 
     return diag