roborock.devices.traits.v1.rooms

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