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 Exception:  # noqa: BLE001
 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, exclude: set[str] | None = None) -> dict:
166        exclude_set = exclude or set()
167        return asdict(
168            self,
169            dict_factory=lambda _fields: {
170                _camelize(key): value.value if isinstance(value, Enum) else value
171                for (key, value) in _fields
172                if value is not None and key not in exclude_set
173            },
174        )
175
176
177@dataclass
178class RoborockBaseTimer(RoborockBase):
179    start_hour: int | None = None
180    start_minute: int | None = None
181    end_hour: int | None = None
182    end_minute: int | None = None
183    enabled: int | None = None
184
185    @property
186    def start_time(self) -> datetime.time | None:
187        return (
188            datetime.time(hour=self.start_hour, minute=self.start_minute)
189            if self.start_hour is not None and self.start_minute is not None
190            else None
191        )
192
193    @property
194    def end_time(self) -> datetime.time | None:
195        return (
196            datetime.time(hour=self.end_hour, minute=self.end_minute)
197            if self.end_hour is not None and self.end_minute is not None
198            else None
199        )
200
201    def as_list(self) -> list:
202        return [self.start_hour, self.start_minute, self.end_hour, self.end_minute]
203
204    def __repr__(self) -> str:
205        return _attr_repr(self)
206
207
208@dataclass
209class Reference(RoborockBase):
210    r: str | None = None
211    a: str | None = None
212    m: str | None = None
213    l: str | None = None  # noqa: E741
214
215
216@dataclass
217class RRiot(RoborockBase):
218    u: str
219    s: str
220    h: str
221    k: str
222    r: Reference
223
224
225@dataclass
226class UserData(RoborockBase):
227    rriot: RRiot
228    uid: int | None = None
229    tokentype: str | None = None
230    token: str | None = None
231    rruid: str | None = None
232    region: str | None = None
233    countrycode: str | None = None
234    country: str | None = None
235    nickname: str | None = None
236    tuya_device_state: int | None = None
237    avatarurl: str | None = None
238
239
240@dataclass
241class HomeDataProductSchema(RoborockBase):
242    id: Any | None = None
243    name: Any | None = None
244    code: Any | None = None
245    mode: Any | None = None
246    type: Any | None = None
247    product_property: Any | None = None
248    property: Any | None = None
249    desc: Any | None = None
250
251
252@dataclass
253class HomeDataProduct(RoborockBase):
254    id: str
255    name: str
256    model: str
257    category: RoborockCategory
258    code: str | None = None
259    icon_url: str | None = None
260    attribute: Any | None = None
261    capability: int | None = None
262    schema: list[HomeDataProductSchema] | None = None
263
264    @property
265    def product_nickname(self) -> RoborockProductNickname:
266        return SHORT_MODEL_TO_ENUM.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS)
267
268    def summary_info(self) -> str:
269        """Return a string with key product information for logging purposes."""
270        return f"{self.name} (model={self.model}, category={self.category})"
271
272    @cached_property
273    def supported_schema_codes(self) -> set[str]:
274        """Return a set of schema codes that are supported by the device.
275
276        These correspond with string field names like "state" or "error_code" that
277        correspond to RoborockDataProtocol or RoborockB01Protocol code values.
278        """
279        if self.schema is None:
280            return set()
281        return {schema.code for schema in self.schema if schema.code is not None}
282
283    @cached_property
284    def supported_schema_ids(self) -> set[int]:
285        """Return a set of schema IDs (DPS integers) that are supported by the device.
286
287        These correspond to RoborockMessageProtocol and RoborockDataProtocol or
288        RoborockB01Protocol enum number values (depends on the device protocol versions).
289        """
290        if self.schema is None:
291            return set()
292        return {int(schema.id) for schema in self.schema if schema.id is not None}
293
294
295@dataclass
296class HomeDataDevice(RoborockBase):
297    duid: str
298    name: str
299    local_key: str
300    product_id: str
301    fv: str | None = None
302    attribute: Any | None = None
303    active_time: int | None = None
304    runtime_env: Any | None = None
305    time_zone_id: str | None = None
306    icon_url: str | None = None
307    lon: Any | None = None
308    lat: Any | None = None
309    share: Any | None = None
310    share_time: Any | None = None
311    online: bool | None = None
312    pv: str | None = None
313    room_id: Any | None = None
314    tuya_uuid: Any | None = None
315    tuya_migrated: bool | None = None
316    extra: Any | None = None
317    sn: str | None = None
318    feature_set: str | None = None
319    new_feature_set: str | None = None
320    device_status: dict | None = None
321    silent_ota_switch: bool | None = None
322    setting: Any | None = None
323    f: bool | None = None
324    create_time: int | None = None
325    cid: str | None = None
326    share_type: Any | None = None
327    share_expired_time: int | None = None
328
329    def summary_info(self) -> str:
330        """Return a string with key device information for logging purposes."""
331        return f"{self.name} (pv={self.pv}, fv={self.fv}, online={self.online})"
332
333
334@dataclass
335class HomeDataRoom(RoborockBase):
336    id: int
337    name: str
338
339    @property
340    def iot_id(self) -> str:
341        """Return the room's ID as a string IOT ID."""
342        return str(self.id)
343
344
345@dataclass
346class HomeDataScene(RoborockBase):
347    id: int
348    name: str
349
350
351@dataclass
352class FirmwareInfo(RoborockBase):
353    """Firmware/OTA info from the cloud (`ota/firmware/{duid}/updatev2`)."""
354
355    version: str | None = None
356    """Latest available firmware version."""
357    current_version: str | None = None
358    """Currently installed firmware version."""
359    updatable: bool | None = None
360    """Whether a newer firmware is available to install."""
361    desc: str | None = None
362    """Release notes / description."""
363    release_time: str | None = None
364    """Release date of the available firmware as an ISO-8601 date (``YYYY-MM-DD``)."""
365    force_update: bool | None = None
366    """Whether the update is mandatory (cannot be skipped)."""
367
368
369@dataclass
370class HomeDataSchedule(RoborockBase):
371    id: int
372    cron: str
373    repeated: bool
374    enabled: bool
375    param: dict | None = None
376
377
378@dataclass
379class HomeData(RoborockBase):
380    id: int
381    name: str
382    products: list[HomeDataProduct] = field(default_factory=list)
383    devices: list[HomeDataDevice] = field(default_factory=list)
384    received_devices: list[HomeDataDevice] = field(default_factory=list)
385    lon: Any | None = None
386    lat: Any | None = None
387    geo_name: Any | None = None
388    rooms: list[HomeDataRoom] = field(default_factory=list)
389
390    def get_all_devices(self) -> list[HomeDataDevice]:
391        devices = []
392        if self.devices is not None:
393            devices += self.devices
394        if self.received_devices is not None:
395            devices += self.received_devices
396        return devices
397
398    @cached_property
399    def product_map(self) -> dict[str, HomeDataProduct]:
400        """Returns a dictionary of product IDs to HomeDataProduct objects."""
401        return {product.id: product for product in self.products}
402
403    @cached_property
404    def device_products(self) -> dict[str, tuple[HomeDataDevice, HomeDataProduct]]:
405        """Returns a dictionary of device DUIDs to HomeDataDeviceProduct objects."""
406        product_map = self.product_map
407        return {
408            device.duid: (device, product)
409            for device in self.get_all_devices()
410            if (product := product_map.get(device.product_id)) is not None
411        }
412
413    @property
414    def rooms_map(self) -> dict[str, HomeDataRoom]:
415        """Returns a dictionary of Room iot_id to rooms"""
416        return {room.iot_id: room for room in self.rooms}
417
418    @property
419    def rooms_name_map(self) -> dict[str, str]:
420        """Returns a dictionary of Room iot_id to room names."""
421        return {room.iot_id: room.name for room in self.rooms}
422
423
424@dataclass
425class LoginData(RoborockBase):
426    user_data: UserData
427    email: str
428    home_data: HomeData | None = None
429
430
431@dataclass
432class DeviceData(RoborockBase):
433    device: HomeDataDevice
434    model: str
435    host: str | None = None
436
437    @property
438    def product_nickname(self) -> RoborockProductNickname:
439        return SHORT_MODEL_TO_ENUM.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS)
440
441    def __repr__(self) -> str:
442        return _attr_repr(self)
443
444
445@dataclass
446class RoomMapping(RoborockBase):
447    segment_id: int
448    iot_id: str
449
450
451@dataclass
452class NamedRoomMapping(RoomMapping):
453    """Dataclass representing a mapping of a room segment to a name.
454
455    The name information is not provided by the device directly, but is provided
456    from the HomeData based on the iot_id from the room.
457    """
458
459    @property
460    def name(self) -> str:
461        """The human-readable name of the room, or a default name if not available."""
462        return self.raw_name or f"Room {self.segment_id}"
463
464    raw_name: str | None = None
465    """The raw name of the room, as provided by the device."""
466
467
468@dataclass
469class CombinedMapInfo(RoborockBase):
470    """Data structure for caching home information.
471
472    This is not provided directly by the API, but is a combination of map data
473    and room data to provide a more useful structure.
474    """
475
476    map_flag: int
477    """The map identifier."""
478
479    name: str
480    """The name of the map from MultiMapsListMapInfo."""
481
482    rooms: list[NamedRoomMapping]
483    """The list of rooms in the map."""
484
485    @property
486    def rooms_map(self) -> dict[int, NamedRoomMapping]:
487        """Returns a mapping of segment_id to NamedRoomMapping."""
488        return {room.segment_id: room for room in self.rooms}
489
490
491@dataclass
492class BroadcastMessage(RoborockBase):
493    duid: str
494    ip: str
495    version: bytes
496
497
498class ServerTimer(NamedTuple):
499    id: str
500    status: str
501    dontknow: int
502
503
504@dataclass
505class RoborockProductStateValue(RoborockBase):
506    value: list
507    desc: dict
508
509
510@dataclass
511class RoborockProductState(RoborockBase):
512    dps: int
513    desc: dict
514    value: list[RoborockProductStateValue]
515
516
517@dataclass
518class RoborockProductSpec(RoborockBase):
519    state: RoborockProductState
520    battery: dict | None = None
521    dry_countdown: dict | None = None
522    extra: dict | None = None
523    offpeak: dict | None = None
524    countdown: dict | None = None
525    mode: dict | None = None
526    ota_nfo: dict | None = None
527    pause: dict | None = None
528    program: dict | None = None
529    shutdown: dict | None = None
530    washing_left: dict | None = None
531
532
533@dataclass
534class RoborockProduct(RoborockBase):
535    id: int | None = None
536    name: str | None = None
537    model: str | None = None
538    packagename: str | None = None
539    ssid: str | None = None
540    picurl: str | None = None
541    cardpicurl: str | None = None
542    mediumCardpicurl: str | None = None
543    resetwifipicurl: str | None = None
544    configPicUrl: str | None = None
545    pluginPicUrl: str | None = None
546    resetwifitext: dict | None = None
547    tuyaid: str | None = None
548    status: int | None = None
549    rriotid: str | None = None
550    pictures: list | None = None
551    ncMode: str | None = None
552    scope: str | None = None
553    product_tags: list | None = None
554    agreements: list | None = None
555    cardspec: str | None = None
556    plugin_pic_url: str | None = None
557
558    @property
559    def product_nickname(self) -> RoborockProductNickname | None:
560        if self.cardspec:
561            return RoborockProductSpec.from_dict(json.loads(self.cardspec).get("data"))
562        return None
563
564    def __repr__(self) -> str:
565        return _attr_repr(self)
566
567
568@dataclass
569class RoborockProductCategory(RoborockBase):
570    id: int
571    display_name: str
572    icon_url: str
573
574
575@dataclass
576class RoborockCategoryDetail(RoborockBase):
577    category: RoborockProductCategory
578    product_list: list[RoborockProduct]
579
580
581@dataclass
582class ProductResponse(RoborockBase):
583    category_detail_list: list[RoborockCategoryDetail]