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, HomeDataRoom, NamedRoomMapping, RoborockBase
  7from roborock.devices.traits.v1 import common
  8from roborock.exceptions import RoborockParsingException
  9from roborock.roborock_typing import RoborockCommand
 10from roborock.web_api import UserWebApiClient
 11
 12_LOGGER = logging.getLogger(__name__)
 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    def with_room_names(self, name_map: dict[str, str]) -> "Rooms":
 30        """Create a new Rooms object with updated room names."""
 31        return Rooms(
 32            rooms=[
 33                NamedRoomMapping(
 34                    segment_id=room.segment_id,
 35                    iot_id=room.iot_id,
 36                    raw_name=name_map.get(room.iot_id),
 37                )
 38                for room in self.rooms or []
 39            ]
 40        )
 41
 42
 43class RoomsConverter(common.V1TraitDataConverter):
 44    """Converts response objects to Rooms."""
 45
 46    def convert(self, response: common.V1ResponseData) -> Rooms:
 47        """Parse the response from the device into a list of NamedRoomMapping."""
 48        if not isinstance(response, list):
 49            raise ValueError(f"Unexpected RoomsTrait response format: {response!r}")
 50        segment_map = self.extract_segment_map(response)
 51        return Rooms(
 52            rooms=[NamedRoomMapping(segment_id=segment_id, iot_id=iot_id) for segment_id, iot_id in segment_map.items()]
 53        )
 54
 55    @staticmethod
 56    def extract_segment_map(response: list) -> dict[int, str]:
 57        """Extract a segment_id -> iot_id mapping from the response.
 58
 59        The response format can be either a flat list of [segment_id, iot_id] or a
 60        list of lists, where each inner list is a pair of [segment_id, iot_id]. This
 61        function normalizes the response into a dict of segment_id to iot_id.
 62
 63        NOTE: We currently only partial samples of the room mapping formats, so
 64        improving test coverage with samples from a real device with this format
 65        would be helpful.
 66        """
 67        if len(response) == 2 and not isinstance(response[0], list):
 68            segment_id, iot_id = response[0], response[1]
 69            return {segment_id: str(iot_id)}
 70
 71        segment_map: dict[int, str] = {}
 72        for part in response:
 73            if not isinstance(part, list) or len(part) < 2:
 74                _LOGGER.warning("Unexpected room mapping entry format: %r", part)
 75                continue
 76            segment_id, iot_id = part[0], part[1]
 77            segment_map[segment_id] = str(iot_id)
 78        return segment_map
 79
 80
 81class RoomsTrait(Rooms, common.V1TraitMixin):
 82    """Trait for managing the room mappings of Roborock devices."""
 83
 84    command = RoborockCommand.GET_ROOM_MAPPING
 85    converter = RoomsConverter()
 86
 87    def __init__(self, home_data: HomeData, device_uid: str, web_api: UserWebApiClient) -> None:
 88        """Initialize the RoomsTrait."""
 89        super().__init__()
 90        self._home_data = home_data
 91        self._device_uid = device_uid
 92        self._shared_device_uid = next(
 93            (device.duid for device in home_data.received_devices if device.duid == device_uid), None
 94        )
 95        self._web_api = web_api
 96        self._discovered_iot_ids: set[str] = set()
 97        self._room_names: dict[str, str] = dict(home_data.rooms_name_map)
 98
 99    @property
100    def _room_name_map(self) -> dict[str, str]:
101        return self._room_names
102
103    async def refresh(self) -> None:
104        """Refresh room mappings and backfill unknown room names from the web API."""
105        response = await self.rpc_channel.send_command(self.command)
106        if not isinstance(response, list):
107            raise RoborockParsingException(
108                trait_name=type(self).__name__,
109                command=self.command,
110                payload=response,
111                inner_error="Unexpected RoomsTrait response format",
112            )
113
114        segment_map = RoomsConverter.extract_segment_map(response)
115        # Track all iot ids seen before. Refresh the room list when new ids are found.
116        new_iot_ids = set(segment_map.values()) - set(self._room_name_map.keys())
117        if new_iot_ids - self._discovered_iot_ids:
118            _LOGGER.debug("Refreshing room list to discover new room names")
119            if updated_rooms := await self._refresh_rooms():
120                _LOGGER.debug("Updating rooms: %s", list(updated_rooms))
121                self._room_names = {room.iot_id: room.name for room in updated_rooms}
122                if self._shared_device_uid is None:
123                    self._home_data.rooms = updated_rooms
124            self._discovered_iot_ids.update(new_iot_ids)
125        try:
126            rooms = self.converter.convert(response)
127        except (TypeError, ValueError) as err:
128            raise RoborockParsingException(
129                trait_name=type(self).__name__,
130                command=self.command,
131                payload=response,
132                inner_error=err,
133            ) from err
134
135        rooms = rooms.with_room_names(self._room_name_map)
136        common.merge_trait_values(self, rooms)
137
138    async def _refresh_rooms(self) -> list[HomeDataRoom]:
139        """Fetch the latest rooms from the web API."""
140        try:
141            if self._shared_device_uid is not None:
142                rooms_by_id = {room.iot_id: room for room in self._home_data.rooms}
143                shared_rooms = await self._web_api.get_shared_device_rooms(self._shared_device_uid)
144                rooms_by_id.update({room.iot_id: room for room in shared_rooms})
145                return list(rooms_by_id.values())
146            return await self._web_api.get_rooms()
147        except Exception:
148            _LOGGER.debug("Failed to fetch rooms from web API", exc_info=True)
149            return []