fix: update DeviceCache interface

Author: allenporter

roborock/cli.py

@@ -42,7 +42,7 @@
 from pyshark.packet.packet import Packet  # type: ignore
 
 from roborock import SHORT_MODEL_TO_ENUM, RoborockCommand
-from roborock.data import CombinedMapInfo, DeviceData, HomeData, NetworkInfo, RoborockBase, UserData
+from roborock.data import DeviceData, RoborockBase, UserData
 from roborock.device_features import DeviceFeatures
 from roborock.devices.cache import Cache, CacheData
 from roborock.devices.device import RoborockDevice
@@ -116,10 +116,8 @@ class ConnectionCache(RoborockBase):
 
     user_data: UserData
     email: str
-    home_data: HomeData | None = None
-    network_info: dict[str, NetworkInfo] | None = None
-    home_map_info: dict[int, CombinedMapInfo] | None = None
-    trait_data: dict[str, Any] | None = None
+    # TODO: Used new APIs for cache file storage
+    cache_data: CacheData | None = None
 
 
 class DeviceConnectionManager:
@@ -134,10 +132,10 @@ def __init__(self, context: "RoborockContext", loop: asyncio.AbstractEventLoop |
     async def ensure_device_manager(self) -> DeviceManager:
         """Ensure device manager is initialized."""
         if self.device_manager is None:
-            cache_data = self.context.cache_data()
+            connection_cache = self.context.connection_cache()
             user_params = UserParams(
-                username=cache_data.email,
-                user_data=cache_data.user_data,
+                username=connection_cache.email,
+                user_data=connection_cache.user_data,
             )
             self.device_manager = await create_device_manager(user_params, cache=self.context)
             # Cache devices for quick lookup
@@ -164,7 +162,8 @@ class RoborockContext(Cache):
     """Context that handles both CLI and session modes internally."""
 
     roborock_file = Path("~/.roborock").expanduser()
-    _cache_data: ConnectionCache | None = None
+    roborock_cache_file = Path("~/.roborock.cache").expanduser()
+    _connection_cache: ConnectionCache | None = None
 
     def __init__(self):
         self.reload()
@@ -177,22 +176,22 @@ def reload(self):
             with open(self.roborock_file) as f:
                 data = json.load(f)
                 if data:
-                    self._cache_data = ConnectionCache.from_dict(data)
+                    self._connection_cache = ConnectionCache.from_dict(data)
 
-    def update(self, cache_data: ConnectionCache):
-        data = json.dumps(cache_data.as_dict(), default=vars, indent=4)
+    def update(self, connection_cache: ConnectionCache):
+        data = json.dumps(connection_cache.as_dict(), default=vars, indent=4)
         with open(self.roborock_file, "w") as f:
             f.write(data)
         self.reload()
 
     def validate(self):
-        if self._cache_data is None:
+        if self._connection_cache is None:
             raise RoborockException("You must login first")
 
-    def cache_data(self) -> ConnectionCache:
+    def connection_cache(self) -> ConnectionCache:
         """Get the cache data."""
         self.validate()
-        return cast(ConnectionCache, self._cache_data)
+        return cast(ConnectionCache, self._connection_cache)
 
     def start_session_mode(self):
         """Start session mode with a background event loop."""
@@ -229,19 +228,21 @@ async def get_device_manager(self) -> DeviceConnectionManager:
 
     async def refresh_devices(self) -> ConnectionCache:
         """Refresh device data from server (always fetches fresh data)."""
-        cache_data = self.cache_data()
-        client = RoborockApiClient(cache_data.email)
-        home_data = await client.get_home_data_v3(cache_data.user_data)
-        cache_data.home_data = home_data
-        self.update(cache_data)
-        return cache_data
+        connection_cache = self.connection_cache()
+        client = RoborockApiClient(connection_cache.email)
+        home_data = await client.get_home_data_v3(connection_cache.user_data)
+        if connection_cache.cache_data is None:
+            connection_cache.cache_data = CacheData()
+        connection_cache.cache_data.home_data = home_data
+        self.update(connection_cache)
+        return connection_cache
 
     async def get_devices(self) -> ConnectionCache:
         """Get device data (uses cache if available, fetches if needed)."""
-        cache_data = self.cache_data()
-        if not cache_data.home_data:
-            cache_data = await self.refresh_devices()
-        return cache_data
+        connection_cache = self.connection_cache()
+        if (connection_cache.cache_data is None) or (connection_cache.cache_data.home_data is None):
+            connection_cache = await self.refresh_devices()
+        return connection_cache
 
     async def cleanup(self):
         """Clean up resources (mainly for session mode)."""
@@ -266,22 +267,16 @@ 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 {},
-            home_map_info=connection_cache.home_map_info,
-            trait_data=connection_cache.trait_data or {},
-        )
+        connection_cache = self.connection_cache()
+        if connection_cache.cache_data is not None:
+            return connection_cache.cache_data
+        return CacheData()
 
     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_map_info = value.home_map_info
