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