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 `requires_schema_code` field metadata attribute is a string of the schema
 52code in HomeDataProduct Schema that is required for the field to be supported.
 53"""
 54
 55import logging
 56from dataclasses import dataclass, field, fields
 57from functools import cache
 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.map.map_parser import MapParserConfig
 65from roborock.protocols.v1_protocol import V1RpcChannel
 66from roborock.web_api import UserWebApiClient
 67
 68from . import (
 69    child_lock,
 70    clean_summary,
 71    command,
 72    common,
 73    consumeable,
 74    device_features,
 75    do_not_disturb,
 76    dust_collection_mode,
 77    flow_led_status,
 78    home,
 79    led_status,
 80    map_content,
 81    maps,
 82    network_info,
 83    rooms,
 84    routines,
 85    smart_wash_params,
 86    status,
 87    valley_electricity_timer,
 88    volume,
 89    wash_towel_mode,
 90)
 91from .child_lock import ChildLockTrait
 92from .clean_summary import CleanSummaryTrait
 93from .command import CommandTrait
 94from .common import V1TraitMixin
 95from .consumeable import ConsumableTrait
 96from .device_features import DeviceFeaturesTrait
 97from .do_not_disturb import DoNotDisturbTrait
 98from .dust_collection_mode import DustCollectionModeTrait
 99from .flow_led_status import FlowLedStatusTrait
100from .home import HomeTrait
101from .led_status import LedStatusTrait
102from .map_content import MapContentTrait
103from .maps import MapsTrait
104from .network_info import NetworkInfoTrait
105from .rooms import RoomsTrait
106from .routines import RoutinesTrait
107from .smart_wash_params import SmartWashParamsTrait
108from .status import StatusTrait
109from .valley_electricity_timer import ValleyElectricityTimerTrait
110from .volume import SoundVolumeTrait
111from .wash_towel_mode import WashTowelModeTrait
112
113_LOGGER = logging.getLogger(__name__)
114
115__all__ = [
116    "PropertiesApi",
117    "child_lock",
118    "clean_summary",
119    "command",
120    "common",
121    "consumeable",
122    "device_features",
123    "do_not_disturb",
124    "dust_collection_mode",
125    "flow_led_status",
126    "home",
127    "led_status",
128    "map_content",
129    "maps",
130    "network_info",
131    "rooms",
132    "routines",
133    "smart_wash_params",
134    "status",
135    "valley_electricity_timer",
136    "volume",
137    "wash_towel_mode",
138]
139
140
141@dataclass
142class PropertiesApi(Trait):
143    """Common properties for V1 devices.
144
145    This class holds all the traits that are common across all V1 devices.
146    """
147
148    # All v1 devices have these traits
149    status: StatusTrait
150    command: CommandTrait
151    dnd: DoNotDisturbTrait
152    clean_summary: CleanSummaryTrait
153    sound_volume: SoundVolumeTrait
154    rooms: RoomsTrait
155    maps: MapsTrait
156    map_content: MapContentTrait
157    consumables: ConsumableTrait
158    home: HomeTrait
159    device_features: DeviceFeaturesTrait
160    network_info: NetworkInfoTrait
161    routines: RoutinesTrait
162
163    # Optional features that may not be supported on all devices
164    child_lock: ChildLockTrait | None = None
165    led_status: LedStatusTrait | None = None
166    flow_led_status: FlowLedStatusTrait | None = None
167    valley_electricity_timer: ValleyElectricityTimerTrait | None = None
168    dust_collection_mode: DustCollectionModeTrait | None = None
169    wash_towel_mode: WashTowelModeTrait | None = None
170    smart_wash_params: SmartWashParamsTrait | None = None
171
172    def __init__(
173        self,
174        device_uid: str,
175        product: HomeDataProduct,
176        home_data: HomeData,
177        rpc_channel: V1RpcChannel,
178        mqtt_rpc_channel: V1RpcChannel,
179        map_rpc_channel: V1RpcChannel,
180        web_api: UserWebApiClient,
181        device_cache: DeviceCache,
182        map_parser_config: MapParserConfig | None = None,
183    ) -> None:
184        """Initialize the V1TraitProps."""
185        self._device_uid = device_uid
186        self._rpc_channel = rpc_channel
187        self._mqtt_rpc_channel = mqtt_rpc_channel
188        self._map_rpc_channel = map_rpc_channel
189        self._web_api = web_api
190        self._device_cache = device_cache
191
192        self.status = StatusTrait(product)
193        self.consumables = ConsumableTrait()
194        self.rooms = RoomsTrait(home_data)
195        self.maps = MapsTrait(self.status)
196        self.map_content = MapContentTrait(map_parser_config)
197        self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, self._device_cache)
198        self.device_features = DeviceFeaturesTrait(product, self._device_cache)
199        self.network_info = NetworkInfoTrait(device_uid, self._device_cache)
200        self.routines = RoutinesTrait(device_uid, web_api)
201
202        # Dynamically create any traits that need to be populated
203        for item in fields(self):
204            if (trait := getattr(self, item.name, None)) is None:
205                # We exclude optional features and them via discover_features
206                if (union_args := get_args(item.type)) is None or len(union_args) > 0:
207                    continue
208                _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
209                trait = item.type()
210                setattr(self, item.name, trait)
211            # This is a hack to allow setting the rpc_channel on all traits. This is
212            # used so we can preserve the dataclass behavior when the values in the
213            # traits are updated, but still want to allow them to have a reference
214            # to the rpc channel for sending commands.
215            trait._rpc_channel = self._get_rpc_channel(trait)
216
217    def _get_rpc_channel(self, trait: V1TraitMixin) -> V1RpcChannel:
218        # The decorator `@common.mqtt_rpc_channel` means that the trait needs
219        # to use the mqtt_rpc_channel (cloud only) instead of the rpc_channel (adaptive)
220        if hasattr(trait, "mqtt_rpc_channel"):
221            return self._mqtt_rpc_channel
222        elif hasattr(trait, "map_rpc_channel"):
223            return self._map_rpc_channel
224        else:
225            return self._rpc_channel
226
227    async def discover_features(self) -> None:
228        """Populate any supported traits that were not initialized in __init__."""
229        _LOGGER.debug("Starting optional trait discovery")
230        await self.device_features.refresh()
231        # Dock type also acts like a device feature for some traits.
232        dock_type = await self._dock_type()
233
234        # Dynamically create any traits that need to be populated
235        for item in fields(self):
236            if (trait := getattr(self, item.name, None)) is not None:
237                continue
238            if (union_args := get_args(item.type)) is None:
239                raise ValueError(f"Unexpected non-union type for trait {item.name}: {item.type}")
240            if len(union_args) != 2 or type(None) not in union_args:
241                raise ValueError(f"Unexpected non-optional type for trait {item.name}: {item.type}")
242
243            # Union args may not be in declared order
244            item_type = union_args[0] if union_args[1] is type(None) else union_args[1]
245            if not self._is_supported(item_type, item.name, dock_type):
246                _LOGGER.debug("Trait '%s' not supported, skipping", item.name)
247                continue
248            _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
249            trait = item_type()
250            setattr(self, item.name, trait)
251            trait._rpc_channel = self._get_rpc_channel(trait)
252
253    def _is_supported(self, trait_type: type[V1TraitMixin], name: str, dock_type: RoborockDockTypeCode) -> bool:
254        """Check if a trait is supported by the device."""
255
256        if (requires_dock_type := getattr(trait_type, "requires_dock_type", None)) is not None:
257            return requires_dock_type(dock_type)
258
259        if (feature_name := getattr(trait_type, "requires_feature", None)) is None:
260            _LOGGER.debug("Optional trait missing 'requires_feature' attribute %s, skipping", name)
261            return False
262        if (is_supported := getattr(self.device_features, feature_name)) is None:
263            raise ValueError(f"Device feature '{feature_name}' on trait '{name}' is unknown")
264        return is_supported
265
266    async def _dock_type(self) -> RoborockDockTypeCode:
267        """Get the dock type from the status trait or cache."""
268        dock_type = await self._get_cached_trait_data("dock_type")
269        if dock_type is not None:
270            _LOGGER.debug("Using cached dock type: %s", dock_type)
271            try:
272                return RoborockDockTypeCode(dock_type)
273            except ValueError:
274                _LOGGER.debug("Cached dock type %s is invalid, refreshing", dock_type)
275
276        _LOGGER.debug("Starting dock type discovery")
277        await self.status.refresh()
278        _LOGGER.debug("Fetched dock type: %s", self.status.dock_type)
279        if self.status.dock_type is None:
280            # Explicitly set so we reuse cached value next type
281            dock_type = RoborockDockTypeCode.no_dock
282        else:
283            dock_type = self.status.dock_type
284        await self._set_cached_trait_data("dock_type", dock_type)
285        return dock_type
286
287    async def _get_cached_trait_data(self, name: str) -> Any:
288        """Get the dock type from the status trait or cache."""
289        cache_data = await self._device_cache.get()
290        if cache_data.trait_data is None:
291            cache_data.trait_data = {}
292        _LOGGER.debug("Cached trait data: %s", cache_data.trait_data)
293        return cache_data.trait_data.get(name)
294
295    async def _set_cached_trait_data(self, name: str, value: Any) -> None:
296        """Set trait-specific cached data."""
297        cache_data = await self._device_cache.get()
298        if cache_data.trait_data is None:
299            cache_data.trait_data = {}
300        cache_data.trait_data[name] = value
301        _LOGGER.debug("Updating cached trait data: %s", cache_data.trait_data)
302        await self._device_cache.set(cache_data)
303
304    def as_dict(self) -> dict[str, Any]:
305        """Return the trait data as a dictionary."""
306        result: dict[str, Any] = {}
307        for item in fields(self):
308            trait = getattr(self, item.name, None)
309            if trait is None or not isinstance(trait, RoborockBase):
310                continue
311            data = trait.as_dict()
312            if data:  # Don't omit unset traits
313                result[item.name] = data
314        return result
315
316
317def create(
318    device_uid: str,
319    product: HomeDataProduct,
320    home_data: HomeData,
321    rpc_channel: V1RpcChannel,
322    mqtt_rpc_channel: V1RpcChannel,
323    map_rpc_channel: V1RpcChannel,
324    web_api: UserWebApiClient,
325    device_cache: DeviceCache,
326    map_parser_config: MapParserConfig | None = None,
327) -> PropertiesApi:
328    """Create traits for V1 devices."""
329    return PropertiesApi(
330        device_uid,
331        product,
332        home_data,
333        rpc_channel,
334        mqtt_rpc_channel,
335        map_rpc_channel,
336        web_api,
337        device_cache,
338        map_parser_config,
339    )