roborock.data.containers

  1import dataclasses
  2import datetime
  3import inspect
  4import json
  5import logging
  6import re
  7import types
  8from dataclasses import asdict, dataclass, field
  9from enum import Enum
 10from functools import cached_property
 11from typing import Any, ClassVar, NamedTuple, get_args, get_origin
 12
 13from .code_mappings import (
 14    SHORT_MODEL_TO_ENUM,
 15    RoborockCategory,
 16    RoborockModeEnum,
 17    RoborockProductNickname,
 18)
 19
 20_LOGGER = logging.getLogger(__name__)
 21
 22
 23def _camelize(s: str):
 24    first, *others = s.split("_")
 25    if len(others) == 0:
 26        return s
 27    return "".join([first.lower(), *map(str.title, others)])
 28
 29
 30def _decamelize(s: str):
 31    # Split before uppercase letters not at the start, and before numbers
 32    s = re.sub(r"(?<=[a-z0-9])([A-Z])", r"_\1", s)
 33    s = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", s)  # Split acronyms followed by normal camelCase
 34    s = re.sub(r"([a-zA-Z])([0-9]+)", r"\1_\2", s)
 35    s = s.lower()
 36    # Temporary fix to avoid breaking any serialization.
 37    s = s.replace("base_64", "base64")
 38    return s
 39
 40
 41def _attr_repr(obj: Any) -> str:
 42    """Return a string representation of the object including specified attributes.
 43
 44    This reproduces the default repr behavior of dataclasses, but also includes
 45    properties. This must be called by the child class's __repr__ method since
 46    the parent RoborockBase class does not know about the child class's attributes.
 47    """
 48    # Reproduce default repr behavior
 49    parts = []
 50    for k in dir(obj):
 51        if k.startswith("_"):
 52            continue
 53        try:
 54            v = getattr(obj, k)
 55        except (RuntimeError, Exception):
 56            continue
 57        if callable(v):
 58            continue
 59        parts.append(f"{k}={v!r}")
 60    return f"{type(obj).__name__}({', '.join(parts)})"
 61
 62
 63def field_metadata(**kwargs):
 64    """Decorator to attach capability check metadata to a property.
 65
 66    This attaches a `_field_metadata` dictionary to the underlying getter function,
 67    which is then preserved when decorated with `@property`.
 68
 69    Supported metadata keys:
 70    - `feature` (str): Name of a capability property on `DeviceFeaturesTrait`.
 71    - `dock_feature` (str): Name of a capability property on `RoborockDockFeatures`.
 72    - `dps` (str/int): RoborockDataProtocol ID to check against supported schema IDs.
 73    """
 74
 75    def decorator(func):
 76        func._field_metadata = kwargs
 77        return func
 78
 79    return decorator
 80
 81
 82@dataclass(repr=False)
 83class RoborockBase:
 84    """Base class for all Roborock data classes."""
 85
 86    _missing_logged: ClassVar[set[str]] = set()
 87
 88    @staticmethod
 89    def _convert_to_class_obj(class_type: type, value):
 90        if get_origin(class_type) is list:
 91            sub_type = get_args(class_type)[0]
 92            return [RoborockBase._convert_to_class_obj(sub_type, obj) for obj in value]
 93        if get_origin(class_type) is dict:
 94            key_type, value_type = get_args(class_type)
 95            if key_type is not None:
 96                return {key_type(k): RoborockBase._convert_to_class_obj(value_type, v) for k, v in value.items()}
 97            return {k: RoborockBase._convert_to_class_obj(value_type, v) for k, v in value.items()}
 98        if inspect.isclass(class_type):
 99            if issubclass(class_type, RoborockBase):
