roborock.devices.traits.v1

  1"""Create traits for V1 devices.
  2
  3Traits are modular components that encapsulate specific features of a Roborock
  4device. This module provides a factory function to create and initialize the
  5appropriate traits for V1 devices based on their capabilities.
  6
  7Using Traits
  8------------
  9Traits are accessed via the `v1_properties` attribute on a device. Each trait
 10represents a specific capability, such as `status`, `consumables`, or `rooms`.
 11
 12Traits serve two main purposes:
 131.  **State**: Traits are dataclasses that hold the current state of the device
 14    feature. You can access attributes directly (e.g., `device.v1_properties.status.battery`).
 152.  **Commands**: Traits provide methods to control the device. For example,
 16    `device.v1_properties.volume.set_volume()`.
 17
 18Additionally, the `command` trait provides a generic way to send any command to the
 19device (e.g. `device.v1_properties.command.send("app_start")`). This is often used
 20for basic cleaning operations like starting, stopping, or docking the vacuum.
 21
 22Most traits have a `refresh()` method that must be called to update their state
 23from the device. The state is not updated automatically in real-time unless
 24specifically implemented by the trait or via polling.
 25
 26Adding New Traits
 27-----------------
 28When adding a new trait, the most common pattern is to subclass `V1TraitMixin`
 29and a `RoborockBase` dataclass. You must define a `command` class variable that
 30specifies the `RoborockCommand` used to fetch the trait data from the device.
 31See `common.py` for more details on common patterns used across traits.
 32
 33There are some additional decorators in `common.py` that can be used to specify which
 34RPC channel to use for the trait (standard, MQTT/cloud, or map-specific).
 35
 36  - `@common.mqtt_rpc_channel` - Use the MQTT RPC channel for this trait.
 37  - `@common.map_rpc_channel` - Use the map RPC channel for this trait.
 38
 39There are also some attributes that specify device feature dependencies for
 40optional traits:
 41
 42    - `requires_feature` - The string name of the device feature that must be supported
 43        for this trait to be enabled. See `DeviceFeaturesTrait` for a list of
 44        available features.
 45    - `requires_dock_type` - If set, this is a function that accepts a `RoborockDockTypeCode`
 46        and returns a boolean indicating whether the trait is supported for that dock type.
 47
 48Additionally, DeviceFeaturesTrait has a method `is_field_supported` that is used to
 49check individual trait field values. This is a more fine grained version to allow
 50optional fields in a dataclass, vs the above feature checks that apply to an entire
 51trait. The `dps` field metadata attribute references a schema code in
 52HomeDataProduct Schema that is required for the field to be supported.
 53"""
 54
 55import logging
 56from collections.abc import Callable
 57from dataclasses import dataclass, field, fields
 58from typing import Any, get_args
 59
 60from roborock.data.containers import HomeData, HomeDataProduct, RoborockBase
 61from roborock.data.v1.v1_code_mappings import RoborockDockTypeCode
 62from roborock.devices.cache import DeviceCache
 63from roborock.devices.traits import Trait
 64from roborock.exceptions import RoborockException
 65from roborock.map.map_parser import MapParserConfig
 66from roborock.protocols.v1_protocol import V1RpcChannel, decode_data_protocol_message
 67from roborock.roborock_message import RoborockDataProtocol, RoborockMessage
 68from roborock.web_api import UserWebApiClient
 69
 70from . import (
 71    child_lock,
 72    clean_summary,
 73    command,
 74    common,
 75    consumeable,
 76    device_features,
 77    do_not_disturb,
 78    dust_collection_mode,
 79    flow_led_status,
 80    home,
 81    led_status,
 82    map_content,
 83    maps,
 84    network_info,
 85    rooms,
 86    routines,
 87    smart_wash_params,
 88    status,
 89    valley_electricity_timer,
 90    volume,
 91    wash_towel_mode,
 92)
 93from .child_lock import ChildLockTrait
 94from .clean_summary import CleanSummaryTrait
 95from .command import CommandTrait
 96from .common import V1TraitMixin
 97from .consumeable import ConsumableTrait
 98from .device_features import DeviceFeaturesTrait
 99from .do_not_disturb import DoNotDisturbTrait
