import dataclasses import datetime import inspect import json import logging import re import types from dataclasses import asdict, dataclass, field from enum import Enum from functools import cached_property from typing import Any, ClassVar, NamedTuple, get_args, get_origin from .code_mappings import ( SHORT_MODEL_TO_ENUM, RoborockCategory, RoborockModeEnum, RoborockProductNickname, ) _LOGGER = logging.getLogger(__name__) def _camelize(s: str): first, *others = s.split("_") if len(others) == 0: return s return "".join([first.lower(), *map(str.title, others)]) def _decamelize(s: str): # Split before uppercase letters not at the start, and before numbers s = re.sub(r"(?<=[a-z0-9])([A-Z])", r"_\1", s) s = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", s) # Split acronyms followed by normal camelCase s = re.sub(r"([a-zA-Z])([0-9]+)", r"\1_\2", s) s = s.lower() # Temporary fix to avoid breaking any serialization. s = s.replace("base_64", "base64") return s def _attr_repr(obj: Any) -> str: """Return a string representation of the object including specified attributes. This reproduces the default repr behavior of dataclasses, but also includes properties. This must be called by the child class's __repr__ method since the parent RoborockBase class does not know about the child class's attributes. """ # Reproduce default repr behavior parts = [] for k in dir(obj): if k.startswith("_"): continue try: v = getattr(obj, k) except (RuntimeError, Exception): continue if callable(v): continue parts.append(f"{k}={v!r}") return f"{type(obj).__name__}({', '.join(parts)})" @dataclass(repr=False) class RoborockBase: """Base class for all Roborock data classes.""" _missing_logged: ClassVar[set[str]] = set() @staticmethod def _convert_to_class_obj(class_type: type, value): if get_origin(class_type) is list: sub_type = get_args(class_type)[0] return [RoborockBase._convert_to_class_obj(sub_type, obj) for obj in value] if get_origin(class_type) is dict: key_type, value_type = get_args(class_type) if key_type is not None: return {key_type(k): RoborockBase._convert_to_class_obj(value_type, v) for k, v in value.items()} return {k: RoborockBase._convert_to_class_obj(value_type, v) for k, v in value.items()} if inspect.isclass(class_type): if issubclass(class_type, RoborockBase): return class_type.from_dict(value) if issubclass(class_type, RoborockModeEnum): return class_type.from_code(value) if class_type is Any or type(class_type) is str: return value return class_type(value) # type: ignore[call-arg] @classmethod def from_dict(cls, data: dict[str, Any]): """Create an instance of the class from a dictionary.""" if not isinstance(data, dict): return None field_types = {field.name: field.type for field in dataclasses.fields(cls)} normalized_data: dict[str, Any] = {} for orig_key, value in data.items(): key = _decamelize(orig_key) if field_types.get(key) is None: if (log_key := f"{cls.__name__}.{key}") not in RoborockBase._missing_logged: _LOGGER.debug( "Key '%s' (decamelized: '%s') not found in %s fields, skipping", orig_key, key, cls.__name__, ) RoborockBase._missing_logged.add(log_key) continue normalized_data[key] = value result = RoborockBase.convert_dict(field_types, normalized_data) return cls(**result) @staticmethod def convert_dict(types_map: dict[Any, type], data: dict[Any, Any]) -> dict[Any, Any]: """Generic helper to convert a dictionary of values based on a schema map of types. This is meant to be used by traits that use dataclass reflection similar to `Roborock.from_dict` to merge in new data updates. """ result: dict[Any, Any] = {} for key, value in data.items(): if key not in types_map: continue field_type = types_map[key] if value == "None" or value is None: result[key] = None continue if isinstance(field_type, types.UnionType): for subtype in get_args(field_type): if subtype is types.NoneType: continue try: result[key] = RoborockBase._convert_to_class_obj(subtype, value) break except Exception: _LOGGER.exception(f"Failed to convert {key} with value {value} to type {subtype}") continue else: try: result[key] = RoborockBase._convert_to_class_obj(field_type, value) except Exception: _LOGGER.exception(f"Failed to convert {key} with value {value} to type {field_type}") continue return result def as_dict(self) -> dict: return asdict( self, dict_factory=lambda _fields: { _camelize(key): value.value if isinstance(value, Enum) else value for (key, value) in _fields if value is not None }, ) @dataclass class RoborockBaseTimer(RoborockBase): start_hour: int | None = None start_minute: int | None = None end_hour: int | None = None end_minute: int | None = None enabled: int | None = None @property def start_time(self) -> datetime.time | None: return ( datetime.time(hour=self.start_hour, minute=self.start_minute) if self.start_hour is not None and self.start_minute is not None else None ) @property def end_time(self) -> datetime.time | None: return ( datetime.time(hour=self.end_hour, minute=self.end_minute) if self.end_hour is not None and self.end_minute is not None else None ) def as_list(self) -> list: return [self.start_hour, self.start_minute, self.end_hour, self.end_minute] def __repr__(self) -> str: return _attr_repr(self) @dataclass class Reference(RoborockBase): r: str | None = None a: str | None = None m: str | None = None l: str | None = None @dataclass class RRiot(RoborockBase): u: str s: str h: str k: str r: Reference @dataclass class UserData(RoborockBase): rriot: RRiot uid: int | None = None tokentype: str | None = None token: str | None = None rruid: str | None = None region: str | None = None countrycode: str | None = None country: str | None = None nickname: str | None = None tuya_device_state: int | None = None avatarurl: str | None = None @dataclass class HomeDataProductSchema(RoborockBase): id: Any | None = None name: Any | None = None code: Any | None = None mode: Any | None = None type: Any | None = None product_property: Any | None = None property: Any | None = None desc: Any | None = None @dataclass class HomeDataProduct(RoborockBase): id: str name: str model: str category: RoborockCategory code: str | None = None icon_url: str | None = None attribute: Any | None = None capability: int | None = None schema: list[HomeDataProductSchema] | None = None @property def product_nickname(self) -> RoborockProductNickname: return SHORT_MODEL_TO_ENUM.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS) def summary_info(self) -> str: """Return a string with key product information for logging purposes.""" return f"{self.name} (model={self.model}, category={self.category})" @cached_property def supported_schema_codes(self) -> set[str]: """Return a set of schema codes that are supported by the device. These correspond with string field names like "state" or "error_code" that correspond to RoborockDataProtocol or RoborockB01Protocol code values. """ if self.schema is None: return set() return {schema.code for schema in self.schema if schema.code is not None} @cached_property def supported_schema_ids(self) -> set[int]: """Return a set of schema IDs (DPS integers) that are supported by the device. These correspond to RoborockMessageProtocol and RoborockDataProtocol or RoborockB01Protocol enum number values (depends on the device protocol versions). """ if self.schema is None: return set() return {int(schema.id) for schema in self.schema if schema.id is not None} @dataclass class HomeDataDevice(RoborockBase): duid: str name: str local_key: str product_id: str fv: str | None = None attribute: Any | None = None active_time: int | None = None runtime_env: Any | None = None time_zone_id: str | None = None icon_url: str | None = None lon: Any | None = None lat: Any | None = None share: Any | None = None share_time: Any | None = None online: bool | None = None pv: str | None = None room_id: Any | None = None tuya_uuid: Any | None = None tuya_migrated: bool | None = None extra: Any | None = None sn: str | None = None feature_set: str | None = None new_feature_set: str | None = None device_status: dict | None = None silent_ota_switch: bool | None = None setting: Any | None = None f: bool | None = None create_time: int | None = None cid: str | None = None share_type: Any | None = None share_expired_time: int | None = None def summary_info(self) -> str: """Return a string with key device information for logging purposes.""" return f"{self.name} (pv={self.pv}, fv={self.fv}, online={self.online})" @dataclass class HomeDataRoom(RoborockBase): id: int name: str @property def iot_id(self) -> str: """Return the room's ID as a string IOT ID.""" return str(self.id) @dataclass class HomeDataScene(RoborockBase): id: int name: str @dataclass class HomeDataSchedule(RoborockBase): id: int cron: str repeated: bool enabled: bool param: dict | None = None @dataclass class HomeData(RoborockBase): id: int name: str products: list[HomeDataProduct] = field(default_factory=lambda: []) devices: list[HomeDataDevice] = field(default_factory=lambda: []) received_devices: list[HomeDataDevice] = field(default_factory=lambda: []) lon: Any | None = None lat: Any | None = None geo_name: Any | None = None rooms: list[HomeDataRoom] = field(default_factory=list) def get_all_devices(self) -> list[HomeDataDevice]: devices = [] if self.devices is not None: devices += self.devices if self.received_devices is not None: devices += self.received_devices return devices @cached_property def product_map(self) -> dict[str, HomeDataProduct]: """Returns a dictionary of product IDs to HomeDataProduct objects.""" return {product.id: product for product in self.products} @cached_property def device_products(self) -> dict[str, tuple[HomeDataDevice, HomeDataProduct]]: """Returns a dictionary of device DUIDs to HomeDataDeviceProduct objects.""" product_map = self.product_map return { device.duid: (device, product) for device in self.get_all_devices() if (product := product_map.get(device.product_id)) is not None } @property def rooms_map(self) -> dict[str, HomeDataRoom]: """Returns a dictionary of Room iot_id to rooms""" return {room.iot_id: room for room in self.rooms} @property def rooms_name_map(self) -> dict[str, str]: """Returns a dictionary of Room iot_id to room names.""" return {room.iot_id: room.name for room in self.rooms} @dataclass class LoginData(RoborockBase): user_data: UserData email: str home_data: HomeData | None = None @dataclass class DeviceData(RoborockBase): device: HomeDataDevice model: str host: str | None = None @property def product_nickname(self) -> RoborockProductNickname: return SHORT_MODEL_TO_ENUM.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS) def __repr__(self) -> str: return _attr_repr(self) @dataclass class RoomMapping(RoborockBase): segment_id: int iot_id: str @dataclass class NamedRoomMapping(RoomMapping): """Dataclass representing a mapping of a room segment to a name. The name information is not provided by the device directly, but is provided from the HomeData based on the iot_id from the room. """ @property def name(self) -> str: """The human-readable name of the room, or a default name if not available.""" return self.raw_name or f"Room {self.segment_id}" raw_name: str | None = None """The raw name of the room, as provided by the device.""" @dataclass class CombinedMapInfo(RoborockBase): """Data structure for caching home information. This is not provided directly by the API, but is a combination of map data and room data to provide a more useful structure. """ map_flag: int """The map identifier.""" name: str """The name of the map from MultiMapsListMapInfo.""" rooms: list[NamedRoomMapping] """The list of rooms in the map.""" @property def rooms_map(self) -> dict[int, NamedRoomMapping]: """Returns a mapping of segment_id to NamedRoomMapping.""" return {room.segment_id: room for room in self.rooms} @dataclass class BroadcastMessage(RoborockBase): duid: str ip: str version: bytes class ServerTimer(NamedTuple): id: str status: str dontknow: int @dataclass class RoborockProductStateValue(RoborockBase): value: list desc: dict @dataclass class RoborockProductState(RoborockBase): dps: int desc: dict value: list[RoborockProductStateValue] @dataclass class RoborockProductSpec(RoborockBase): state: RoborockProductState battery: dict | None = None dry_countdown: dict | None = None extra: dict | None = None offpeak: dict | None = None countdown: dict | None = None mode: dict | None = None ota_nfo: dict | None = None pause: dict | None = None program: dict | None = None shutdown: dict | None = None washing_left: dict | None = None @dataclass class RoborockProduct(RoborockBase): id: int | None = None name: str | None = None model: str | None = None packagename: str | None = None ssid: str | None = None picurl: str | None = None cardpicurl: str | None = None mediumCardpicurl: str | None = None resetwifipicurl: str | None = None configPicUrl: str | None = None pluginPicUrl: str | None = None resetwifitext: dict | None = None tuyaid: str | None = None status: int | None = None rriotid: str | None = None pictures: list | None = None ncMode: str | None = None scope: str | None = None product_tags: list | None = None agreements: list | None = None cardspec: str | None = None plugin_pic_url: str | None = None @property def product_nickname(self) -> RoborockProductNickname | None: if self.cardspec: return RoborockProductSpec.from_dict(json.loads(self.cardspec).get("data")) return None def __repr__(self) -> str: return _attr_repr(self) @dataclass class RoborockProductCategory(RoborockBase): id: int display_name: str icon_url: str @dataclass class RoborockCategoryDetail(RoborockBase): category: RoborockProductCategory product_list: list[RoborockProduct] @dataclass class ProductResponse(RoborockBase): category_detail_list: list[RoborockCategoryDetail]