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.roborock_typing import RoborockCommand
  9from roborock.web_api import UserWebApiClient
 10
 11_LOGGER = logging.getLogger(__name__)
 12
 13
 14@dataclass
 15class Rooms(RoborockBase):
 16    """Dataclass representing a collection of room mappings."""
 17
 18    rooms: list[NamedRoomMapping] | None = None
 19    """List of room mappings."""
 20
 21    @property
 22    def room_map(self) -> dict[int, NamedRoomMapping]:
 23        """Returns a mapping of segment_id to NamedRoomMapping."""
 24        if self.rooms is None:
 25            return {}
 26        return {room.segment_id: room for room in self.rooms}
 27
 28    def with_room_names(self, name_map: dict[str, str]) -> "Rooms":
 29        """Create a new Rooms object with updated room names."""
 30        return Rooms(
 31            rooms=[
 32                NamedRoomMapping(
 33                    segment_id=room.segment_id,
 34                    iot_id=room.iot_id,
 35                    raw_name=name_map.get(room.iot_id),
 36                )
 37                for room in self.rooms or []
 38            ]
 39        )
 40
 41
 42class RoomsConverter(common.V1TraitDataConverter):
 43    """Converts response objects to Rooms."""
 44
 45    def convert(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        segment_map = self.extract_segment_map(response)
 50        return Rooms(
 51            rooms=[NamedRoomMapping(segment_id=segment_id, iot_id=iot_id) for segment_id, iot_id in segment_map.items()]
 52        )
 53
 54    @staticmethod
 55    def extract_segment_map(response: list) -> dict[int, str]:
 56        """Extract a segment_id -> iot_id mapping from the response.
 57
 58        The response format can be either a flat list of [segment_id, iot_id] or a
 59        list of lists, where each inner list is a pair of [segment_id, iot_id]. This
 60        function normalizes the response into a dict of segment_id to iot_id.
 61
 62        NOTE: We currently only partial samples of the room mapping formats, so
 63        improving test coverage with samples from a real device with this format
 64        would be helpful.
 65        """
 66        if len(response) == 2 and not isinstance(response[0], list):
 67            segment_id, iot_id = response[0], response[1]
 68            return {segment_id: str(iot_id)}
 69
 70        segment_map: dict[int, str] = {}
 71        for part in response:
 72            if not isinstance(part, list) or len(part) < 2:
 73                _LOGGER.warning("Unexpected room mapping entry format: %r", part)
 74                continue
 75            segment_id, iot_id = part[0], part[1]
 76            segment_map[segment_id] = str(iot_id)
 77        return segment_map
 78
 79
 80class RoomsTrait(Rooms, common.V1TraitMixin):
 81    """Trait for managing the room mappings of Roborock devices."""
 82
 83    command = RoborockCommand.GET_ROOM_MAPPING
 84    converter = RoomsConverter()
 85
 86    def __init__(self, home_data: HomeData, web_api: UserWebApiClient) -> None:
 87        """Initialize the RoomsTrait."""
 88        super().__init__()
 89        self._home_data = home_data
 90        self._web_api = web_api
 91        self._discovered_iot_ids: set[str] = set()
 92
 93    async def refresh(self) -> None:
 94        """Refresh room mappings and backfill unknown room names from the web API."""
 95        response = await self.rpc_channel.send_command(self.command)
 96        if not isinstance(response, list):
 97            raise ValueError(f"Unexpected RoomsTrait response format: {response!r}")
 98
 99        segment_map = RoomsConverter.extract_segment_map(response)
100        # Track all iot ids seen before. Refresh the room list when new ids are found.
101        new_iot_ids = set(segment_map.values()) - set(self._home_data.rooms_map.keys())
102        if new_iot_ids - self._discovered_iot_ids:
103            _LOGGER.debug("Refreshing room list to discover new room names")
104            if updated_rooms := await self._refresh_rooms():
105                _LOGGER.debug("Updating rooms: %s", list(updated_rooms))
106                self._home_data.rooms = updated_rooms
107            self._discovered_iot_ids.update(new_iot_ids)
108
109        rooms = self.converter.convert(response)
110        rooms = rooms.with_room_names(self._home_data.rooms_name_map)
111        common.merge_trait_values(self, rooms)
112
113    async def _refresh_rooms(self) -> list[HomeDataRoom]:
114        """Fetch the latest rooms from the web API."""
115        try:
116            return await self._web_api.get_rooms()
117        except Exception:
118            _LOGGER.debug("Failed to fetch rooms from web API", exc_info=True)
119            return []