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