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