Skip to content

Commit 8bf70f4

Browse files
authored
fix: refactor cache and call get_status after changing mop mode (#105)
1 parent 2da7b38 commit 8bf70f4

File tree

2 files changed

+175
-139
lines changed

2 files changed

+175
-139
lines changed

roborock/api.py

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
import aiohttp
1919

2020
from .code_mappings import RoborockDockTypeCode
21-
from .command_cache import CacheableAttribute, CommandType, RoborockAttribute, create_cache_map, parse_method
21+
from .command_cache import CacheableAttribute, CommandType, RoborockAttribute, find_cacheable_attribute, get_cache_map
2222
from .containers import (
2323
ChildLockStatus,
2424
CleanRecord,
@@ -152,6 +152,9 @@ async def close_value(self, params=None):
152152
await self._async_value()
153153
return response
154154

155+
async def refresh_value(self):
156+
await self._async_value()
157+
155158

156159
device_cache: dict[str, dict[CacheableAttribute, AttributeCache]] = {}
157160

@@ -171,8 +174,7 @@ def __init__(self, endpoint: str, device_info: DeviceData) -> None:
171174
cache = device_cache.get(device_info.device.duid)
172175
if not cache:
173176
cache = {
174-
cacheable_attribute: AttributeCache(attr, self)
175-
for cacheable_attribute, attr in create_cache_map().items()
177+
cacheable_attribute: AttributeCache(attr, self) for cacheable_attribute, attr in get_cache_map().items()
176178
}
177179
device_cache[device_info.device.duid] = cache
178180
self.cache: dict[CacheableAttribute, AttributeCache] = cache
@@ -260,7 +262,8 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
260262
elif data_protocol in ROBOROCK_DATA_CONSUMABLE_PROTOCOL:
261263
if self.cache[CacheableAttribute.consumable].value is None:
262264
self._logger.debug(
263-
f"Got consumable update({data_protocol.name}) before get_status was called."
265+
f"Got consumable update({data_protocol.name})"
266+
+ "before get_consumable was called."
264267
)
265268
self.cache[CacheableAttribute.consumable]._value = {}
266269
value = self.cache[CacheableAttribute.consumable].value
@@ -375,18 +378,22 @@ async def send_command(
375378
params: Optional[list | dict] = None,
376379
return_type: Optional[Type[RT]] = None,
377380
) -> RT:
378-
parsed_method = parse_method(method)
381+
cacheable_attribute_result = find_cacheable_attribute(method)
382+
379383
cache = None
380-
if parsed_method is not None:
381-
cache = self.cache[parsed_method.attribute]
384+
command_type = None
385+
if cacheable_attribute_result is not None:
386+
cache = self.cache[cacheable_attribute_result.attribute]
387+
command_type = cacheable_attribute_result.type
388+
382389
response: Any = None
383-
if cache is not None:
384-
if parsed_method.type == CommandType.GET:
385-
response = await cache.async_value()
386-
elif parsed_method.type == CommandType.SET:
387-
response = await cache.update_value(params)
390+
if cache is not None and command_type == CommandType.GET:
391+
response = await cache.async_value()
388392
else:
389393
response = await self._send_command(method, params)
394+
if cache is not None and command_type == CommandType.CHANGE:
395+
await cache.refresh_value()
396+
390397
if return_type:
391398
return return_type.from_dict(response)
392399
return response

roborock/command_cache.py

Lines changed: 156 additions & 127 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from __future__ import annotations
22

3-
from dataclasses import dataclass
3+
from dataclasses import dataclass, field
44
from enum import Enum
55
from typing import Mapping, Optional
66

@@ -41,143 +41,172 @@ class RoborockAttribute:
4141
add_command: Optional[RoborockCommand] = None
4242
set_command: Optional[RoborockCommand] = None
4343
close_command: Optional[RoborockCommand] = None
44+
additional_change_commands: list[RoborockCommand] = field(default_factory=list)
4445

4546

46-
def create_cache_map():
47-
cache_map: Mapping[CacheableAttribute, RoborockAttribute] = {
48-
CacheableAttribute.status: RoborockAttribute(
49-
attribute="status",
50-
get_command=RoborockCommand.GET_STATUS,
51-
),
52-
CacheableAttribute.consumable: RoborockAttribute(
53-
attribute="consumable",
54-
get_command=RoborockCommand.GET_CONSUMABLE,
55-
),
56-
CacheableAttribute.sound_volume: RoborockAttribute(
57-
attribute="sound_volume",
58-
get_command=RoborockCommand.GET_SOUND_VOLUME,
59-
set_command=RoborockCommand.CHANGE_SOUND_VOLUME,
60-
),
61-
CacheableAttribute.camera_status: RoborockAttribute(
62-
attribute="camera_status",
63-
get_command=RoborockCommand.GET_CAMERA_STATUS,
64-
set_command=RoborockCommand.SET_CAMERA_STATUS,
65-
),
66-
CacheableAttribute.carpet_clean_mode: RoborockAttribute(
67-
attribute="carpet_clean_mode",
68-
get_command=RoborockCommand.GET_CARPET_CLEAN_MODE,
69-
set_command=RoborockCommand.SET_CARPET_CLEAN_MODE,
70-
),
71-
CacheableAttribute.carpet_mode: RoborockAttribute(
72-
attribute="carpet_mode",
73-
get_command=RoborockCommand.GET_CARPET_MODE,
74-
set_command=RoborockCommand.SET_CARPET_MODE,
75-
),
76-
CacheableAttribute.child_lock_status: RoborockAttribute(
77-
attribute="child_lock_status",
78-
get_command=RoborockCommand.GET_CHILD_LOCK_STATUS,
79-
set_command=RoborockCommand.SET_CHILD_LOCK_STATUS,
80-
),
81-
CacheableAttribute.collision_avoid_status: RoborockAttribute(
82-
attribute="collision_avoid_status",
83-
get_command=RoborockCommand.GET_COLLISION_AVOID_STATUS,
84-
set_command=RoborockCommand.SET_COLLISION_AVOID_STATUS,
85-
),
86-
CacheableAttribute.customize_clean_mode: RoborockAttribute(
87-
attribute="customize_clean_mode",
88-
get_command=RoborockCommand.GET_CUSTOMIZE_CLEAN_MODE,
89-
set_command=RoborockCommand.SET_CUSTOMIZE_CLEAN_MODE,
90-
),
91-
CacheableAttribute.custom_mode: RoborockAttribute(
92-
attribute="custom_mode",
93-
get_command=RoborockCommand.GET_CUSTOM_MODE,
94-
set_command=RoborockCommand.SET_CUSTOM_MODE,
95-
),
96-
CacheableAttribute.dnd_timer: RoborockAttribute(
97-
attribute="dnd_timer",
98-
get_command=RoborockCommand.GET_DND_TIMER,
99-
set_command=RoborockCommand.SET_DND_TIMER,
100-
close_command=RoborockCommand.CLOSE_DND_TIMER,
101-
),
102-
CacheableAttribute.dust_collection_mode: RoborockAttribute(
103-
attribute="dust_collection_mode",
104-
get_command=RoborockCommand.GET_DUST_COLLECTION_MODE,
105-
set_command=RoborockCommand.SET_DUST_COLLECTION_MODE,
106-
),
107-
CacheableAttribute.flow_led_status: RoborockAttribute(
108-
attribute="flow_led_status",
109-
get_command=RoborockCommand.GET_FLOW_LED_STATUS,
110-
set_command=RoborockCommand.SET_FLOW_LED_STATUS,
111-
),
112-
CacheableAttribute.identify_furniture_status: RoborockAttribute(
113-
attribute="identify_furniture_status",
114-
get_command=RoborockCommand.GET_IDENTIFY_FURNITURE_STATUS,
115-
set_command=RoborockCommand.SET_IDENTIFY_FURNITURE_STATUS,
116-
),
117-
CacheableAttribute.identify_ground_material_status: RoborockAttribute(
118-
attribute="identify_ground_material_status",
119-
get_command=RoborockCommand.GET_IDENTIFY_GROUND_MATERIAL_STATUS,
120-
set_command=RoborockCommand.SET_IDENTIFY_GROUND_MATERIAL_STATUS,
121-
),
122-
CacheableAttribute.led_status: RoborockAttribute(
123-
attribute="led_status",
124-
get_command=RoborockCommand.GET_LED_STATUS,
125-
set_command=RoborockCommand.SET_LED_STATUS,
126-
),
127-
CacheableAttribute.server_timer: RoborockAttribute(
128-
attribute="server_timer",
129-
get_command=RoborockCommand.GET_SERVER_TIMER,
130-
add_command=RoborockCommand.SET_SERVER_TIMER,
131-
set_command=RoborockCommand.UPD_SERVER_TIMER,
132-
close_command=RoborockCommand.DEL_SERVER_TIMER,
133-
),
134-
CacheableAttribute.smart_wash_params: RoborockAttribute(
135-
attribute="smart_wash_params",
136-
get_command=RoborockCommand.GET_SMART_WASH_PARAMS,
137-
set_command=RoborockCommand.SET_SMART_WASH_PARAMS,
138-
),
139-
CacheableAttribute.timezone: RoborockAttribute(
140-
attribute="timezone", get_command=RoborockCommand.GET_TIMEZONE, set_command=RoborockCommand.SET_TIMEZONE
141-
),
142-
CacheableAttribute.valley_electricity_timer: RoborockAttribute(
143-
attribute="valley_electricity_timer",
144-
get_command=RoborockCommand.GET_VALLEY_ELECTRICITY_TIMER,
145-
set_command=RoborockCommand.SET_VALLEY_ELECTRICITY_TIMER,
146-
close_command=RoborockCommand.CLOSE_VALLEY_ELECTRICITY_TIMER,
147-
),
148-
CacheableAttribute.wash_towel_mode: RoborockAttribute(
149-
attribute="wash_towel_mode",
150-
get_command=RoborockCommand.GET_WASH_TOWEL_MODE,
151-
set_command=RoborockCommand.SET_WASH_TOWEL_MODE,
152-
),
153-
}
47+
cache_map: Mapping[CacheableAttribute, RoborockAttribute] = {
48+
CacheableAttribute.status: RoborockAttribute(
49+
attribute="status",
50+
get_command=RoborockCommand.GET_STATUS,
51+
additional_change_commands=[
52+
RoborockCommand.SET_WATER_BOX_CUSTOM_MODE,
53+
RoborockCommand.SET_MOP_MODE,
54+
],
55+
),
56+
CacheableAttribute.consumable: RoborockAttribute(
57+
attribute="consumable",
58+
get_command=RoborockCommand.GET_CONSUMABLE,
59+
),
60+
CacheableAttribute.sound_volume: RoborockAttribute(
61+
attribute="sound_volume",
62+
get_command=RoborockCommand.GET_SOUND_VOLUME,
63+
set_command=RoborockCommand.CHANGE_SOUND_VOLUME,
64+
),
65+
CacheableAttribute.camera_status: RoborockAttribute(
66+
attribute="camera_status",
67+
get_command=RoborockCommand.GET_CAMERA_STATUS,
68+
set_command=RoborockCommand.SET_CAMERA_STATUS,
69+
),
70+
CacheableAttribute.carpet_clean_mode: RoborockAttribute(
71+
attribute="carpet_clean_mode",
72+
get_command=RoborockCommand.GET_CARPET_CLEAN_MODE,
73+
set_command=RoborockCommand.SET_CARPET_CLEAN_MODE,
74+
),
75+
CacheableAttribute.carpet_mode: RoborockAttribute(
76+
attribute="carpet_mode",
77+
get_command=RoborockCommand.GET_CARPET_MODE,
78+
set_command=RoborockCommand.SET_CARPET_MODE,
79+
),
80+
CacheableAttribute.child_lock_status: RoborockAttribute(
81+
attribute="child_lock_status",
82+
get_command=RoborockCommand.GET_CHILD_LOCK_STATUS,
83+
set_command=RoborockCommand.SET_CHILD_LOCK_STATUS,
84+
),
85+
CacheableAttribute.collision_avoid_status: RoborockAttribute(
86+
attribute="collision_avoid_status",
87+
get_command=RoborockCommand.GET_COLLISION_AVOID_STATUS,
88+
set_command=RoborockCommand.SET_COLLISION_AVOID_STATUS,
89+
),
90+
CacheableAttribute.customize_clean_mode: RoborockAttribute(
91+
attribute="customize_clean_mode",
92+
get_command=RoborockCommand.GET_CUSTOMIZE_CLEAN_MODE,
93+
set_command=RoborockCommand.SET_CUSTOMIZE_CLEAN_MODE,
94+
),
95+
CacheableAttribute.custom_mode: RoborockAttribute(
96+
attribute="custom_mode",
97+
get_command=RoborockCommand.GET_CUSTOM_MODE,
98+
set_command=RoborockCommand.SET_CUSTOM_MODE,
99+
),
100+
CacheableAttribute.dnd_timer: RoborockAttribute(
101+
attribute="dnd_timer",
102+
get_command=RoborockCommand.GET_DND_TIMER,
103+
set_command=RoborockCommand.SET_DND_TIMER,
104+
close_command=RoborockCommand.CLOSE_DND_TIMER,
105+
),
106+
CacheableAttribute.dust_collection_mode: RoborockAttribute(
107+
attribute="dust_collection_mode",
108+
get_command=RoborockCommand.GET_DUST_COLLECTION_MODE,
109+
set_command=RoborockCommand.SET_DUST_COLLECTION_MODE,
110+
),
111+
CacheableAttribute.flow_led_status: RoborockAttribute(
112+
attribute="flow_led_status",
113+
get_command=RoborockCommand.GET_FLOW_LED_STATUS,
114+
set_command=RoborockCommand.SET_FLOW_LED_STATUS,
115+
),
116+
CacheableAttribute.identify_furniture_status: RoborockAttribute(
117+
attribute="identify_furniture_status",
118+
get_command=RoborockCommand.GET_IDENTIFY_FURNITURE_STATUS,
119+
set_command=RoborockCommand.SET_IDENTIFY_FURNITURE_STATUS,
120+
),
121+
CacheableAttribute.identify_ground_material_status: RoborockAttribute(
122+
attribute="identify_ground_material_status",
123+
get_command=RoborockCommand.GET_IDENTIFY_GROUND_MATERIAL_STATUS,
124+
set_command=RoborockCommand.SET_IDENTIFY_GROUND_MATERIAL_STATUS,
125+
),
126+
CacheableAttribute.led_status: RoborockAttribute(
127+
attribute="led_status",
128+
get_command=RoborockCommand.GET_LED_STATUS,
129+
set_command=RoborockCommand.SET_LED_STATUS,
130+
),
131+
CacheableAttribute.server_timer: RoborockAttribute(
132+
attribute="server_timer",
133+
get_command=RoborockCommand.GET_SERVER_TIMER,
134+
add_command=RoborockCommand.SET_SERVER_TIMER,
135+
set_command=RoborockCommand.UPD_SERVER_TIMER,
136+
close_command=RoborockCommand.DEL_SERVER_TIMER,
137+
),
138+
CacheableAttribute.smart_wash_params: RoborockAttribute(
139+
attribute="smart_wash_params",
140+
get_command=RoborockCommand.GET_SMART_WASH_PARAMS,
141+
set_command=RoborockCommand.SET_SMART_WASH_PARAMS,
142+
),
143+
CacheableAttribute.timezone: RoborockAttribute(
144+
attribute="timezone", get_command=RoborockCommand.GET_TIMEZONE, set_command=RoborockCommand.SET_TIMEZONE
145+
),
146+
CacheableAttribute.valley_electricity_timer: RoborockAttribute(
147+
attribute="valley_electricity_timer",
148+
get_command=RoborockCommand.GET_VALLEY_ELECTRICITY_TIMER,
149+
set_command=RoborockCommand.SET_VALLEY_ELECTRICITY_TIMER,
150+
close_command=RoborockCommand.CLOSE_VALLEY_ELECTRICITY_TIMER,
151+
),
152+
CacheableAttribute.wash_towel_mode: RoborockAttribute(
153+
attribute="wash_towel_mode",
154+
get_command=RoborockCommand.GET_WASH_TOWEL_MODE,
155+
set_command=RoborockCommand.SET_WASH_TOWEL_MODE,
156+
),
157+
}
158+
159+
160+
def get_change_commands(attr: RoborockAttribute) -> list[RoborockCommand]:
161+
commands = [
162+
attr.add_command,
163+
attr.set_command,
164+
attr.close_command,
165+
*attr.additional_change_commands,
166+
]
167+
168+
return [command for command in commands if command is not None]
169+
170+
171+
cache_map_by_get_command: dict[RoborockCommand, CacheableAttribute] = {
172+
attribute.get_command: cacheable_attribute for cacheable_attribute, attribute in cache_map.items()
173+
}
174+
175+
cache_map_by_change_command: dict[RoborockCommand, CacheableAttribute] = {
176+
command: cacheable_attribute
177+
for cacheable_attribute, attribute in cache_map.items()
178+
for command in get_change_commands(attribute)
179+
}
180+
181+
182+
def get_cache_map():
154183
return cache_map
155184

156185

157186
class CommandType(Enum):
158187
OTHER = -1
159188
GET = 0
160-
SET = 1
189+
CHANGE = 1
161190

162191

163192
@dataclass
164-
class ParserCommand:
165-
type: CommandType
193+
class CacheableAttributeResult:
166194
attribute: CacheableAttribute
195+
type: CommandType
196+
197+
198+
def find_cacheable_attribute(method: RoborockCommand) -> CacheableAttributeResult | None:
199+
if method is None:
200+
return None
167201

202+
cacheable_attribute = None
203+
command_type = CommandType.OTHER
204+
if cacheable_attribute := cache_map_by_get_command.get(method, None):
205+
command_type = CommandType.GET
206+
elif cacheable_attribute := cache_map_by_change_command.get(method, None):
207+
command_type = CommandType.CHANGE
168208

169-
def parse_method(method: str):
170-
if method is not None:
171-
attribute = method.lower()
172-
command_type = CommandType.OTHER
173-
if attribute.startswith(GET_PREFIX):
174-
attribute = attribute.removeprefix(GET_PREFIX)
175-
command_type = CommandType.GET
176-
elif attribute.startswith(SET_PREFIX):
177-
for prefix in SET_PREFIX:
178-
attribute = attribute.removeprefix(prefix)
179-
command_type = CommandType.SET
180-
cacheable_attribute = next((attr for attr in CacheableAttribute if attr == attribute), None)
181-
if cacheable_attribute:
182-
return ParserCommand(type=command_type, attribute=CacheableAttribute(cacheable_attribute))
183-
return None
209+
if cacheable_attribute:
210+
return CacheableAttributeResult(attribute=CacheableAttribute(cacheable_attribute), type=command_type)
211+
else:
212+
return None

0 commit comments

Comments
 (0)