|
| 1 | +"""Trait for managing room mappings on Roborock devices.""" |
| 2 | + |
| 3 | +import logging |
| 4 | +from dataclasses import dataclass |
| 5 | + |
| 6 | +from roborock.containers import HomeData, RoborockBase, RoomMapping |
| 7 | +from roborock.devices.traits.v1 import common |
| 8 | +from roborock.roborock_typing import RoborockCommand |
| 9 | + |
| 10 | +_LOGGER = logging.getLogger(__name__) |
| 11 | + |
| 12 | +_DEFAULT_NAME = "Unknown" |
| 13 | + |
| 14 | + |
| 15 | +@dataclass |
| 16 | +class NamedRoomMapping(RoomMapping): |
| 17 | + """Dataclass representing a mapping of a room segment to a name. |
| 18 | +
|
| 19 | + The name information is not provided by the device directly, but is provided |
| 20 | + from the HomeData based on the iot_id from the room. |
| 21 | + """ |
| 22 | + |
| 23 | + name: str | None = None |
| 24 | + """The human-readable name of the room, if available.""" |
| 25 | + |
| 26 | + |
| 27 | +@dataclass |
| 28 | +class Rooms(RoborockBase): |
| 29 | + """Dataclass representing a collection of room mappings.""" |
| 30 | + |
| 31 | + rooms: list[NamedRoomMapping] | None = None |
| 32 | + """List of room mappings.""" |
| 33 | + |
| 34 | + @property |
| 35 | + def room_map(self) -> dict[int, NamedRoomMapping]: |
| 36 | + """Returns a mapping of segment_id to NamedRoomMapping.""" |
| 37 | + if self.rooms is None: |
| 38 | + return {} |
| 39 | + return {room.segment_id: room for room in self.rooms} |
| 40 | + |
| 41 | + |
| 42 | +@common.mqtt_rpc_channel |
| 43 | +class RoomsTrait(Rooms, common.V1TraitMixin): |
| 44 | + """Trait for managing the room mappings of Roborock devices.""" |
| 45 | + |
| 46 | + command = RoborockCommand.GET_ROOM_MAPPING |
| 47 | + |
| 48 | + def __init__(self, home_data: HomeData) -> None: |
| 49 | + """Initialize the RoomsTrait.""" |
| 50 | + super().__init__() |
| 51 | + self._home_data = home_data |
| 52 | + |
| 53 | + @property |
| 54 | + def _iot_id_room_name_map(self) -> dict[str, str]: |
| 55 | + """Returns a dictionary of Room IOT IDs to room names.""" |
| 56 | + return {str(room.id): room.name for room in self._home_data.rooms or ()} |
| 57 | + |
| 58 | + def _parse_response(self, response: common.V1ResponseData) -> Rooms: |
| 59 | + """Parse the response from the device into a list of NamedRoomMapping.""" |
| 60 | + if not isinstance(response, list): |
| 61 | + raise ValueError(f"Unexpected RoomsTrait response format: {response!r}") |
| 62 | + name_map = self._iot_id_room_name_map |
| 63 | + segment_pairs = _extract_segment_pairs(response) |
| 64 | + return Rooms( |
| 65 | + rooms=[ |
| 66 | + NamedRoomMapping(segment_id=segment_id, iot_id=iot_id, name=name_map.get(iot_id, _DEFAULT_NAME)) |
| 67 | + for segment_id, iot_id in segment_pairs |
| 68 | + ] |
| 69 | + ) |
| 70 | + |
| 71 | + |
| 72 | +def _extract_segment_pairs(response: list) -> list[tuple[int, str]]: |
| 73 | + """Extract segment_id and iot_id pairs from the response. |
| 74 | +
|
| 75 | + The response format can be either a flat list of [segment_id, iot_id] or a |
| 76 | + list of lists, where each inner list is a pair of [segment_id, iot_id]. This |
| 77 | + function normalizes the response into a list of (segment_id, iot_id) tuples |
| 78 | +
|
| 79 | + NOTE: We currently only have samples of the list of lists format in |
| 80 | + tests/protocols/testdata so improving test coverage with samples from a real |
| 81 | + device with this format would be helpful. |
| 82 | + """ |
| 83 | + if len(response) == 2 and not isinstance(response[0], list): |
| 84 | + segment_id, iot_id = response[0], response[1] |
| 85 | + return [(segment_id, iot_id)] |
| 86 | + |
| 87 | + segment_pairs: list[tuple[int, str]] = [] |
| 88 | + for part in response: |
| 89 | + if not isinstance(part, list) or len(part) != 2: |
| 90 | + _LOGGER.warning("Unexpected room mapping entry format: %r", part) |
| 91 | + continue |
| 92 | + segment_id, iot_id = part[0], part[1] |
| 93 | + segment_pairs.append((segment_id, iot_id)) |
| 94 | + return segment_pairs |
0 commit comments