Add a Home trait that is responsible for caching information about maps and rooms

roborock/cli.py

@@ -42,7 +42,7 @@
 from pyshark.packet.packet import Packet  # type: ignore
 
 from roborock import SHORT_MODEL_TO_ENUM, DeviceFeatures, RoborockCommand, RoborockException
-from roborock.containers import DeviceData, HomeData, NetworkInfo, RoborockBase, UserData
+from roborock.containers import CombinedMapInfo, DeviceData, HomeData, NetworkInfo, RoborockBase, UserData
 from roborock.devices.cache import Cache, CacheData
 from roborock.devices.device import RoborockDevice
 from roborock.devices.device_manager import DeviceManager, create_device_manager, create_home_data_api
@@ -116,6 +116,7 @@ class ConnectionCache(RoborockBase):
     email: str
     home_data: HomeData | None = None
     network_info: dict[str, NetworkInfo] | None = None
+    home_cache: dict[int, CombinedMapInfo] | None = None
 
 
 class DeviceConnectionManager:
@@ -258,14 +259,21 @@ def finish_session(self) -> None:
 
     async def get(self) -> CacheData:
         """Get cached value."""
+        _LOGGER.debug("Getting cache data")
         connection_cache = self.cache_data()
-        return CacheData(home_data=connection_cache.home_data, network_info=connection_cache.network_info or {})
+        return CacheData(
+            home_data=connection_cache.home_data,
+            network_info=connection_cache.network_info or {},
+            home_cache=connection_cache.home_cache,
+        )
 
     async def set(self, value: CacheData) -> None:
         """Set value in the cache."""
+        _LOGGER.debug("Setting cache data")
         connection_cache = self.cache_data()
         connection_cache.home_data = value.home_data
         connection_cache.network_info = value.network_info
+        connection_cache.home_cache = value.home_cache
         self.update(connection_cache)
 
 
@@ -533,6 +541,42 @@ async def rooms(ctx, device_id: str):
     await _display_v1_trait(context, device_id, lambda v1: v1.rooms)
 
 
+@session.command()
+@click.option("--device_id", required=True)
+@click.option("--refresh", is_flag=True, default=False, help="Refresh status before discovery.")
+@click.pass_context
+@async_command
+async def home(ctx, device_id: str, refresh: bool):
+    """Discover and cache home layout (maps and rooms)."""
+    context: RoborockContext = ctx.obj
+    device_manager = await context.get_device_manager()
+    device = await device_manager.get_device(device_id)
+    if device.v1_properties is None:
+        raise RoborockException(f"Device {device.name} does not support V1 protocol")
+
+    # Ensure we have the latest status before discovery
+    await device.v1_properties.status.refresh()
+
+    home_trait = device.v1_properties.home
+    await home_trait.discover_home()
+    if refresh:
+        await home_trait.refresh()
+
+    # Display the discovered home cache
+    if home_trait.home_cache:
+        cache_summary = {
+            map_flag: {
+                "name": map_data.name,
+                "room_count": len(map_data.rooms),
+                "rooms": [{"segment_id": room.segment_id, "name": room.name} for room in map_data.rooms],
+            }
+            for map_flag, map_data in home_trait.home_cache.items()
+        }
+        click.echo(dump_json(cache_summary))
+    else:
+        click.echo("No maps discovered")
+
+
 @click.command()
 @click.option("--device_id", required=True)
 @click.option("--cmd", required=True)
@@ -780,6 +824,7 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur
 cli.add_command(consumables)
 cli.add_command(reset_consumable)
 cli.add_command(rooms)
+cli.add_command(home)
 
 
 def main():

roborock/containers.py

@@ -880,6 +880,36 @@ class RoomMapping(RoborockBase):
     iot_id: str
 
 
+@dataclass
+class NamedRoomMapping(RoomMapping):
+    """Dataclass representing a mapping of a room segment to a name.
+
+    The name information is not provided by the device directly, but is provided
+    from the HomeData based on the iot_id from the room.
+    """
+
+    name: str
+    """The human-readable name of the room, if available."""
+
+
+@dataclass
+class CombinedMapInfo(RoborockBase):
+    """Data structure for caching home information.
+
+    This is not provided directly by the API, but is a combination of map data
+    and room data to provide a more useful structure.
+    """
+
+    map_flag: int
+    """The map identifier."""
+
+    name: str
+    """The name of the map from MultiMapsListMapInfo."""
+
+    rooms: list[NamedRoomMapping]
+    """The list of rooms in the map."""
+
+
 @dataclass
 class ChildLockStatus(RoborockBase):
     lock_status: int

roborock/devices/cache.py

@@ -8,7 +8,7 @@
 from dataclasses import dataclass, field
 from typing import Protocol
 
-from roborock.containers import HomeData, NetworkInfo
+from roborock.containers import CombinedMapInfo, HomeData, NetworkInfo
 
 
 @dataclass
@@ -21,6 +21,9 @@ class CacheData:
     network_info: dict[str, NetworkInfo] = field(default_factory=dict)
     """Network information indexed by device DUID."""
 
