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 fields that are supported by the device."""
255        if self.schema is None:
256            return set()
257        return {schema.code for schema in self.schema if schema.code is not None}
258
259
260@dataclass
261class HomeDataDevice(RoborockBase):
262    duid: str
263    name: str
264    local_key: str
265    product_id: str
266    fv: str | None = None
267    attribute: Any | None = None
268    active_time: int | None = None
269    runtime_env: Any | None = None
270    time_zone_id: str | None = None
271    icon_url: str | None = None
272    lon: Any | None = None
273    lat: Any | None = None
274    share: Any | None = None
275    share_time: Any | None = None
276    online: bool | None = None
277    pv: str | None = None
278    room_id: Any | None = None
279    tuya_uuid: Any | None = None
280    tuya_migrated: bool | None = None
281    extra: Any | None = None
282    sn: str | None = None
283    feature_set: str | None = None
284    new_feature_set: str | None = None
285    device_status: dict | None = None
286    silent_ota_switch: bool | None = None
287    setting: Any | None = None
288    f: bool | None = None
289    create_time: int | None = None
290    cid: str | None = None
291    share_type: Any | None = None
292    share_expired_time: int | None = None
293
294    def summary_info(self) -> str:
295        """Return a string with key device information for logging purposes."""
296        return f"{self.name} (pv={self.pv}, fv={self.fv}, online={self.online})"
297
298
299@dataclass
300class HomeDataRoom(RoborockBase):
301    id: int
302    name: str
303
304    @property
305    def iot_id(self) -> str:
306        """Return the room's ID as a string IOT ID."""
307        return str(self.id)
308
309
310@dataclass
311class HomeDataScene(RoborockBase):
312    id: int
313    name: str
314
315
316@dataclass
317class HomeDataSchedule(RoborockBase):
318    id: int
319    cron: str
320    repeated: bool
321    enabled: bool
322    param: dict | None = None
323
324
325@dataclass
326class HomeData(RoborockBase):
327    id: int
328    name: str
329    products: list[HomeDataProduct] = field(default_factory=lambda: [])
330    devices: list[HomeDataDevice] = field(default_factory=lambda: [])
331    received_devices: list[HomeDataDevice] = field(default_factory=lambda: [])
332    lon: Any | None = None
333    lat: Any | None = None
334    geo_name: Any | None = None
335    rooms: list[HomeDataRoom] = field(default_factory=list)
336
337    def get_all_devices(self) -> list[HomeDataDevice]:
338        devices = []
339        if self.devices is not None:
340            devices += self.devices
341        if self.received_devices is not None:
342            devices += self.received_devices
343        return devices
344
345    @cached_property
346    def product_map(self) -> dict[str, HomeDataProduct]:
347        """Returns a dictionary of product IDs to HomeDataProduct objects."""
348        return {product.id: product for product in self.products}
349
350    @cached_property
351    def device_products(self) -> dict[str, tuple[HomeDataDevice, HomeDataProduct]]:
352        """Returns a dictionary of device DUIDs to HomeDataDeviceProduct objects."""
353        product_map = self.product_map
354        return {
355            device.duid: (device, product)
356            for device in self.get_all_devices()
357            if (product := product_map.get(device.product_id)) is not None
358        }
359
360    @property
361    def rooms_map(self) -> dict[str, HomeDataRoom]:
362        """Returns a dictionary of Room iot_id to rooms"""
363        return {room.iot_id: room for room in self.rooms}
364
365    @property
366    def rooms_name_map(self) -> dict[str, str]:
367        """Returns a dictionary of Room iot_id to room names."""
368        return {room.iot_id: room.name for room in self.rooms}
369
370
371@dataclass
372class LoginData(RoborockBase):
373    user_data: UserData
374    email: str
375    home_data: HomeData | None = None
376
377
378@dataclass
379class DeviceData(RoborockBase):
380    device: HomeDataDevice
381    model: str
382    host: str | None = None
383
384    @property
385    def product_nickname(self) -> RoborockProductNickname:
386        return SHORT_MODEL_TO_ENUM.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS)
387
388    def __repr__(self) -> str:
389        return _attr_repr(self)
390
391
392@dataclass
393class RoomMapping(RoborockBase):
394    segment_id: int
395    iot_id: str
396
397
398@dataclass
399class NamedRoomMapping(RoomMapping):
400    """Dataclass representing a mapping of a room segment to a name.
401
402    The name information is not provided by the device directly, but is provided
403    from the HomeData based on the iot_id from the room.
404    """
405
406    @property
407    def name(self) -> str:
408        """The human-readable name of the room, or a default name if not available."""
409        return self.raw_name or f"Room {self.segment_id}"
410
411    raw_name: str | None = None
412    """The raw name of the room, as provided by the device."""
413
414
415@dataclass
416class CombinedMapInfo(RoborockBase):
417    """Data structure for caching home information.
418
419    This is not provided directly by the API, but is a combination of map data
420    and room data to provide a more useful structure.
421    """
422
423    map_flag: int
424    """The map identifier."""
425
426    name: str
427    """The name of the map from MultiMapsListMapInfo."""
428
429    rooms: list[NamedRoomMapping]
430    """The list of rooms in the map."""
431
432    @property
433    def rooms_map(self) -> dict[int, NamedRoomMapping]:
434        """Returns a mapping of segment_id to NamedRoomMapping."""
435        return {room.segment_id: room for room in self.rooms}
436
437
438@dataclass
439class BroadcastMessage(RoborockBase):
440    duid: str
441    ip: str
442    version: bytes
443
444
445class ServerTimer(NamedTuple):
446    id: str
447    status: str
448    dontknow: int
449
450
451@dataclass
452class RoborockProductStateValue(RoborockBase):
453    value: list
454    desc: dict
455
456
457@dataclass
458class RoborockProductState(RoborockBase):
459    dps: int
460    desc: dict
461    value: list[RoborockProductStateValue]
462
463
464@dataclass
465class RoborockProductSpec(RoborockBase):
466    state: RoborockProductState
467    battery: dict | None = None
468    dry_countdown: dict | None = None
469    extra: dict | None = None
470    offpeak: dict | None = None
471    countdown: dict | None = None
472    mode: dict | None = None
473    ota_nfo: dict | None = None
474    pause: dict | None = None
475    program: dict | None = None
476    shutdown: dict | None = None
477    washing_left: dict | None = None
478
479
480@dataclass
481class RoborockProduct(RoborockBase):
482    id: int | None = None
483    name: str | None = None
484    model: str | None = None
485    packagename: str | None = None
486    ssid: str | None = None
487    picurl: str | None = None
488    cardpicurl: str | None = None
489    mediumCardpicurl: str | None = None
490    resetwifipicurl: str | None = None
491    configPicUrl: str | None = None
492    pluginPicUrl: str | None = None
493    resetwifitext: dict | None = None
494    tuyaid: str | None = None
495    status: int | None = None
496    rriotid: str | None = None
497    pictures: list | None = None
498    ncMode: str | None = None
499    scope: str | None = None
500    product_tags: list | None = None
501    agreements: list | None = None
502    cardspec: str | None = None
503    plugin_pic_url: str | None = None
504
505    @property
506    def product_nickname(self) -> RoborockProductNickname | None:
507        if self.cardspec:
508            return RoborockProductSpec.from_dict(json.loads(self.cardspec).get("data"))
509        return None
510
511    def __repr__(self) -> str:
512        return _attr_repr(self)
513
514
515@dataclass
516class RoborockProductCategory(RoborockBase):
517    id: int
518    display_name: str
519    icon_url: str
520
521
522@dataclass
523class RoborockCategoryDetail(RoborockBase):
524    category: RoborockProductCategory
525    product_list: list[RoborockProduct]
526
527
528@dataclass
529class ProductResponse(RoborockBase):
530    category_detail_list: list[RoborockCategoryDetail]