100from .dust_collection_mode import DustCollectionModeTrait
101from .flow_led_status import FlowLedStatusTrait
102from .home import HomeTrait
103from .led_status import LedStatusTrait
104from .map_content import MapContentTrait
105from .maps import MapsTrait
106from .network_info import NetworkInfoTrait
107from .rooms import RoomsTrait
108from .routines import RoutinesTrait
109from .smart_wash_params import SmartWashParamsTrait
110from .status import StatusTrait
111from .valley_electricity_timer import ValleyElectricityTimerTrait
112from .volume import SoundVolumeTrait
113from .wash_towel_mode import WashTowelModeTrait
114
115_LOGGER = logging.getLogger(__name__)
116
117__all__ = [
118    "PropertiesApi",
119    "child_lock",
120    "clean_summary",
121    "command",
122    "common",
123    "consumeable",
124    "device_features",
125    "do_not_disturb",
126    "dust_collection_mode",
127    "flow_led_status",
128    "home",
129    "led_status",
130    "map_content",
131    "maps",
132    "network_info",
133    "rooms",
134    "routines",
135    "smart_wash_params",
136    "status",
137    "valley_electricity_timer",
138    "volume",
139    "wash_towel_mode",
140]
141
142
143@dataclass
144class PropertiesApi(Trait):
145    """Common properties for V1 devices.
146
147    This class holds all the traits that are common across all V1 devices.
148    """
149
150    # All v1 devices have these traits
151    status: StatusTrait
152    command: CommandTrait
153    dnd: DoNotDisturbTrait
154    clean_summary: CleanSummaryTrait
155    sound_volume: SoundVolumeTrait
156    rooms: RoomsTrait
157    maps: MapsTrait
158    map_content: MapContentTrait
159    consumables: ConsumableTrait
160    home: HomeTrait
161    device_features: DeviceFeaturesTrait
162    network_info: NetworkInfoTrait
163    routines: RoutinesTrait
164
165    # Optional features that may not be supported on all devices
166    child_lock: ChildLockTrait | None = None
167    led_status: LedStatusTrait | None = None
168    flow_led_status: FlowLedStatusTrait | None = None
169    valley_electricity_timer: ValleyElectricityTimerTrait | None = None
170    dust_collection_mode: DustCollectionModeTrait | None = None
171    wash_towel_mode: WashTowelModeTrait | None = None
172    smart_wash_params: SmartWashParamsTrait | None = None
173
174    def __init__(
175        self,
176        device_uid: str,
177        product: HomeDataProduct,
178        home_data: HomeData,
179        rpc_channel: V1RpcChannel,
180        mqtt_rpc_channel: V1RpcChannel,
181        map_rpc_channel: V1RpcChannel,
182        add_dps_listener: Callable[[Callable[[dict[RoborockDataProtocol, Any]], None]], Callable[[], None]],
183        web_api: UserWebApiClient,
184        device_cache: DeviceCache,
185        map_parser_config: MapParserConfig | None = None,
186        region: str | None = None,
187    ) -> None:
188        """Initialize the V1TraitProps."""
189        self._device_uid = device_uid
190        self._rpc_channel = rpc_channel
191        self._mqtt_rpc_channel = mqtt_rpc_channel
192        self._map_rpc_channel = map_rpc_channel
193        self._web_api = web_api
194        self._device_cache = device_cache
195        self._region = region
196        self._unsub: Callable[[], None] | None = None
197        self._add_dps_listener = add_dps_listener
198
199        self.device_features = DeviceFeaturesTrait(product, self._device_cache)
200        self.status = StatusTrait(self.device_features, region=self._region)
201        self.consumables = ConsumableTrait()
202        self.rooms = RoomsTrait(home_data, device_uid, web_api)
203        self.maps = MapsTrait(self.status)
204        self.map_content = MapContentTrait(map_parser_config)
205        self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, self._device_cache)
206        self.network_info = NetworkInfoTrait(device_uid, self._device_cache)
207        self.routines = RoutinesTrait(device_uid, web_api)
208
209        # Dynamically create any traits that need to be populated
210        for item in fields(self):
211            if (trait := getattr(self, item.name, None)) is None:
212                # We exclude optional features and them via discover_features
213                if (union_args := get_args(item.type)) is None or len(union_args) > 0:
214                    continue
215                _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
216                if not callable(item.type):
217                    continue
218                trait = item.type()
219                setattr(self, item.name, trait)
220            # This is a hack to allow setting the rpc_channel on all traits. This is
221            # used so we can preserve the dataclass behavior when the values in the
222            # traits are updated, but still want to allow them to have a reference
223            # to the rpc channel for sending commands.
224            trait._rpc_channel = self._get_rpc_channel(trait)
225
226    def _get_rpc_channel(self, trait: V1TraitMixin) -> V1RpcChannel:
227        # The decorator `@common.mqtt_rpc_channel` means that the trait needs
228        # to use the mqtt_rpc_channel (cloud only) instead of the rpc_channel (adaptive)
229        if hasattr(trait, "mqtt_rpc_channel"):
230            return self._mqtt_rpc_channel
231        elif hasattr(trait, "map_rpc_channel"):
232            return self._map_rpc_channel
233        else:
234            return self._rpc_channel
235
236    async def start(self) -> None:
237        """Start the properties API and discover features."""
238        if self._unsub:
239            return
240        await self.discover_features()
241        self._unsub = self._add_dps_listener(self._on_dps_update)
242
243    def close(self) -> None:
244        if self._unsub:
245            self._unsub()
246            self._unsub = None
247
248    def _on_dps_update(self, dps: dict[RoborockDataProtocol, Any]) -> None:
249        """Handle incoming messages from the device.
250
251        This will notify all traits of the new values. This can be improved in
252        the future to be dynamic when we have more traits that support dynamic
253        updates but for now we just invoke them manually.
254        """
255        _LOGGER.debug("Received message from device: %s", dps)
256        self.status.update_from_dps(dps)
257        self.consumables.update_from_dps(dps)
258
259    async def discover_features(self) -> None:
260        """Populate any supported traits that were not initialized in __init__."""
261        _LOGGER.debug("Starting optional trait discovery")
262        await self.device_features.refresh()
263        # Dock type also acts like a device feature for some traits.
264        dock_type = await self._dock_type()
265
266        # Initialize traits with special arguments before the generic loop
267        if self.wash_towel_mode is None and self._is_supported(WashTowelModeTrait, "wash_towel_mode", dock_type):
268            wash_towel_mode = WashTowelModeTrait(self.device_features)
269            wash_towel_mode._rpc_channel = self._get_rpc_channel(wash_towel_mode)  # type: ignore[assignment]
270            self.wash_towel_mode = wash_towel_mode
271
272        # Dynamically create any traits that need to be populated
273        for item in fields(self):
274            if (trait := getattr(self, item.name, None)) is not None:
275                continue
276            if (union_args := get_args(item.type)) is None:
277                raise ValueError(f"Unexpected non-union type for trait {item.name}: {item.type}")
278            if len(union_args) != 2 or type(None) not in union_args:
279                raise ValueError(f"Unexpected non-optional type for trait {item.name}: {item.type}")
280
281            # Union args may not be in declared order
282            item_type = union_args[0] if union_args[1] is type(None) else union_args[1]
283            if not self._is_supported(item_type, item.name, dock_type):
284                _LOGGER.debug("Trait '%s' not supported, skipping", item.name)
285                continue
286            _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
287            trait = item_type()
288            setattr(self, item.name, trait)
289            trait._rpc_channel = self._get_rpc_channel(trait)
290
291    def _is_supported(self, trait_type: type[V1TraitMixin], name: str, dock_type: RoborockDockTypeCode) -> bool:
292        """Check if a trait is supported by the device."""
293
294        if (requires_dock_type := getattr(trait_type, "requires_dock_type", None)) is not None:
295            return requires_dock_type(dock_type)
296
297        if (feature_name := getattr(trait_type, "requires_feature", None)) is None:
298            _LOGGER.debug("Optional trait missing 'requires_feature' attribute %s, skipping", name)
299            return False
300        if (is_supported := getattr(self.device_features, feature_name)) is None:
301            raise ValueError(f"Device feature '{feature_name}' on trait '{name}' is unknown")
302        return is_supported
303
304    async def _dock_type(self) -> RoborockDockTypeCode:
305        """Get the dock type from the status trait or cache."""
306        dock_type = await self._get_cached_trait_data("dock_type")
307        if dock_type is not None:
308            _LOGGER.debug("Using cached dock type: %s", dock_type)
309            try:
310                return RoborockDockTypeCode(dock_type)
311            except ValueError:
312                _LOGGER.debug("Cached dock type %s is invalid, refreshing", dock_type)
313
314        _LOGGER.debug("Starting dock type discovery")
315        await self.status.refresh()
316        _LOGGER.debug("Fetched dock type: %s", self.status.dock_type)
317        if self.status.dock_type is None:
318            # Explicitly set so we reuse cached value next type
319            dock_type = RoborockDockTypeCode.no_dock
320        else:
321            dock_type = self.status.dock_type
322        await self._set_cached_trait_data("dock_type", dock_type)
323        return dock_type
324
325    async def _get_cached_trait_data(self, name: str) -> Any:
326        """Get the dock type from the status trait or cache."""
327        cache_data = await self._device_cache.get()
328        if cache_data.trait_data is None:
329            cache_data.trait_data = {}
330        _LOGGER.debug("Cached trait data: %s", cache_data.trait_data)
331        return cache_data.trait_data.get(name)
332
333    async def _set_cached_trait_data(self, name: str, value: Any) -> None:
334        """Set trait-specific cached data."""
335        cache_data = await self._device_cache.get()
336        if cache_data.trait_data is None:
337            cache_data.trait_data = {}
338        cache_data.trait_data[name] = value
339        _LOGGER.debug("Updating cached trait data: %s", cache_data.trait_data)
340        await self._device_cache.set(cache_data)
341
342    def as_dict(self) -> dict[str, Any]:
343        """Return the trait data as a dictionary."""
344        result: dict[str, Any] = {}
345        for item in fields(self):
346            trait = getattr(self, item.name, None)
347            if trait is None or not isinstance(trait, RoborockBase):
348                continue
349            data = trait.as_dict()
350            if data:  # Don't omit unset traits
351                result[item.name] = data
352        return result
353
354
355def create(
356    device_uid: str,
357    product: HomeDataProduct,
358    home_data: HomeData,
359    rpc_channel: V1RpcChannel,
360    mqtt_rpc_channel: V1RpcChannel,
361    map_rpc_channel: V1RpcChannel,
362    add_dps_listener: Callable[[Callable[[dict[RoborockDataProtocol, Any]], None]], Callable[[], None]],
363    web_api: UserWebApiClient,
364    device_cache: DeviceCache,
365    map_parser_config: MapParserConfig | None = None,
366    region: str | None = None,
367) -> PropertiesApi:
368    """Create traits for V1 devices."""
369    return PropertiesApi(
370        device_uid,
371        product,
372        home_data,
373        rpc_channel,
374        mqtt_rpc_channel,
375        map_rpc_channel,
376        add_dps_listener,
377        web_api,
378        device_cache,
379        map_parser_config,
380        region=region,
381    )