+    home_cache: dict[int, CombinedMapInfo] = field(default_factory=dict)
+    """Home cache information indexed by map_flag."""
+
 
 class Cache(Protocol):
     """Protocol for a cache that can store and retrieve values."""

roborock/devices/device_manager.py

@@ -157,6 +157,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                     channel.rpc_channel,
                     channel.mqtt_rpc_channel,
                     channel.map_rpc_channel,
+                    cache,
                     map_parser_config=map_parser_config,
                 )
             case DeviceVersion.A01:

roborock/devices/traits/v1/__init__.py

@@ -4,6 +4,7 @@
 from dataclasses import dataclass, field, fields
 
 from roborock.containers import HomeData, HomeDataProduct
+from roborock.devices.cache import Cache
 from roborock.devices.traits import Trait
 from roborock.devices.v1_rpc_channel import V1RpcChannel
 from roborock.map.map_parser import MapParserConfig
@@ -12,6 +13,7 @@
 from .common import V1TraitMixin
 from .consumeable import ConsumableTrait
 from .do_not_disturb import DoNotDisturbTrait
+from .home import HomeTrait
 from .map_content import MapContentTrait
 from .maps import MapsTrait
 from .rooms import RoomsTrait
@@ -30,6 +32,7 @@
     "MapsTrait",
     "MapContentTrait",
     "ConsumableTrait",
+    "HomeTrait",
 ]
 
 
@@ -49,6 +52,7 @@ class PropertiesApi(Trait):
     maps: MapsTrait
     map_content: MapContentTrait
     consumables: ConsumableTrait
+    home: HomeTrait
 
     # In the future optional fields can be added below based on supported features
 
@@ -59,13 +63,15 @@ def __init__(
         rpc_channel: V1RpcChannel,
         mqtt_rpc_channel: V1RpcChannel,
         map_rpc_channel: V1RpcChannel,
+        cache: Cache,
         map_parser_config: MapParserConfig | None = None,
     ) -> None:
         """Initialize the V1TraitProps."""
         self.status = StatusTrait(product)
         self.rooms = RoomsTrait(home_data)
         self.maps = MapsTrait(self.status)
         self.map_content = MapContentTrait(map_parser_config)
+        self.home = HomeTrait(self.maps, self.rooms, cache)
         # This is a hack to allow setting the rpc_channel on all traits. This is
         # used so we can preserve the dataclass behavior when the values in the
         # traits are updated, but still want to allow them to have a reference
@@ -90,7 +96,8 @@ def create(
     rpc_channel: V1RpcChannel,
     mqtt_rpc_channel: V1RpcChannel,
     map_rpc_channel: V1RpcChannel,
+    cache: Cache,
     map_parser_config: MapParserConfig | None = None,
 ) -> PropertiesApi:
     """Create traits for V1 devices."""
-    return PropertiesApi(product, home_data, rpc_channel, mqtt_rpc_channel, map_rpc_channel, map_parser_config)
+    return PropertiesApi(product, home_data, rpc_channel, mqtt_rpc_channel, map_rpc_channel, cache, map_parser_config)

roborock/devices/traits/v1/home.py