100                return class_type.from_dict(value)
101            if issubclass(class_type, RoborockModeEnum):
102                return class_type.from_code(value)
103        if class_type is Any or type(class_type) is str:
104            return value
105        return class_type(value)  # type: ignore[call-arg]
106
107    @classmethod
108    def from_dict(cls, data: dict[str, Any]):
109        """Create an instance of the class from a dictionary."""
110        if not isinstance(data, dict):
111            return None
112        field_types = {field.name: field.type for field in dataclasses.fields(cls)}
113        normalized_data: dict[str, Any] = {}
114        for orig_key, value in data.items():
115            key = _decamelize(orig_key)
116            if field_types.get(key) is None:
117                if (log_key := f"{cls.__name__}.{key}") not in RoborockBase._missing_logged:
118                    _LOGGER.debug(
119                        "Key '%s' (decamelized: '%s') not found in %s fields, skipping",
120                        orig_key,
121                        key,
122                        cls.__name__,
123                    )
124                    RoborockBase._missing_logged.add(log_key)
125                continue
126            normalized_data[key] = value
127
128        result = RoborockBase.convert_dict(field_types, normalized_data)
129        return cls(**result)
130
131    @staticmethod
132    def convert_dict(types_map: dict[Any, type], data: dict[Any, Any]) -> dict[Any, Any]:
133        """Generic helper to convert a dictionary of values based on a schema map of types.
134
135        This is meant to be used by traits that use dataclass reflection similar to
136        `Roborock.from_dict` to merge in new data updates.
137        """
138        result: dict[Any, Any] = {}
139        for key, value in data.items():
140            if key not in types_map:
141                continue
142            field_type = types_map[key]
143            if value == "None" or value is None:
144                result[key] = None
145                continue
146            if isinstance(field_type, types.UnionType):
147                for subtype in get_args(field_type):
148                    if subtype is types.NoneType:
149                        continue
150                    try:
151                        result[key] = RoborockBase._convert_to_class_obj(subtype, value)
152                        break
153                    except Exception:
154                        _LOGGER.exception(f"Failed to convert {key} with value {value} to type {subtype}")
155                        continue
156            else:
157                try:
158                    result[key] = RoborockBase._convert_to_class_obj(field_type, value)
159                except Exception:
160                    _LOGGER.exception(f"Failed to convert {key} with value {value} to type {field_type}")
161                    continue
162
163        return result
164
165    def as_dict(self) -> dict:
166        return asdict(
167            self,
168            dict_factory=lambda _fields: {
169                _camelize(key): value.value if isinstance(value, Enum) else value
170                for (key, value) in _fields
171                if value is not None
172            },
173        )
174
175
176@dataclass
177class RoborockBaseTimer(RoborockBase):
178    start_hour: int | None = None
179    start_minute: int | None = None
180    end_hour: int | None = None
181    end_minute: int | None = None
182    enabled: int | None = None
183
184    @property
185    def start_time(self) -> datetime.time | None:
186        return (
187            datetime.time(hour=self.start_hour, minute=self.start_minute)
188            if self.start_hour is not None and self.start_minute is not None
189            else None
190        )
191
192    @property
193    def end_time(self) -> datetime.time | None:
194        return (
195            datetime.time(hour=self.end_hour, minute=self.end_minute)
196            if self.end_hour is not None and self.end_minute is not None
197            else None
198        )
199
200    def as_list(self) -> list:
201        return [self.start_hour, self.start_minute, self.end_hour, self.end_minute]
202
203    def __repr__(self) -> str:
204        return _attr_repr(self)
205
206
207@dataclass
208class Reference(RoborockBase):
209    r: str | None = None
210    a: str | None = None
211    m: str | None = None
212    l: str | None = None
213
214
215@dataclass
216class RRiot(RoborockBase):
217    u: str
218    s: str
219    h: str
220    k: str
221    r: Reference
222
223
224@dataclass
225class UserData(RoborockBase):
226    rriot: RRiot
227    uid: int | None = None
228    tokentype: str | None = None
229    token: str | None = None
230    rruid: str | None = None
231    region: str | None = None
232    countrycode: str | None = None
233    country: str | None = None
234    nickname: str | None = None
235    tuya_device_state: int | None = None
236    avatarurl: str | None = None
237
238
239@dataclass
240class HomeDataProductSchema(RoborockBase):
241    id: Any | None = None
242    name: Any | None = None
243    code: Any | None = None
244    mode: Any | None = None
245    type: Any | None = None
246    product_property: Any | None = None
247    property: Any | None = None
248    desc: Any | None = None
249
250
251@dataclass
252class HomeDataProduct(RoborockBase):
253    id: str
254    name: str
255    model: str
256    category: RoborockCategory
257    code: str | None = None
258    icon_url: str | None = None
259    attribute: Any | None = None
260    capability: int | None = None
261    schema: list[HomeDataProductSchema] | None = None
262
263    @property
264    def product_nickname(self) -> RoborockProductNickname:
265        return SHORT_MODEL_TO_ENUM.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS)
266
267    def summary_info(self) -> str:
268        """Return a string with key product information for logging purposes."""
269        return f"{self.name} (model={self.model}, category={self.category})"
270
271    @cached_property
272    def supported_schema_codes(self) -> set[str]:
273        """Return a set of schema codes that are supported by the device.
274
275        These correspond with string field names like "state" or "error_code" that
276        correspond to RoborockDataProtocol or RoborockB01Protocol code values.
277        """
278        if self.schema is None:
279            return set()
280        return {schema.code for schema in self.schema if schema.code is not None}
281
282    @cached_property
283    def supported_schema_ids(self) -> set[int]:
284        """Return a set of schema IDs (DPS integers) that are supported by the device.
285
286        These correspond to RoborockMessageProtocol and RoborockDataProtocol or
287        RoborockB01Protocol enum number values (depends on the device protocol versions).
288        """
289        if self.schema is None:
290            return set()
291        return {int(schema.id) for schema in self.schema if schema.id is not None}
292
293
294@dataclass
295class HomeDataDevice(RoborockBase):
296    duid: str
297    name: str
298    local_key: str
299    product_id: str
300    fv: str | None = None
301    attribute: Any | None = None
302    active_time: int | None = None
303    runtime_env: Any | None = None
304    time_zone_id: str | None = None
305    icon_url: str | None = None
306    lon: Any | None = None
307    lat: Any | None = None
308    share: Any | None = None
309    share_time: Any | None = None
310    online: bool | None = None
311    pv: str | None = None
312    room_id: Any | None = None
313    tuya_uuid: Any | None = None
314    tuya_migrated: bool | None = None
315    extra: Any | None = None
316    sn: str | None = None
317    feature_set: str | None = None
318    new_feature_set: str | None = None
319    device_status: dict | None = None
320    silent_ota_switch: bool | None = None
321    setting: Any | None = None
322    f: bool | None = None
323    create_time: int | None = None
324    cid: str | None = None
325    share_type: Any | None = None
326    share_expired_time: int | None = None
327
328    def summary_info(self) -> str:
329        """Return a string with key device information for logging purposes."""
330        return f"{self.name} (pv={self.pv}, fv={self.fv}, online={self.online})"
331
332
333@dataclass
334class HomeDataRoom(RoborockBase):
335    id: int
336    name: str
337
338    @property
339    def iot_id(self) -> str:
340        """Return the room's ID as a string IOT ID."""
341        return str(self.id)
342
343
344@dataclass
345class HomeDataScene(RoborockBase):
346    id: int
347    name: str
348
349
350@dataclass
351class HomeDataSchedule(RoborockBase):
352    id: int
353    cron: str
354    repeated: bool
355    enabled: bool
356    param: dict | None = None
357
358
359@dataclass
360class HomeData(RoborockBase):
361    id: int
362    name: str
363    products: list[HomeDataProduct] = field(default_factory=lambda: [])
364    devices: list[HomeDataDevice] = field(default_factory=lambda: [])
365    received_devices: list[HomeDataDevice] = field(default_factory=lambda: [])
366    lon: Any | None = None
367    lat: Any | None = None
368    geo_name: Any | None = None
369    rooms: list[HomeDataRoom] = field(default_factory=list)
370
371    def get_all_devices(self) -> list[HomeDataDevice]:
372        devices = []
373        if self.devices is not None:
374            devices += self.devices
375        if self.received_devices is not None:
376            devices += self.received_devices
377        return devices
378
379    @cached_property
380    def product_map(self) -> dict[str, HomeDataProduct]:
381        """Returns a dictionary of product IDs to HomeDataProduct objects."""
382        return {product.id: product for product in self.products}
383
384    @cached_property
385    def device_products(self) -> dict[str, tuple[HomeDataDevice, HomeDataProduct]]:
386        """Returns a dictionary of device DUIDs to HomeDataDeviceProduct objects."""
387        product_map = self.product_map
388        return {
389            device.duid: (device, product)
390            for device in self.get_all_devices()
391            if (product := product_map.get(device.product_id)) is not None
392        }
393
394    @property
395    def rooms_map(self) -> dict[str, HomeDataRoom]:
396        """Returns a dictionary of Room iot_id to rooms"""
397        return {room.iot_id: room for room in self.rooms}
398
399    @property
400    def rooms_name_map(self) -> dict[str, str]:
401        """Returns a dictionary of Room iot_id to room names."""
402        return {room.iot_id: room.name for room in self.rooms}
403
404
405@dataclass
406class LoginData(RoborockBase):
407    user_data: UserData
408    email: str
409    home_data: HomeData | None = None
410
411
412@dataclass
413class DeviceData(RoborockBase):
414    device: HomeDataDevice
415    model: str
416    host: str | None = None
417
418    @property
419    def product_nickname(self) -> RoborockProductNickname:
420        return SHORT_MODEL_TO_ENUM.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS)
421
422    def __repr__(self) -> str:
423        return _attr_repr(self)
424
425
426@dataclass
427class RoomMapping(RoborockBase):
428    segment_id: int
429    iot_id: str
430
431
432@dataclass
433class NamedRoomMapping(RoomMapping):
434    """Dataclass representing a mapping of a room segment to a name.
435
436    The name information is not provided by the device directly, but is provided
437    from the HomeData based on the iot_id from the room.
438    """
439
440    @property
441    def name(self) -> str:
442        """The human-readable name of the room, or a default name if not available."""
443        return self.raw_name or f"Room {self.segment_id}"
444
445    raw_name: str | None = None
446    """The raw name of the room, as provided by the device."""
447
448
449@dataclass
450class CombinedMapInfo(RoborockBase):
451    """Data structure for caching home information.
452
453    This is not provided directly by the API, but is a combination of map data
454    and room data to provide a more useful structure.
455    """
456
457    map_flag: int
458    """The map identifier."""
459
460    name: str
461    """The name of the map from MultiMapsListMapInfo."""
462
463    rooms: list[NamedRoomMapping]
464    """The list of rooms in the map."""
465
466    @property
467    def rooms_map(self) -> dict[int, NamedRoomMapping]:
468        """Returns a mapping of segment_id to NamedRoomMapping."""
469        return {room.segment_id: room for room in self.rooms}
470
471
472@dataclass
473class BroadcastMessage(RoborockBase):
474    duid: str
475    ip: str
476    version: bytes
477
478
479class ServerTimer(NamedTuple):
480    id: str
481    status: str
482    dontknow: int
483
484
485@dataclass
486class RoborockProductStateValue(RoborockBase):
487    value: list
488    desc: dict
489
490
491@dataclass
492class RoborockProductState(RoborockBase):
493    dps: int
494    desc: dict
495    value: list[RoborockProductStateValue]
496
497
498@dataclass
499class RoborockProductSpec(RoborockBase):
500    state: RoborockProductState
501    battery: dict | None = None
502    dry_countdown: dict | None = None
503    extra: dict | None = None
504    offpeak: dict | None = None
505    countdown: dict | None = None
506    mode: dict | None = None
507    ota_nfo: dict | None = None
508    pause: dict | None = None
509    program: dict | None = None
510    shutdown: dict | None = None
511    washing_left: dict | None = None
512
513
514@dataclass
515class RoborockProduct(RoborockBase):
516    id: int | None = None
517    name: str | None = None
518    model: str | None = None
519    packagename: str | None = None
520    ssid: str | None = None
521    picurl: str | None = None
522    cardpicurl: str | None = None
523    mediumCardpicurl: str | None = None
524    resetwifipicurl: str | None = None
525    configPicUrl: str | None = None
526    pluginPicUrl: str | None = None
527    resetwifitext: dict | None = None
528    tuyaid: str | None = None
529    status: int | None = None
530    rriotid: str | None = None
531    pictures: list | None = None
532    ncMode: str | None = None
533    scope: str | None = None
534    product_tags: list | None = None
535    agreements: list | None = None
536    cardspec: str | None = None
537    plugin_pic_url: str | None = None
538
539    @property
540    def product_nickname(self) -> RoborockProductNickname | None:
541        if self.cardspec:
542            return RoborockProductSpec.from_dict(json.loads(self.cardspec).get("data"))
543        return None
544
545    def __repr__(self) -> str:
546        return _attr_repr(self)
547
548
549@dataclass
550class RoborockProductCategory(RoborockBase):
551    id: int
552    display_name: str
553    icon_url: str
554
555
556@dataclass
557class RoborockCategoryDetail(RoborockBase):
558    category: RoborockProductCategory
559    product_list: list[RoborockProduct]
560
561
562@dataclass
563class ProductResponse(RoborockBase):
564    category_detail_list: list[RoborockCategoryDetail]