roborock.devices.traits.v1.consumeable

 1"""Trait for managing consumable attributes.
 2
 3A consumable attribute is one that is expected to be replaced or refilled
 4periodically, such as filters, brushes, etc.
 5"""
 6
 7import logging
 8from enum import StrEnum
 9from typing import Any, Self
10
11from roborock.data import Consumable
12from roborock.devices.traits.common import DpsDataConverter, TraitUpdateListener
13from roborock.devices.traits.v1 import common
14from roborock.roborock_message import RoborockDataProtocol
15from roborock.roborock_typing import RoborockCommand
16
17__all__ = [
18    "ConsumableTrait",
19]
20
21_LOGGER = logging.getLogger(__name__)
22
23_DPS_CONVERTER = DpsDataConverter.from_dataclass(Consumable)
24
25
26class ConsumableAttribute(StrEnum):
27    """Enum for consumable attributes."""
28
29    SENSOR_DIRTY_TIME = "sensor_dirty_time"
30    FILTER_WORK_TIME = "filter_work_time"
31    SIDE_BRUSH_WORK_TIME = "side_brush_work_time"
32    MAIN_BRUSH_WORK_TIME = "main_brush_work_time"
33    STRAINER_WORK_TIME = "strainer_work_time"
34    CLEANING_BRUSH_WORK_TIME = "cleaning_brush_work_time"
35
36    @classmethod
37    def from_str(cls, value: str) -> Self:
38        """Create a ConsumableAttribute from a string value."""
39        for member in cls:
40            if member.value == value:
41                return member
42        raise ValueError(f"Unknown ConsumableAttribute: {value}")
43
44
45class ConsumableTrait(Consumable, common.V1TraitMixin, TraitUpdateListener):
46    """Trait for managing consumable attributes on Roborock devices.
47
48    After the first refresh, you can tell what consumables are supported by
49    checking which attributes are not None.
50    """
51
52    command = RoborockCommand.GET_CONSUMABLE
53    converter = common.DefaultConverter(Consumable)
54
55    def __init__(self) -> None:
56        """Initialize the consumable trait."""
57        super().__init__()
58        TraitUpdateListener.__init__(self, logger=_LOGGER)
59
60    async def reset_consumable(self, consumable: ConsumableAttribute) -> None:
61        """Reset a specific consumable attribute on the device."""
62        await self.rpc_channel.send_command(RoborockCommand.RESET_CONSUMABLE, params=[consumable.value])
63        await self.refresh()
64
65    def update_from_dps(self, decoded_dps: dict[RoborockDataProtocol, Any]) -> None:
66        """Update the trait from data protocol push message data.
67
68        This handles unsolicited status updates pushed by the device
69        via RoborockDataProtocol codes (e.g. STATE=121, BATTERY=122).
70        """
71        if _DPS_CONVERTER.update_from_dps(self, decoded_dps):
72            self._notify_update()