@@ -0,0 +1,157 @@
+"""Trait that represents a full view of the home layout.
+
+This trait combines information about maps and rooms to provide a comprehensive
+view of the home layout, including room names and their corresponding segment
+on the map. It also makes it straight forward to fetch the map image and data.
+
+This trait depends on the MapsTrait and RoomsTrait to gather the necessary
+information. It provides properties to access the current map, the list of
+rooms with names, and the map image and data.
+"""
+
+import asyncio
+import logging
+from typing import Self
+
+from roborock.containers import CombinedMapInfo, RoborockBase
+from roborock.devices.cache import Cache
+from roborock.devices.traits.v1 import common
+from roborock.exceptions import RoborockException
+from roborock.roborock_typing import RoborockCommand
+
+from .maps import MapsTrait
+from .rooms import RoomsTrait
+
+_LOGGER = logging.getLogger(__name__)
+
+MAP_SLEEP = 3
+
+
+class HomeTrait(RoborockBase, common.V1TraitMixin):
+    """Trait that represents a full view of the home layout."""
+
+    command = RoborockCommand.GET_MAP_V1  # This is not used
+
+    def __init__(self, maps_trait: MapsTrait, rooms_trait: RoomsTrait, cache: Cache) -> None:
+        """Initialize the HomeTrait.
+
+        We keep track of the MapsTrait and RoomsTrait to provide a comprehensive
+        view of the home layout. This also depends on the StatusTrait to determine
+        the current map. See comments in MapsTrait for details on that dependency.
+
+        The cache is used to store discovered home data to minimize map switching
+        and improve performance. The cache should be persisted by the caller to
+        ensure data is retained across restarts.
+
+        After initial discovery, only information for the current map is refreshed
+        to keep data up to date without excessive map switching. However, as
+        users switch rooms, the current map's data will be updated to ensure
+        accuracy.
+        """
+        super().__init__()
+        self._maps_trait = maps_trait
+        self._rooms_trait = rooms_trait
+        self._cache = cache
+        self._home_cache: dict[int, CombinedMapInfo] | None = None
+
+    async def discover_home(self) -> None:
+        """Iterate through all maps to discover rooms and cache them.
+
+        This will be a no-op if the home cache is already populated.
+        """
+        cache_data = await self._cache.get()
+        if cache_data.home_cache:
+            _LOGGER.debug("Home cache already populated, skipping discovery")
+            self._home_cache = cache_data.home_cache
+            return
+
+        await self._maps_trait.refresh()
+        if self._maps_trait.current_map_info is None:
+            raise RoborockException("Cannot perform home discovery without current map info")
+
+        home_cache = await self._build_home_cache()
+        _LOGGER.debug("Home discovery complete, caching data for %d maps", len(home_cache))
+        cache_data = await self._cache.get()
+        cache_data.home_cache = home_cache
+        await self._cache.set(cache_data)
+        self._home_cache = home_cache
+
+    async def _refresh_map_data(self, map_info) -> CombinedMapInfo:
+        """Collect room data for a specific map and return CombinedMapInfo."""
+        await self._rooms_trait.refresh()
+        return CombinedMapInfo(
+            map_flag=map_info.map_flag,
+            name=map_info.name,
+            rooms=self._rooms_trait.rooms or [],
+        )
+
+    async def _build_home_cache(self) -> dict[int, CombinedMapInfo]:
+        """Perform the actual discovery and caching of home data."""
+        home_cache: dict[int, CombinedMapInfo] = {}
+
+        # Sort map_info to process the current map last, reducing map switching.
+        # False (non-original maps) sorts before True (original map). We ensure
+        # we load the original map last.
+        sorted_map_infos = sorted(
+            self._maps_trait.map_info or [],
+            key=lambda mi: mi.map_flag == self._maps_trait.current_map,
+            reverse=False,
+        )
+        _LOGGER.debug("Building home cache for maps: %s", [mi.map_flag for mi in sorted_map_infos])
+        for map_info in sorted_map_infos:
+            # We need to load each map to get its room data
+            if len(sorted_map_infos) > 1:
+                _LOGGER.debug("Loading map %s", map_info.map_flag)
+                await self._maps_trait.set_current_map(map_info.map_flag)
+                await asyncio.sleep(MAP_SLEEP)
+
+            map_data = await self._refresh_map_data(map_info)
+            home_cache[map_info.map_flag] = map_data
+        return home_cache
+
+    async def refresh(self) -> Self:
+        """Refresh current map's underlying map and room data, updating cache as needed.
+
+        This will only refresh the current map's data and will not populate non
+        active maps or re-discover the home. It is expected that this will keep
+        information up to date for the current map as users switch to that map.
+        """
+        if self._home_cache is None:
+            raise RoborockException("Cannot refresh home data without home cache, did you call discover_home()?")
+
+        # Refresh the list of map names/info
+        await self._maps_trait.refresh()
+        if (current_map_info := self._maps_trait.current_map_info) is None or (
+            map_flag := self._maps_trait.current_map
+        ) is None:
+            raise RoborockException("Cannot refresh home data without current map info")
+
+        # Refresh the current map's room data
+        current_map_data = self._home_cache.get(map_flag)
+        if current_map_data:
+            map_data = await self._refresh_map_data(current_map_info)
+            if map_data != current_map_data:
+                self._home_cache[map_flag] = map_data
+                # Persist to cache
+                cache_data = await self._cache.get()
+                cache_data.home_cache = self._home_cache
+                await self._cache.set(cache_data)
+
+        return self
+
+    @property
+    def home_cache(self) -> dict[int, CombinedMapInfo] | None:
+        """Returns the map information for all cached maps."""
+        return self._home_cache
+
+    @property
+    def current_map_data(self) -> CombinedMapInfo | None:
+        """Returns the map data for the current map."""
+        current_map_flag = self._maps_trait.current_map
+        if current_map_flag is None or self._home_cache is None:
+            return None
+        return self._home_cache.get(current_map_flag)
+
+    def _parse_response(self, response: common.V1ResponseData) -> Self:
+        """This trait does not parse responses directly."""
+        raise NotImplementedError("HomeTrait does not support direct command responses")

roborock/devices/traits/v1/maps.py

@@ -22,6 +22,14 @@ class MapsTrait(MultiMapsList, common.V1TraitMixin):
 
     A device may have multiple maps, each identified by a unique map_flag.
     Each map can have multiple rooms associated with it, in a `RoomMapping`.
+
+    The MapsTrait depends on the StatusTrait to determine the currently active
+    map. It is the responsibility of the caller to ensure that the StatusTrait
+    is up to date before using this trait. However, there is a possibility of
+    races if another client changes the current map between the time the
+    StatusTrait is refreshed and when the MapsTrait is used. This is mitigated
+    by the fact that the map list is unlikely to change frequently, and the
+    current map is only changed when the user explicitly switches maps.
     """
 
     command = RoborockCommand.GET_MULTI_MAPS_LIST