-        connection_cache.trait_data = value.trait_data
+        connection_cache = self.connection_cache()
+        connection_cache.cache_data = value
         self.update(connection_cache)
 
 
@@ -367,9 +362,9 @@ async def discover(ctx):
 @async_command
 async def list_devices(ctx):
     context: RoborockContext = ctx.obj
-    cache_data = await context.get_devices()
+    connection_cache = await context.get_devices()
 
-    home_data = cache_data.home_data
+    home_data = connection_cache.cache_data.home_data
 
     device_name_id = {device.name: device.duid for device in home_data.get_all_devices()}
     click.echo(json.dumps(device_name_id, indent=4))

roborock/devices/cache.py

@@ -20,11 +20,11 @@ class DeviceCacheData(RoborockBase):
     network_info: NetworkInfo | None = None
     """Network information for the device"""
 
-    home_map_info: CombinedMapInfo | None = None
-    """Home map information for the device."""
+    home_map_info: dict[int, CombinedMapInfo] | None = None
+    """Home map information for the device by map_flag."""
 
-    home_map_content_base64: str | None = None
-    """Home cache content for the device (encoded base64)."""
+    home_map_content_base64: dict[int, str] | None = None
+    """Home cache content for the device (encoded base64) by map_flag."""
 
     device_features: DeviceFeatures | None = None
     """Device features information."""
@@ -91,19 +91,33 @@ async def set(self, value: CacheData) -> None:
         """Set value in the cache."""
         ...
 
-    async def get_device_info(self, duid: str) -> DeviceCacheData:
+
+@dataclass
+class DeviceCache(RoborockBase):
+    """Provides a cache interface for a specific device.
+
+    This is a convenience wrapper around a general Cache implementation to
+    provide device-specific caching functionality.
+    """
+
+    def __init__(self, duid: str, cache: Cache) -> None:
+        """Initialize the device cache with the given cache implementation."""
+        self._duid = duid
+        self._cache = cache
+
+    async def get(self) -> DeviceCacheData:
         """Get cached device-specific information."""
-        cache_data = await self.get()
-        if duid not in cache_data.device_info:
-            cache_data.device_info[duid] = DeviceCacheData()
-            await self.set(cache_data)
-        return cache_data.device_info[duid]
+        cache_data = await self._cache.get()
+        if self._duid not in cache_data.device_info:
+            cache_data.device_info[self._duid] = DeviceCacheData()
+            await self._cache.set(cache_data)
+        return cache_data.device_info[self._duid]
 
-    async def set_device_info(self, duid: str, device_cache_data: DeviceCacheData) -> None:
+    async def set(self, device_cache_data: DeviceCacheData) -> None:
         """Set cached device-specific information."""
-        cache_data = await self.get()
-        cache_data.device_info[duid] = device_cache_data
-        await self.set(cache_data)
+        cache_data = await self._cache.get()
+        cache_data.device_info[self._duid] = device_cache_data
+        await self._cache.set(cache_data)
 
 
 class InMemoryCache(Cache):

roborock/devices/device_manager.py

@@ -21,7 +21,7 @@
 from roborock.protocol import create_mqtt_params
 from roborock.web_api import RoborockApiClient, UserWebApiClient
 
-from .cache import Cache, NoCache
+from .cache import Cache, DeviceCache, NoCache
 from .channel import Channel
 from .mqtt_channel import create_mqtt_channel
 from .traits import Trait, a01, b01, v1
@@ -177,9 +177,10 @@ async def create_device_manager(
     def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice:
         channel: Channel
         trait: Trait
+        device_cache: DeviceCache = DeviceCache(device.duid, cache)
         match device.pv:
             case DeviceVersion.V1:
-                channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
+                channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, device_cache)
                 trait = v1.create(
                     device.duid,
                     product,
@@ -188,7 +189,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                     channel.mqtt_rpc_channel,
                     channel.map_rpc_channel,
                     web_api,
-                    cache,
+                    device_cache=device_cache,
                     map_parser_config=map_parser_config,
                 )
             case DeviceVersion.A01:

roborock/devices/traits/v1/__init__.py

@@ -32,11 +32,12 @@
 
 import logging
 from dataclasses import dataclass, field, fields
+from functools import cache
 from typing import Any, get_args
 
 from roborock.data.containers import HomeData, HomeDataProduct, RoborockBase
 from roborock.data.v1.v1_code_mappings import RoborockDockTypeCode
-from roborock.devices.cache import Cache
+from roborock.devices.cache import DeviceCache
 from roborock.devices.traits import Trait
 from roborock.map.map_parser import MapParserConfig
 from roborock.protocols.v1_protocol import V1RpcChannel
