Skip to content
29 changes: 27 additions & 2 deletions roborock/data/containers.py
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,11 @@ 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):
Expand Down Expand Up @@ -352,6 +357,16 @@ def device_products(self) -> dict[str, tuple[HomeDataDevice, HomeDataProduct]]:
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):
Expand Down Expand Up @@ -388,8 +403,13 @@ class NamedRoomMapping(RoomMapping):
from the HomeData based on the iot_id from the room.
"""

name: str
"""The human-readable name of the room, if available."""
@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
Expand All @@ -409,6 +429,11 @@ class CombinedMapInfo(RoborockBase):
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):
Expand Down
22 changes: 21 additions & 1 deletion roborock/data/v1/v1_containers.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
)
from roborock.exceptions import RoborockException

from ..containers import RoborockBase, RoborockBaseTimer, _attr_repr
from ..containers import NamedRoomMapping, RoborockBase, RoborockBaseTimer, _attr_repr
from .v1_code_mappings import (
CleanFluidStatus,
ClearWaterBoxStatus,
Expand Down Expand Up @@ -686,6 +686,17 @@ class MultiMapsListRoom(RoborockBase):
iot_name_id: str | None = None
iot_name: str | None = None

@property
def named_room_mapping(self) -> NamedRoomMapping | None:
"""Returns a NamedRoomMapping object if valid."""
if self.id is None or self.iot_name_id is None:
return None
return NamedRoomMapping(
segment_id=self.id,
iot_id=self.iot_name_id,
raw_name=self.iot_name,
)


@dataclass
class MultiMapsListMapInfoBakMaps(RoborockBase):
Expand All @@ -707,6 +718,15 @@ def mapFlag(self) -> int:
"""Alias for map_flag, returns the map flag as an integer."""
return self.map_flag

@property
def rooms_map(self) -> dict[int, NamedRoomMapping]:
"""Returns a dictionary of room mappings by segment id."""
return {
room.id: room.named_room_mapping
for room in self.rooms or ()
if room.id is not None and room.named_room_mapping is not None
Comment thread
allenporter marked this conversation as resolved.
Outdated
}


@dataclass
class MultiMapsList(RoborockBase):
Expand Down
43 changes: 19 additions & 24 deletions roborock/devices/traits/v1/home.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@
import logging
from typing import Self

from roborock.data import CombinedMapInfo, NamedRoomMapping, RoborockBase
from roborock.data import CombinedMapInfo, MultiMapsListMapInfo, RoborockBase
from roborock.data.containers import NamedRoomMapping
Comment thread
allenporter marked this conversation as resolved.
Outdated
from roborock.data.v1.v1_code_mappings import RoborockStateCode
from roborock.devices.cache import DeviceCache
from roborock.devices.traits.v1 import common
Expand Down Expand Up @@ -114,35 +115,29 @@ async def discover_home(self) -> None:
self._discovery_completed = True
await self._update_home_cache(home_map_info, home_map_content)

async def _refresh_map_info(self, map_info) -> CombinedMapInfo:
async def _refresh_map_info(self, map_info: MultiMapsListMapInfo) -> CombinedMapInfo:
"""Collect room data for a specific map and return CombinedMapInfo."""
await self._rooms_trait.refresh()

rooms: dict[int, NamedRoomMapping] = {}
if map_info.rooms:
# Not all vacuums respond with rooms inside map_info.
# If we can determine if all vacuums will return everything with get_rooms, we could remove this step.
for room in map_info.rooms:
if room.id is not None and room.iot_name_id is not None:
rooms[room.id] = NamedRoomMapping(
segment_id=room.id,
iot_id=room.iot_name_id,
name=room.iot_name or f"Room {room.id}",
)

# Add rooms from rooms_trait.
# Keep existing names from map_info unless they are fallback names.
if self._rooms_trait.rooms:
for room in self._rooms_trait.rooms:
if room.segment_id is not None and room.name:
existing_room = rooms.get(room.segment_id)
if existing_room is None or existing_room.name == f"Room {room.segment_id}":
rooms[room.segment_id] = room

# We have room names from two sources. The map_info.rooms which we just
# received from the MultiMapsList or the self._rooms_trait.rooms which
# comes from the GET_ROOM_MAPPING command. We merge them, giving priority
Comment thread
allenporter marked this conversation as resolved.
Outdated
# to map_info rooms, and excluding unset rooms from the trait.
rooms: list[NamedRoomMapping] = list(
Comment thread
allenporter marked this conversation as resolved.
Outdated
{
**map_info.rooms_map,
**{
room.segment_id: room
for room in self._rooms_trait.rooms or ()
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could it be worthwhile to set .rooms to a property? I'm figuring we do this other places as well, may be easier to just know we can always assume it is a list. Not needed for this PR though, just a thought

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be current_rooms?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that makes sense!

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This depends on one of the map infos returned from an RPC before it is saved so I don't think this works. I still added the attribute, but it won't make this code simpler.

if (existing := map_info.rooms_map.get(room.segment_id)) is None
or (room.raw_name and existing.raw_name is None)
Comment thread
allenporter marked this conversation as resolved.
Outdated
},
}.values()
)
return CombinedMapInfo(
map_flag=map_info.map_flag,
name=map_info.name,
rooms=list(rooms.values()),
rooms=rooms,
)

async def _refresh_map_content(self) -> MapContent:
Expand Down
52 changes: 22 additions & 30 deletions roborock/devices/traits/v1/rooms.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import logging
from dataclasses import dataclass

from roborock.data import HomeData, NamedRoomMapping, RoborockBase
from roborock.data import HomeData, HomeDataRoom, NamedRoomMapping, RoborockBase
from roborock.devices.traits.v1 import common
from roborock.roborock_typing import RoborockCommand
from roborock.web_api import UserWebApiClient
Expand Down Expand Up @@ -36,7 +36,7 @@ def __init__(self, home_data: HomeData, web_api: UserWebApiClient) -> None:
super().__init__()
self._home_data = home_data
self._web_api = web_api
self._seen_unknown_room_iot_ids: set[str] = set()
self._discovered_iot_ids: set[str] = set()

async def refresh(self) -> None:
"""Refresh room mappings and backfill unknown room names from the web API."""
Expand All @@ -45,47 +45,39 @@ async def refresh(self) -> None:
raise ValueError(f"Unexpected RoomsTrait response format: {response!r}")

segment_map = _extract_segment_map(response)
await self._populate_missing_home_data_rooms(segment_map)

new_data = self._parse_response(response, segment_map)
# Track all iot ids seen before. Refresh the room list when new ids are found.
new_iot_ids = set(segment_map.values()) - set(self._home_data.rooms_map.keys())
if new_iot_ids - self._discovered_iot_ids:
_LOGGER.debug("Refreshing room list to discover new room names")
if updated_rooms := await self._refresh_rooms():
_LOGGER.debug("Updating rooms: %s", list(updated_rooms))
self._home_data.rooms = updated_rooms
self._discovered_iot_ids.update(new_iot_ids)

new_data = self._parse_rooms(segment_map, self._home_data.rooms_name_map)
self._update_trait_values(new_data)
_LOGGER.debug("Refreshed %s: %s", self.__class__.__name__, new_data)

@property
def _iot_id_room_name_map(self) -> dict[str, str]:
"""Returns a dictionary of Room IOT IDs to room names."""
return {str(room.id): room.name for room in self._home_data.rooms or ()}

def _parse_response(self, response: common.V1ResponseData, segment_map: dict[int, str] | None = None) -> Rooms:
@staticmethod
def _parse_rooms(
segment_map: dict[int, str],
name_map: dict[str, str],
) -> Rooms:
"""Parse the response from the device into a list of NamedRoomMapping."""
if not isinstance(response, list):
raise ValueError(f"Unexpected RoomsTrait response format: {response!r}")
if segment_map is None:
segment_map = _extract_segment_map(response)
name_map = self._iot_id_room_name_map
return Rooms(
rooms=[
NamedRoomMapping(segment_id=segment_id, iot_id=iot_id, name=name_map.get(iot_id, f"Room {segment_id}"))
NamedRoomMapping(segment_id=segment_id, iot_id=iot_id, raw_name=name_map.get(iot_id))
for segment_id, iot_id in segment_map.items()
]
)

async def _populate_missing_home_data_rooms(self, segment_map: dict[int, str]) -> None:
"""Load missing room names into home data for newly-seen unknown room ids."""
missing_room_iot_ids = set(segment_map.values()) - set(self._iot_id_room_name_map.keys())
new_missing_room_iot_ids = missing_room_iot_ids - self._seen_unknown_room_iot_ids
if not new_missing_room_iot_ids:
return

async def _refresh_rooms(self) -> list[HomeDataRoom]:
"""Merge the existing rooms with new rooms returned from the API."""
Comment thread
allenporter marked this conversation as resolved.
Outdated
try:
web_rooms = await self._web_api.get_rooms()
return await self._web_api.get_rooms()
except Exception:
_LOGGER.debug("Failed to fetch rooms from web API", exc_info=True)
else:
if isinstance(web_rooms, list) and web_rooms:
self._home_data.rooms = web_rooms

self._seen_unknown_room_iot_ids.update(missing_room_iot_ids)
return []


def _extract_segment_map(response: list) -> dict[int, str]:
Expand Down
Loading
Loading