feat: adding command cache (#77)

* feat: adding command cache

* chore: typo

* fix: dependencies

* feat: adding cache evict time

Author: humbertogontijo

pyproject.toml

@@ -23,10 +23,10 @@ python = "^3.10"
 click = ">=8"
 aiohttp = "^3.8.2"
 async-timeout = "*"
-pycryptodome = ">=3.17,<3.19"
-pycryptodomex = {version = ">=3.17,<3.19", markers = "sys_platform == 'darwin'"}
-paho-mqtt = "~1.6.1"
-dacite = "~1.8.0"
+pycryptodome = "^3.18"
+pycryptodomex = {version = "^3.18", markers = "sys_platform == 'darwin'"}
+paho-mqtt = "^1.6.1"
+dacite = "^1.8.0"
 construct = "^2.10.68"
 alexapy = "^1.26.8"
 

roborock/api.py

@@ -13,11 +13,12 @@
 import struct
 import time
 from random import randint
-from typing import Any, Callable, Coroutine, Optional, Type, TypeVar
+from typing import Any, Callable, Coroutine, Optional, Type, TypeVar, final
 
 import aiohttp
 
 from .code_mappings import RoborockDockTypeCode
+from .command_cache import AttributeCache, CacheableAttribute, CommandType, parse_method
 from .containers import (
     ChildLockStatus,
     CleanRecord,
@@ -56,7 +57,7 @@
 from .roborock_future import RoborockFuture
 from .roborock_message import RoborockDataProtocol, RoborockMessage
 from .roborock_typing import DeviceProp, DockSummary, RoborockCommand
-from .util import fallback_cache, unpack_list
+from .util import unpack_list
 
 _LOGGER = logging.getLogger(__name__)
 KEEPALIVE = 60
@@ -103,8 +104,9 @@ def __init__(self, endpoint: str, device_info: DeviceData) -> None:
         self._last_disconnection = self.time_func()
         self.keep_alive = KEEPALIVE
         self._diagnostic_data: dict[str, dict[str, Any]] = {}
-        self.dnd_timer: DnDTimer | None = None
-        self.valley_timer: ValleyElectricityTimer | None = None
+        self.cache: dict[CacheableAttribute, AttributeCache] = {
+            cacheable_attribute: AttributeCache(cacheable_attribute) for cacheable_attribute in CacheableAttribute
+        }
 
     def __del__(self) -> None:
         self.sync_disconnect()
@@ -240,38 +242,51 @@ def _get_payload(
         )
         return request_id, timestamp, payload
 
+    async def _send_command(
+        self,
+        method: RoborockCommand,
+        params: Optional[list | dict] = None,
+    ):
+        raise NotImplementedError
+
+    @final
     async def send_command(
         self,
         method: RoborockCommand,
         params: Optional[list | dict] = None,
         return_type: Optional[Type[RT]] = None,
     ) -> RT:
-        raise NotImplementedError
+        parsed_method = parse_method(method)
+        if parsed_method is not None and parsed_method.type == CommandType.GET:
+            cache = self.cache[parsed_method.attribute].value
+            if cache is not None:
+                if return_type:
+                    return return_type.from_dict(cache)
+                return cache
+
+        response = await self._send_command(method, params)
+
+        if parsed_method is not None:
+            if parsed_method.type == CommandType.SET:
+                self.cache[parsed_method.attribute].evict()
+            elif parsed_method.type == CommandType.GET:
+                self.cache[parsed_method.attribute].load(response)
+        if return_type:
+            return return_type.from_dict(response)
+        return response
 
-    @fallback_cache
     async def get_status(self) -> Status | None:
         _cls: Type[Status] = ModelStatus.get(
             self.device_info.model, S7MaxVStatus
         )  # Default to S7 MAXV if we don't have the data
         return await self.send_command(RoborockCommand.GET_STATUS, return_type=_cls)
 
-    @fallback_cache
     async def get_dnd_timer(self) -> DnDTimer | None:
-        result = await self.send_command(RoborockCommand.GET_DND_TIMER, return_type=DnDTimer)
-        if result is not None:
-            self.dnd_timer = result
-        return result
+        return await self.send_command(RoborockCommand.GET_DND_TIMER, return_type=DnDTimer)
 
-    @fallback_cache
     async def get_valley_electricity_timer(self) -> ValleyElectricityTimer | None:
-        result = await self.send_command(
-            RoborockCommand.GET_VALLEY_ELECTRICITY_TIMER, return_type=ValleyElectricityTimer
-        )
-        if result is not None:
-            self.valley_timer = result
-        return result
+        return await self.send_command(RoborockCommand.GET_VALLEY_ELECTRICITY_TIMER, return_type=ValleyElectricityTimer)
 
-    @fallback_cache
     async def get_clean_summary(self) -> CleanSummary | None:
         clean_summary: dict | list | int = await self.send_command(RoborockCommand.GET_CLEAN_SUMMARY)
         if isinstance(clean_summary, dict):
@@ -288,27 +303,21 @@ async def get_clean_summary(self) -> CleanSummary | None:
             return CleanSummary(clean_time=clean_summary)
         return None
 
-    @fallback_cache
     async def get_clean_record(self, record_id: int) -> CleanRecord | None:
         return await self.send_command(RoborockCommand.GET_CLEAN_RECORD, [record_id], return_type=CleanRecord)
 
-    @fallback_cache
     async def get_consumable(self) -> Consumable | None:
         return await self.send_command(RoborockCommand.GET_CONSUMABLE, return_type=Consumable)
 
-    @fallback_cache
     async def get_wash_towel_mode(self) -> WashTowelMode | None:
         return await self.send_command(RoborockCommand.GET_WASH_TOWEL_MODE, return_type=WashTowelMode)
 
-    @fallback_cache
     async def get_dust_collection_mode(self) -> DustCollectionMode | None:
         return await self.send_command(RoborockCommand.GET_DUST_COLLECTION_MODE, return_type=DustCollectionMode)
 
-    @fallback_cache
     async def get_smart_wash_params(self) -> SmartWashParams | None:
         return await self.send_command(RoborockCommand.GET_SMART_WASH_PARAMS, return_type=SmartWashParams)
 
-    @fallback_cache
     async def get_dock_summary(self, dock_type: RoborockDockTypeCode) -> DockSummary | None:
         """Gets the status summary from the dock with the methods available for a given dock.
 
@@ -330,7 +339,6 @@ async def get_dock_summary(self, dock_type: RoborockDockTypeCode) -> DockSummary
         )
         return DockSummary(dust_collection_mode, wash_towel_mode, smart_wash_params)
 
-    @fallback_cache
     async def get_prop(self) -> DeviceProp | None:
         """Gets device general properties."""
         [status, clean_summary, consumable] = await asyncio.gather(
@@ -356,15 +364,12 @@ async def get_prop(self) -> DeviceProp | None:
             )
         return None
 
-    @fallback_cache
     async def get_multi_maps_list(self) -> MultiMapsList | None:
         return await self.send_command(RoborockCommand.GET_MULTI_MAPS_LIST, return_type=MultiMapsList)
 
-    @fallback_cache
     async def get_networking(self) -> NetworkInfo | None:
         return await self.send_command(RoborockCommand.GET_NETWORK_INFO, return_type=NetworkInfo)
 
-    @fallback_cache
     async def get_room_mapping(self) -> list[RoomMapping] | None:
         """Gets the mapping from segment id -> iot id. Only works on local api."""
         mapping: list = await self.send_command(RoborockCommand.GET_ROOM_MAPPING)
@@ -375,17 +380,14 @@ async def get_room_mapping(self) -> list[RoomMapping] | None:
             ]
         return None
 
-    @fallback_cache
     async def get_child_lock_status(self) -> ChildLockStatus | None:
         """Gets current child lock status."""
         return await self.send_command(RoborockCommand.GET_CHILD_LOCK_STATUS, return_type=ChildLockStatus)
 
-    @fallback_cache
     async def get_flow_led_status(self) -> FlowLedStatus | None:
         """Gets current flow led status."""
         return await self.send_command(RoborockCommand.GET_FLOW_LED_STATUS, return_type=FlowLedStatus)
 
-    @fallback_cache
     async def get_sound_volume(self) -> int | None:
         """Gets current volume level."""
         return await self.send_command(RoborockCommand.GET_SOUND_VOLUME)

roborock/cloud_api.py

@@ -5,12 +5,12 @@
 import threading
 import uuid
 from asyncio import Lock
-from typing import Optional, Type
+from typing import Optional
 from urllib.parse import urlparse
 
 import paho.mqtt.client as mqtt
 
-from .api import COMMANDS_SECURED, KEEPALIVE, RT, RoborockClient, md5hex
+from .api import COMMANDS_SECURED, KEEPALIVE, RoborockClient, md5hex
 from .containers import DeviceData, UserData
 from .exceptions import CommandVacuumError, RoborockException, VacuumError
 from .protocol import MessageParser, Utils
@@ -149,12 +149,11 @@ def _send_msg_raw(self, msg: bytes) -> None:
         if info.rc != mqtt.MQTT_ERR_SUCCESS:
             raise RoborockException(f"Failed to publish ({mqtt.error_string(info.rc)})")
 
-    async def send_command(
+    async def _send_command(
         self,
         method: RoborockCommand,
         params: Optional[list | dict] = None,
-        return_type: Optional[Type[RT]] = None,
-    ) -> RT:
+    ):
         await self.validate_connection()
         request_id, timestamp, payload = super()._get_payload(method, params, True)
         _LOGGER.debug(f"id={request_id} Requesting method {method} with {params}")
@@ -172,8 +171,6 @@ async def send_command(
             _LOGGER.debug(f"id={request_id} Response from {method}: {len(response)} bytes")
         else:
             _LOGGER.debug(f"id={request_id} Response from {method}: {response}")
-        if return_type:
-            return return_type.from_dict(response)
         return response
 
     async def get_map_v1(self):

roborock/command_cache.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from enum import Enum
+
+GET_PREFIX = "get_"
+SET_PREFIX = ("set_", "change_", "close_")
+
+
+class CacheableAttribute(str, Enum):
+    sound_volume = "sound_volume"
+    camera_status = "camera_status"
+    carpet_clean_mode = "carpet_clean_mode"
+    carpet_mode = "carpet_mode"
+    child_lock_status = "child_lock_status"
+    collision_avoid_status = "collision_avoid_status"
+    customize_clean_mode = "customize_clean_mode"
+    custom_mode = "custom_mode"
+    dnd_timer = "dnd_timer"
+    dust_collection_mode = "dust_collection_mode"
+    flow_led_status = "flow_led_status"
+    identify_furniture_status = "identify_furniture_status"
+    identify_ground_material_status = "identify_ground_material_status"
+    led_status = "led_status"
+    server_timer = "server_timer"
+    smart_wash_params = "smart_wash_params"
+    timezone = "timezone"
+    valley_electricity_timer = "valley_electricity_timer"
+    wash_towel_mode = "wash_towel_mode"
+
+
+EVICT_TIME = 60
+
+
+class AttributeCache:
+    _value: dict | None = None
+    last_update: float | None = None
+
+    def __init__(self, attribute: str):
+        self.attribute = attribute
+
+    @property
+    def value(self):
+        if self.last_update is None or time.monotonic() - self.last_update > EVICT_TIME:
+            self._value = None
+        return self._value
+
+    def load(self, value: dict):
+        self._value = value
+        self.last_update = time.monotonic()
+        return self._value
+
+    def evict(self):
+        self._value = None
+
+
+class CommandType(Enum):
+    OTHER = -1
+    GET = 0
+    SET = 1
+
+
+@dataclass
+class ParserCommand:
+    type: CommandType
+    attribute: CacheableAttribute
+
+
+def parse_method(method: str):
+    if method is not None:
+        attribute = method.lower()
+        command_type = CommandType.OTHER
+        if attribute.startswith(GET_PREFIX):
+            attribute = attribute.removeprefix(GET_PREFIX)
+            command_type = CommandType.GET
+        elif attribute.startswith(SET_PREFIX):
+            for prefix in SET_PREFIX:
+                attribute = attribute.removeprefix(prefix)
+            command_type = CommandType.SET
+        try:
+            cacheable_attribute = CacheableAttribute(attribute)
+            return ParserCommand(type=command_type, attribute=cacheable_attribute)
+        except ValueError:
+            pass
+    return None

roborock/local_api.py

@@ -3,12 +3,12 @@
 import asyncio
 import logging
 from asyncio import Lock, TimerHandle, Transport
-from typing import Optional, Type
+from typing import Optional
 
 import async_timeout
 
 from . import DeviceData
-from .api import COMMANDS_SECURED, QUEUE_TIMEOUT, RT, RoborockClient
+from .api import COMMANDS_SECURED, QUEUE_TIMEOUT, RoborockClient
 from .exceptions import CommandVacuumError, RoborockConnectionException, RoborockException
 from .protocol import MessageParser
 from .roborock_message import RoborockMessage, RoborockMessageProtocol
@@ -123,17 +123,13 @@ async def ping(self):
             )
         )
 
-    async def send_command(
+    async def _send_command(
         self,
         method: RoborockCommand,
         params: Optional[list | dict] = None,
-        return_type: Optional[Type[RT]] = None,
     ):
         roborock_message = self.build_roborock_message(method, params)
-        response = (await self.send_message(roborock_message))[0]
-        if return_type:
-            return return_type.from_dict(response)
-        return response
+        return (await self.send_message(roborock_message))[0]
 
     async def async_local_response(self, roborock_message: RoborockMessage):
         method = roborock_message.get_method()

roborock/roborock_typing.py

@@ -127,16 +127,6 @@ class RoborockCommand(str, Enum):
     UPD_SERVER_TIMER = "upd_server_timer"
 
 
-CacheableCommands = [
-    attribute
-    for attribute in [
-        command.removeprefix("get_").removeprefix("set_").removeprefix("change_").upper() for command in RoborockCommand
-    ]
-    if "GET_" + attribute in RoborockCommand.__members__
-    and ("SET_" + attribute in RoborockCommand.__members__ or "CHANGE_" + attribute in RoborockCommand.__members__)
-]
-
-
 @dataclass
 class CommandInfo:
     params: Optional[list | dict] = None