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_features` - If set, this is a function that accepts a `RoborockDockFeatures`
 46        and returns a boolean indicating whether the trait is supported for that dock.
 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.device_features import RoborockDockFeatures
 63from roborock.devices.cache import DeviceCache
 64from roborock.devices.traits import Trait
 65from roborock.exceptions import RoborockException
 66from roborock.map.map_parser import MapParserConfig
 67from roborock.protocols.v1_protocol import V1RpcChannel, decode_data_protocol_message
 68from roborock.roborock_message import RoborockDataProtocol, RoborockMessage
 69from roborock.web_api import UserWebApiClient
 70
 71from . import (
 72    child_lock,
 73    clean_summary,
 74    command,
 75    common,
 76    consumeable,
 77    customize_clean_mode,
 78    device_features,
 79    do_not_disturb,
 80    dust_collection_mode,
 81    flow_led_status,
 82    home,
 83    led_status,
 84    map_content,
 85    maps,
 86    mop_dryer,
 87    network_info,
 88    obstacle_photos,
 89    rooms,
 90    routines,
 91    smart_wash_params,
 92    status,
 93    valley_electricity_timer,
 94    volume,
 95    wash_towel_mode,
 96)
 97from .child_lock import ChildLockTrait
 98from .clean_summary import CleanSummaryTrait
 99from .command import CommandTrait
100from .common import V1TraitMixin
101from .consumeable import ConsumableTrait
102from .customize_clean_mode import CustomizeCleanModeTrait
103from .device_features import DeviceFeaturesTrait
104from .do_not_disturb import DoNotDisturbTrait
105from .dust_collection_mode import DustCollectionModeTrait
106from .flow_led_status import FlowLedStatusTrait
107from .home import HomeTrait
108from .led_status import LedStatusTrait
109from .map_content import MapContentTrait
110from .maps import MapsTrait
111from .mop_dryer import MopDryerTrait
112from .network_info import NetworkInfoTrait
113from .obstacle_photos import ObstaclePhotoTrait
114from .rooms import RoomsTrait
115from .routines import RoutinesTrait
116from .smart_wash_params import SmartWashParamsTrait
117from .status import StatusTrait
118from .valley_electricity_timer import ValleyElectricityTimerTrait
119from .volume import SoundVolumeTrait
120from .wash_towel_mode import WashTowelModeTrait
121
122_LOGGER = logging.getLogger(__name__)
123
124__all__ = [
125    "PropertiesApi",
126    "child_lock",
127    "clean_summary",
128    "command",
129    "common",
130    "consumeable",
131    "customize_clean_mode",
132    "device_features",
133    "do_not_disturb",
134    "dust_collection_mode",
135    "flow_led_status",
136    "home",
137    "led_status",
138    "map_content",
139    "maps",
140    "mop_dryer",
141    "network_info",
142    "obstacle_photos",
143    "rooms",
144    "routines",
145    "smart_wash_params",
146    "status",
147    "valley_electricity_timer",
148    "volume",
149    "wash_towel_mode",
150]
151
152
153@dataclass
154class PropertiesApi(Trait):
155    """Common properties for V1 devices.
156
157    This class holds all the traits that are common across all V1 devices.
158    """
159
160    # All v1 devices have these traits
161    status: StatusTrait
162    command: CommandTrait
163    dnd: DoNotDisturbTrait
164    clean_summary: CleanSummaryTrait
165    sound_volume: SoundVolumeTrait
166    rooms: RoomsTrait
167    maps: MapsTrait
168    map_content: MapContentTrait
169    consumables: ConsumableTrait
170    home: HomeTrait
171    device_features: DeviceFeaturesTrait
172    network_info: NetworkInfoTrait
173    routines: RoutinesTrait
174
175    # Optional features that may not be supported on all devices
176    child_lock: ChildLockTrait | None = None
177    led_status: LedStatusTrait | None = None
178    flow_led_status: FlowLedStatusTrait | None = None
179    valley_electricity_timer: ValleyElectricityTimerTrait | None = None
180    dust_collection_mode: DustCollectionModeTrait | None = None
181    wash_towel_mode: WashTowelModeTrait | None = None
182    smart_wash_params: SmartWashParamsTrait | None = None
183    obstacle_photos: ObstaclePhotoTrait | None = None
184    mop_dryer: MopDryerTrait | None = None
185    customize_clean_mode: CustomizeCleanModeTrait | None = None
186
187    def __init__(
188        self,
189        device_uid: str,
190        product: HomeDataProduct,
191        home_data: HomeData,
192        rpc_channel: V1RpcChannel,
193        mqtt_rpc_channel: V1RpcChannel,
194        map_rpc_channel: V1RpcChannel,
195        blob_rpc_channel: V1RpcChannel,
196        add_dps_listener: Callable[[Callable[[dict[RoborockDataProtocol, Any]], None]], Callable[[], None]],
197        web_api: UserWebApiClient,
198        device_cache: DeviceCache,
199        map_parser_config: MapParserConfig | None = None,
200        region: str | None = None,
201    ) -> None:
202        """Initialize the V1TraitProps."""
203        self._device_uid = device_uid
204        self._rpc_channel = rpc_channel
205        self._mqtt_rpc_channel = mqtt_rpc_channel
206        self._map_rpc_channel = map_rpc_channel
207        self._blob_rpc_channel = blob_rpc_channel
208        self._web_api = web_api
209        self._device_cache = device_cache
210        self._region = region
211        self._unsub: Callable[[], None] | None = None
212        self._add_dps_listener = add_dps_listener
213
214        self.device_features = DeviceFeaturesTrait(product, self._device_cache)
215        self.status = StatusTrait(self.device_features, region=self._region)
216        self.consumables = ConsumableTrait()
217        self.rooms = RoomsTrait(home_data, device_uid, web_api)
218        self.maps = MapsTrait(self.status)
219        self.map_content = MapContentTrait(map_parser_config)
220        self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, self._device_cache)
221        self.network_info = NetworkInfoTrait(device_uid, self._device_cache)
222        self.routines = RoutinesTrait(device_uid, web_api)
223
224        # Dynamically create any traits that need to be populated
225        for item in fields(self):
226            if (trait := getattr(self, item.name, None)) is None:
227                # We exclude optional features and them via discover_features
228                if (union_args := get_args(item.type)) is None or len(union_args) > 0:
229                    continue
230                _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
231                if not callable(item.type):
232                    continue
233                trait = item.type()
234                setattr(self, item.name, trait)
235            # This is a hack to allow setting the rpc_channel on all traits. This is
236            # used so we can preserve the dataclass behavior when the values in the
237            # traits are updated, but still want to allow them to have a reference
238            # to the rpc channel for sending commands.
239            trait._rpc_channel = self._get_rpc_channel(trait)
240
241    def _get_rpc_channel(self, trait: V1TraitMixin) -> V1RpcChannel:
242        # The decorator `@common.mqtt_rpc_channel` means that the trait needs
243        # to use the mqtt_rpc_channel (cloud only) instead of the rpc_channel (adaptive)
244        if hasattr(trait, "mqtt_rpc_channel"):
245            return self._mqtt_rpc_channel
246        elif hasattr(trait, "blob_rpc_channel"):
247            return self._blob_rpc_channel
248        elif hasattr(trait, "map_rpc_channel"):
249            return self._map_rpc_channel
250        else:
251            return self._rpc_channel
252
253    async def start(self) -> None:
254        """Start the properties API and discover features."""
255        if self._unsub:
256            return
257        await self.discover_features()
258        self._unsub = self._add_dps_listener(self._on_dps_update)
259
260    def close(self) -> None:
261        if self._unsub:
262            self._unsub()
263            self._unsub = None
264
265    def _on_dps_update(self, dps: dict[RoborockDataProtocol, Any]) -> None:
266        """Handle incoming messages from the device.
267
268        This will notify all traits of the new values. This can be improved in
269        the future to be dynamic when we have more traits that support dynamic
270        updates but for now we just invoke them manually.
271        """
272        _LOGGER.debug("Received message from device: %s", dps)
273        self.status.update_from_dps(dps)
274        self.consumables.update_from_dps(dps)
275
276    async def discover_features(self) -> None:
277        """Populate any supported traits that were not initialized in __init__."""
278        _LOGGER.debug("Starting optional trait discovery")
279        await self.device_features.refresh()
280        # Dock type also acts like a device feature for some traits.
281        dock_type = await self._dock_type()
282        dock_features = RoborockDockFeatures.from_dock_type(dock_type, has_am=self.status.has_am)
283        self.device_features.dock_features = dock_features
284
285        # Initialize traits with special arguments before the generic loop
286        if self.wash_towel_mode is None and self._is_supported(WashTowelModeTrait, "wash_towel_mode", dock_features):
287            wash_towel_mode = WashTowelModeTrait(self.device_features)
288            wash_towel_mode._rpc_channel = self._get_rpc_channel(wash_towel_mode)  # type: ignore[assignment]
289            self.wash_towel_mode = wash_towel_mode
290
291        if self.customize_clean_mode is None and self._is_supported(
292            CustomizeCleanModeTrait, "customize_clean_mode", dock_features
293        ):
294            customize_clean_mode = CustomizeCleanModeTrait(self.device_features)
295            customize_clean_mode._rpc_channel = self._get_rpc_channel(customize_clean_mode)  # type: ignore[assignment]
296            self.customize_clean_mode = customize_clean_mode
297
298        if self.obstacle_photos is None and self._is_supported(ObstaclePhotoTrait, "obstacle_photos", dock_features):
299            obstacle_photos = ObstaclePhotoTrait(self._rpc_channel)
300            obstacle_photos._rpc_channel = self._get_rpc_channel(obstacle_photos)
301            self.obstacle_photos = obstacle_photos
302
303        # Dynamically create any traits that need to be populated
304        for item in fields(self):
305            if (trait := getattr(self, item.name, None)) is not None:
306                continue
307            if (union_args := get_args(item.type)) is None:
308                raise ValueError(f"Unexpected non-union type for trait {item.name}: {item.type}")
309            if len(union_args) != 2 or type(None) not in union_args:
310                raise ValueError(f"Unexpected non-optional type for trait {item.name}: {item.type}")
311
312            # Union args may not be in declared order
313            item_type = union_args[0] if union_args[1] is type(None) else union_args[1]
314            if item_type is ObstaclePhotoTrait:
315                continue
316            if not self._is_supported(item_type, item.name, dock_features):
317                _LOGGER.debug("Trait '%s' not supported, skipping", item.name)
318                continue
319            _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
320            trait = item_type()
321            setattr(self, item.name, trait)
322            trait._rpc_channel = self._get_rpc_channel(trait)
323
324    def _is_supported(self, trait_type: type[V1TraitMixin], name: str, dock_features: RoborockDockFeatures) -> bool:
325        """Check if a trait is supported by the device."""
326
327        if (requires_dock_features := getattr(trait_type, "requires_dock_features", None)) is not None:
328            return requires_dock_features(dock_features)
329
330        if (feature_name := getattr(trait_type, "requires_feature", None)) is None:
331            _LOGGER.debug("Optional trait missing 'requires_feature' attribute %s, skipping", name)
332            return False
333        if (is_supported := getattr(self.device_features, feature_name)) is None:
334            raise ValueError(f"Device feature '{feature_name}' on trait '{name}' is unknown")
335        return is_supported
336
337    async def _dock_type(self) -> RoborockDockTypeCode:
338        """Get the dock type from the status trait or cache."""
339        dock_type = await self._get_cached_trait_data("dock_type")
340        if dock_type is not None:
341            _LOGGER.debug("Using cached dock type: %s", dock_type)
342            try:
343                dock_type = RoborockDockTypeCode(dock_type)
344            except ValueError:
345                _LOGGER.debug("Cached dock type %s is invalid, refreshing", dock_type)
346            else:
347                if self.status.dss is None:
348                    await self.status.refresh()
349                    if self.status.dock_type is not None:
350                        dock_type = self.status.dock_type
351                        await self._set_cached_trait_data("dock_type", dock_type)
352                return dock_type
353
354        _LOGGER.debug("Starting dock type discovery")
355        await self.status.refresh()
356        _LOGGER.debug("Fetched dock type: %s", self.status.dock_type)
357        if self.status.dock_type is None:
358            # Explicitly set so we reuse cached value next type
359            dock_type = RoborockDockTypeCode.o0_dock
360        else:
361            dock_type = self.status.dock_type
362        await self._set_cached_trait_data("dock_type", dock_type)
363        return dock_type
364
365    async def _get_cached_trait_data(self, name: str) -> Any:
366        """Get the dock type from the status trait or cache."""
367        cache_data = await self._device_cache.get()
368        if cache_data.trait_data is None:
369            cache_data.trait_data = {}
370        _LOGGER.debug("Cached trait data: %s", cache_data.trait_data)
371        return cache_data.trait_data.get(name)
372
373    async def _set_cached_trait_data(self, name: str, value: Any) -> None:
374        """Set trait-specific cached data."""
375        cache_data = await self._device_cache.get()
376        if cache_data.trait_data is None:
377            cache_data.trait_data = {}
378        cache_data.trait_data[name] = value
379        _LOGGER.debug("Updating cached trait data: %s", cache_data.trait_data)
380        await self._device_cache.set(cache_data)
381
382    def as_dict(self) -> dict[str, Any]:
383        """Return the trait data as a dictionary."""
384        result: dict[str, Any] = {}
385        for item in fields(self):
386            trait = getattr(self, item.name, None)
387            if trait is None or not isinstance(trait, RoborockBase):
388                continue
389            data = trait.as_dict()
390            if data:  # Don't omit unset traits
391                result[item.name] = data
392        return result
393
394
395def create(
396    device_uid: str,
397    product: HomeDataProduct,
398    home_data: HomeData,
399    rpc_channel: V1RpcChannel,
400    mqtt_rpc_channel: V1RpcChannel,
401    map_rpc_channel: V1RpcChannel,
402    blob_rpc_channel: V1RpcChannel,
403    add_dps_listener: Callable[[Callable[[dict[RoborockDataProtocol, Any]], None]], Callable[[], None]],
404    web_api: UserWebApiClient,
405    device_cache: DeviceCache,
406    map_parser_config: MapParserConfig | None = None,
407    region: str | None = None,
408) -> PropertiesApi:
409    """Create traits for V1 devices."""
410    return PropertiesApi(
411        device_uid,
412        product,
413        home_data,
414        rpc_channel,
415        mqtt_rpc_channel,
416        map_rpc_channel,
417        blob_rpc_channel,
418        add_dps_listener,
419        web_api,
420        device_cache,
421        map_parser_config,
422        region=region,
423    )