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