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