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        result: dict[str, Any] = {}
 95        for orig_key, value in data.items():
 96            key = _decamelize(orig_key)
 97            if (field_type := 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            if value == "None" or value is None:
108                result[key] = None
109                continue
110            if isinstance(field_type, types.UnionType):
111                for subtype in get_args(field_type):
112                    if subtype is types.NoneType:
113                        continue
114                    try:
115                        result[key] = RoborockBase._convert_to_class_obj(subtype, value)
116                        break
117                    except Exception:
118                        _LOGGER.exception(f"Failed to convert {key} with value {value} to type {subtype}")
119                        continue
120            else:
121                try:
122                    result[key] = RoborockBase._convert_to_class_obj(field_type, value)
123                except Exception:
124                    _LOGGER.exception(f"Failed to convert {key} with value {value} to type {field_type}")
125                    continue
126
127        return cls(**result)
128
129    def as_dict(self) -> dict:
130        return asdict(
131            self,
132            dict_factory=lambda _fields: {
133                _camelize(key): value.value if isinstance(value, Enum) else value
134                for (key, value) in _fields
135                if value is not None
136            },
137        )
138
139
140@dataclass
141class RoborockBaseTimer(RoborockBase):
142    start_hour: int | None = None
143    start_minute: int | None = None
144    end_hour: int | None = None
145    end_minute: int | None = None
146    enabled: int | None = None
147
148    @property
149    def start_time(self) -> datetime.time | None:
150        return (
151            datetime.time(hour=self.start_hour, minute=self.start_minute)
152            if self.start_hour is not None and self.start_minute is not None
153            else None
154        )
155
156    @property
157    def end_time(self) -> datetime.time | None:
158        return (
159            datetime.time(hour=self.end_hour, minute=self.end_minute)
160            if self.end_hour is not None and self.end_minute is not None
161            else None
162        )
163
164    def as_list(self) -> list:
165        return [self.start_hour, self.start_minute, self.end_hour, self.end_minute]
166
167    def __repr__(self) -> str:
168        return _attr_repr(self)
169
170
171@dataclass
172class Reference(RoborockBase):
173    r: str | None = None
174    a: str | None = None
175    m: str | None = None
176    l: str | None = None
177
178
179@dataclass
180class RRiot(RoborockBase):
181    u: str
182    s: str
183    h: str
184    k: str
185    r: Reference
186
187
188@dataclass
189class UserData(RoborockBase):
190    rriot: RRiot
191    uid: int | None = None
192    tokentype: str | None = None
193    token: str | None = None
194    rruid: str | None = None
195    region: str | None = None
196    countrycode: str | None = None
197    country: str | None = None
198    nickname: str | None = None
199    tuya_device_state: int | None = None
200    avatarurl: str | None = None
201
202
203@dataclass
204class HomeDataProductSchema(RoborockBase):
205    id: Any | None = None
206    name: Any | None = None
207    code: Any | None = None
208    mode: Any | None = None
209    type: Any | None = None
210    product_property: Any | None = None
211    property: Any | None = None
212    desc: Any | None = None
213
214
215@dataclass
216class HomeDataProduct(RoborockBase):
217    id: str
218    name: str
219    model: str
220    category: RoborockCategory
221    code: str | None = None
222    icon_url: str | None = None
223    attribute: Any | None = None
224    capability: int | None = None
225    schema: list[HomeDataProductSchema] | None = None
226
227    @property
228    def product_nickname(self) -> RoborockProductNickname:
229        return SHORT_MODEL_TO_ENUM.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS)
230
231    def summary_info(self) -> str:
232        """Return a string with key product information for logging purposes."""
233        return f"{self.name} (model={self.model}, category={self.category})"
234
235
236@dataclass
237class HomeDataDevice(RoborockBase):
238    duid: str
239    name: str
240    local_key: str
241    product_id: str
242    fv: str | None = None
243    attribute: Any | None = None
244    active_time: int | None = None
245    runtime_env: Any | None = None
246    time_zone_id: str | None = None
247    icon_url: str | None = None
248    lon: Any | None = None
249    lat: Any | None = None
250    share: Any | None = None
251    share_time: Any | None = None
252    online: bool | None = None
253    pv: str | None = None
254    room_id: Any | None = None
255    tuya_uuid: Any | None = None
256    tuya_migrated: bool | None = None
257    extra: Any | None = None
258    sn: str | None = None
259    feature_set: str | None = None
260    new_feature_set: str | None = None
261    device_status: dict | None = None
262    silent_ota_switch: bool | None = None
263    setting: Any | None = None
264    f: bool | None = None
265    create_time: int | None = None
266    cid: str | None = None
267    share_type: Any | None = None
268    share_expired_time: int | None = None
269
270    def summary_info(self) -> str:
271        """Return a string with key device information for logging purposes."""
272        return f"{self.name} (pv={self.pv}, fv={self.fv}, online={self.online})"
273
274
275@dataclass
276class HomeDataRoom(RoborockBase):
277    id: int
278    name: str
279
280
281@dataclass
282class HomeDataScene(RoborockBase):
283    id: int
284    name: str
285
286
287@dataclass
288class HomeDataSchedule(RoborockBase):
289    id: int
290    cron: str
291    repeated: bool
292    enabled: bool
293    param: dict | None = None
294
295
296@dataclass
297class HomeData(RoborockBase):
298    id: int
299    name: str
300    products: list[HomeDataProduct] = field(default_factory=lambda: [])
301    devices: list[HomeDataDevice] = field(default_factory=lambda: [])
302    received_devices: list[HomeDataDevice] = field(default_factory=lambda: [])
303    lon: Any | None = None
304    lat: Any | None = None
305    geo_name: Any | None = None
306    rooms: list[HomeDataRoom] = field(default_factory=list)
307
308    def get_all_devices(self) -> list[HomeDataDevice]:
309        devices = []
310        if self.devices is not None:
311            devices += self.devices
312        if self.received_devices is not None:
313            devices += self.received_devices
314        return devices
315
316    @cached_property
317    def product_map(self) -> dict[str, HomeDataProduct]:
318        """Returns a dictionary of product IDs to HomeDataProduct objects."""
319        return {product.id: product for product in self.products}
320
321    @cached_property
322    def device_products(self) -> dict[str, tuple[HomeDataDevice, HomeDataProduct]]:
323        """Returns a dictionary of device DUIDs to HomeDataDeviceProduct objects."""
324        product_map = self.product_map
325        return {
326            device.duid: (device, product)
327            for device in self.get_all_devices()
328            if (product := product_map.get(device.product_id)) is not None
329        }
330
331
332@dataclass
333class LoginData(RoborockBase):
334    user_data: UserData
335    email: str
336    home_data: HomeData | None = None
337
338
339@dataclass
340class DeviceData(RoborockBase):
341    device: HomeDataDevice
342    model: str
343    host: str | None = None
344
345    @property
346    def product_nickname(self) -> RoborockProductNickname:
347        return SHORT_MODEL_TO_ENUM.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS)
348
349    def __repr__(self) -> str:
350        return _attr_repr(self)
351
352
353@dataclass
354class RoomMapping(RoborockBase):
355    segment_id: int
356    iot_id: str
357
358
359@dataclass
360class NamedRoomMapping(RoomMapping):
361    """Dataclass representing a mapping of a room segment to a name.
362
363    The name information is not provided by the device directly, but is provided
364    from the HomeData based on the iot_id from the room.
365    """
366
367    name: str
368    """The human-readable name of the room, if available."""
369
370
371@dataclass
372class CombinedMapInfo(RoborockBase):
373    """Data structure for caching home information.
374
375    This is not provided directly by the API, but is a combination of map data
376    and room data to provide a more useful structure.
377    """
378
379    map_flag: int
380    """The map identifier."""
381
382    name: str
383    """The name of the map from MultiMapsListMapInfo."""
384
385    rooms: list[NamedRoomMapping]
386    """The list of rooms in the map."""
387
388
389@dataclass
390class BroadcastMessage(RoborockBase):
391    duid: str
392    ip: str
393    version: bytes
394
395
396class ServerTimer(NamedTuple):
397    id: str
398    status: str
399    dontknow: int
400
401
402@dataclass
403class RoborockProductStateValue(RoborockBase):
404    value: list
405    desc: dict
406
407
408@dataclass
409class RoborockProductState(RoborockBase):
410    dps: int
411    desc: dict
412    value: list[RoborockProductStateValue]
413
414
415@dataclass
416class RoborockProductSpec(RoborockBase):
417    state: RoborockProductState
418    battery: dict | None = None
419    dry_countdown: dict | None = None
420    extra: dict | None = None
421    offpeak: dict | None = None
422    countdown: dict | None = None
423    mode: dict | None = None
424    ota_nfo: dict | None = None
425    pause: dict | None = None
426    program: dict | None = None
427    shutdown: dict | None = None
428    washing_left: dict | None = None
429
430
431@dataclass
432class RoborockProduct(RoborockBase):
433    id: int | None = None
434    name: str | None = None
435    model: str | None = None
436    packagename: str | None = None
437    ssid: str | None = None
438    picurl: str | None = None
439    cardpicurl: str | None = None
440    mediumCardpicurl: str | None = None
441    resetwifipicurl: str | None = None
442    configPicUrl: str | None = None
443    pluginPicUrl: str | None = None
444    resetwifitext: dict | None = None
445    tuyaid: str | None = None
446    status: int | None = None
447    rriotid: str | None = None
448    pictures: list | None = None
449    ncMode: str | None = None
450    scope: str | None = None
451    product_tags: list | None = None
452    agreements: list | None = None
453    cardspec: str | None = None
454    plugin_pic_url: str | None = None
455
456    @property
457    def product_nickname(self) -> RoborockProductNickname | None:
458        if self.cardspec:
459            return RoborockProductSpec.from_dict(json.loads(self.cardspec).get("data"))
460        return None
461
462    def __repr__(self) -> str:
463        return _attr_repr(self)
464
465
466@dataclass
467class RoborockProductCategory(RoborockBase):
468    id: int
469    display_name: str
470    icon_url: str
471
472
473@dataclass
474class RoborockCategoryDetail(RoborockBase):
475    category: RoborockProductCategory
476    product_list: list[RoborockProduct]
477
478
479@dataclass
480class ProductResponse(RoborockBase):
481    category_detail_list: list[RoborockCategoryDetail]