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