@@ -131,7 +132,7 @@ def __init__(
         mqtt_rpc_channel: V1RpcChannel,
         map_rpc_channel: V1RpcChannel,
         web_api: UserWebApiClient,
-        cache: Cache,
+        device_cache: DeviceCache,
         map_parser_config: MapParserConfig | None = None,
     ) -> None:
         """Initialize the V1TraitProps."""
@@ -140,16 +141,16 @@ def __init__(
         self._mqtt_rpc_channel = mqtt_rpc_channel
         self._map_rpc_channel = map_rpc_channel
         self._web_api = web_api
-        self._cache = cache
+        self._device_cache = device_cache
 
         self.status = StatusTrait(product)
         self.consumables = ConsumableTrait()
         self.rooms = RoomsTrait(home_data)
         self.maps = MapsTrait(self.status)
         self.map_content = MapContentTrait(map_parser_config)
-        self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, cache)
-        self.device_features = DeviceFeaturesTrait(product.product_nickname, cache)
-        self.network_info = NetworkInfoTrait(device_uid, cache)
+        self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, self._device_cache)
+        self.device_features = DeviceFeaturesTrait(product.product_nickname, self._device_cache)
+        self.network_info = NetworkInfoTrait(device_uid, self._device_cache)
         self.routines = RoutinesTrait(device_uid, web_api)
 
         # Dynamically create any traits that need to be populated
@@ -239,20 +240,20 @@ async def _dock_type(self) -> RoborockDockTypeCode:
 
     async def _get_cached_trait_data(self, name: str) -> Any:
         """Get the dock type from the status trait or cache."""
-        cache_data = await self._cache.get()
+        cache_data = await self._device_cache.get()
         if cache_data.trait_data is None:
             cache_data.trait_data = {}
         _LOGGER.debug("Cached trait data: %s", cache_data.trait_data)
         return cache_data.trait_data.get(name)
 
     async def _set_cached_trait_data(self, name: str, value: Any) -> None:
         """Set trait-specific cached data."""
-        cache_data = await self._cache.get()
+        cache_data = await self._device_cache.get()
         if cache_data.trait_data is None:
             cache_data.trait_data = {}
         cache_data.trait_data[name] = value
         _LOGGER.debug("Updating cached trait data: %s", cache_data.trait_data)
-        await self._cache.set(cache_data)
+        await self._device_cache.set(cache_data)
 
     def as_dict(self) -> dict[str, Any]:
         """Return the trait data as a dictionary."""
@@ -275,7 +276,7 @@ def create(
     mqtt_rpc_channel: V1RpcChannel,
     map_rpc_channel: V1RpcChannel,
     web_api: UserWebApiClient,
-    cache: Cache,
+    device_cache: DeviceCache,
     map_parser_config: MapParserConfig | None = None,
 ) -> PropertiesApi:
     """Create traits for V1 devices."""
@@ -287,6 +288,6 @@ def create(
         mqtt_rpc_channel,
         map_rpc_channel,
         web_api,
-        cache,
+        device_cache,
         map_parser_config,
     )

roborock/devices/traits/v1/device_features.py

@@ -2,7 +2,7 @@
 
 from roborock.data import AppInitStatus, RoborockProductNickname
 from roborock.device_features import DeviceFeatures
-from roborock.devices.cache import Cache
+from roborock.devices.cache import DeviceCache
 from roborock.devices.traits.v1 import common
 from roborock.roborock_typing import RoborockCommand
 
@@ -12,10 +12,10 @@ class DeviceFeaturesTrait(DeviceFeatures, common.V1TraitMixin):
 
     command = RoborockCommand.APP_GET_INIT_STATUS
 
-    def __init__(self, product_nickname: RoborockProductNickname, cache: Cache) -> None:  # pylint: disable=super-init-not-called
+    def __init__(self, product_nickname: RoborockProductNickname, device_cache: DeviceCache) -> None:  # pylint: disable=super-init-not-called
         """Initialize MapContentTrait."""
         self._nickname = product_nickname
-        self._cache = cache
+        self._device_cache = device_cache
         # All fields of DeviceFeatures are required. Initialize them to False
         # so we have some known state.
         for field in fields(self):
@@ -28,14 +28,14 @@ async def refresh(self) -> None:
         change often and this avoids unnecessary RPC calls. This would only
         ever change with a firmware update, so caching is appropriate.
         """
-        cache_data = await self._cache.get()
+        cache_data = await self._device_cache.get()
         if cache_data.device_features is not None:
             self._update_trait_values(cache_data.device_features)
             return
         # Save cached device features
         await super().refresh()
         cache_data.device_features = self
-        await self._cache.set(cache_data)
+        await self._device_cache.set(cache_data)
 
     def _parse_response(self, response: common.V1ResponseData) -> DeviceFeatures:
         """Parse the response from the device into a MapContentTrait instance."""