Skip to content

Commit 2462505

Browse files
committed
feat: Add initial Q10 support for Status Trait
For now there is a central refresh mechanism for all traits which can be revisited in the future. The properties API listens for updates and fans them out to traits, which listen for apporopriate DPS values they are responsible for. This adds a way to map dps values to dataclass fields, and reuses the exisitng Roborock base data for parsing updates.
1 parent a69286f commit 2462505

File tree

9 files changed

+412
-11
lines changed

9 files changed

+412
-11
lines changed
Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,53 +1,102 @@
1-
from ..containers import RoborockBase
1+
"""Data container classes for Q10 B01 devices.
2+
3+
Many of these classes use the `field(metadata={"dps": ...})` convention to map
4+
dataclass fields to device Data Points (DPS). This metadata is utilized by the
5+
`update_from_dps` helper in `roborock.devices.traits.b01.q10.common` to
6+
automatically update objects from raw device responses.
7+
"""
28

9+
from dataclasses import dataclass, field
310

11+
from ..containers import RoborockBase
12+
from .b01_q10_code_mappings import (
13+
B01_Q10_DP,
14+
YXBackType,
15+
YXDeviceCleanTask,
16+
YXDeviceState,
17+
YXDeviceWorkMode,
18+
YXFanLevel,
19+
YXWaterLevel,
20+
)
21+
22+
23+
@dataclass
424
class dpCleanRecord(RoborockBase):
525
op: str
626
result: int
727
id: str
828
data: list
929

1030

31+
@dataclass
1132
class dpMultiMap(RoborockBase):
1233
op: str
1334
result: int
1435
data: list
1536

1637

38+
@dataclass
1739
class dpGetCarpet(RoborockBase):
1840
op: str
1941
result: int
2042
data: str
2143

2244

45+
@dataclass
2346
class dpSelfIdentifyingCarpet(RoborockBase):
2447
op: str
2548
result: int
2649
data: str
2750

2851

52+
@dataclass
2953
class dpNetInfo(RoborockBase):
3054
wifiName: str
3155
ipAdress: str
3256
mac: str
3357
signal: int
3458

3559

60+
@dataclass
3661
class dpNotDisturbExpand(RoborockBase):
3762
disturb_dust_enable: int
3863
disturb_light: int
3964
disturb_resume_clean: int
4065
disturb_voice: int
4166

4267

68+
@dataclass
4369
class dpCurrentCleanRoomIds(RoborockBase):
4470
room_id_list: list
4571

4672

73+
@dataclass
4774
class dpVoiceVersion(RoborockBase):
4875
version: int
4976

5077

78+
@dataclass
5179
class dpTimeZone(RoborockBase):
5280
timeZoneCity: str
5381
timeZoneSec: int
82+
83+
84+
@dataclass
85+
class Q10Status(RoborockBase):
86+
"""Status for Q10 devices.
87+
88+
Fields are mapped to DPS values using metadata. Objects of this class can be
89+
automatically updated using the `update_from_dps` helper.
90+
"""
91+
92+
clean_time: int | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_TIME})
93+
clean_area: int | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_AREA})
94+
battery: int | None = field(default=None, metadata={"dps": B01_Q10_DP.BATTERY})
95+
status: YXDeviceState | None = field(default=None, metadata={"dps": B01_Q10_DP.STATUS})
96+
fun_level: YXFanLevel | None = field(default=None, metadata={"dps": B01_Q10_DP.FUN_LEVEL})
97+
water_level: YXWaterLevel | None = field(default=None, metadata={"dps": B01_Q10_DP.WATER_LEVEL})
98+
clean_count: int | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_COUNT})
99+
clean_mode: YXDeviceWorkMode | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_MODE})
100+
clean_task_type: YXDeviceCleanTask | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_TASK_TYPE})
101+
back_type: YXBackType | None = field(default=None, metadata={"dps": B01_Q10_DP.BACK_TYPE})
102+
cleaning_progress: int | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEANING_PROGRESS})

roborock/data/containers.py

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -91,10 +91,10 @@ def from_dict(cls, data: dict[str, Any]):
9191
if not isinstance(data, dict):
9292
return None
9393
field_types = {field.name: field.type for field in dataclasses.fields(cls)}
94-
result: dict[str, Any] = {}
94+
normalized_data: dict[str, Any] = {}
9595
for orig_key, value in data.items():
9696
key = _decamelize(orig_key)
97-
if (field_type := field_types.get(key)) is None:
97+
if field_types.get(key) is None:
9898
if (log_key := f"{cls.__name__}.{key}") not in RoborockBase._missing_logged:
9999
_LOGGER.debug(
100100
"Key '%s' (decamelized: '%s') not found in %s fields, skipping",
@@ -104,6 +104,23 @@ def from_dict(cls, data: dict[str, Any]):
104104
)
105105
RoborockBase._missing_logged.add(log_key)
106106
continue
107+
normalized_data[key] = value
108+
109+
result = RoborockBase.convert_dict(field_types, normalized_data)
110+
return cls(**result)
111+
112+
@staticmethod
113+
def convert_dict(types_map: dict[Any, type], data: dict[Any, Any]) -> dict[Any, Any]:
114+
"""Generic helper to convert a dictionary of values based on a schema map of types.
115+
116+
This is meant to be used by traits that use dataclass reflection similar to
117+
`Roborock.from_dict` to merge in new data updates.
118+
"""
119+
result: dict[Any, Any] = {}
120+
for key, value in data.items():
121+
if key not in types_map:
122+
continue
123+
field_type = types_map[key]
107124
if value == "None" or value is None:
108125
result[key] = None
109126
continue
@@ -124,7 +141,7 @@ def from_dict(cls, data: dict[str, Any]):
124141
_LOGGER.exception(f"Failed to convert {key} with value {value} to type {field_type}")
125142
continue
126143

127-
return cls(**result)
144+
return result
128145

129146
def as_dict(self) -> dict:
130147
return asdict(

roborock/devices/device.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -197,12 +197,14 @@ async def connect(self) -> None:
197197
if self._unsub:
198198
raise ValueError("Already connected to the device")
199199
unsub = await self._channel.subscribe(self._on_message)
200-
if self.v1_properties is not None:
201-
try:
200+
try:
201+
if self.v1_properties is not None:
202202
await self.v1_properties.discover_features()
203-
except RoborockException:
204-
unsub()
205-
raise
203+
elif self.b01_q10_properties is not None:
204+
await self.b01_q10_properties.start()
205+
except RoborockException:
206+
unsub()
207+
raise
206208
self._logger.info("Connected to device")
207209
self._unsub = unsub
208210

@@ -214,6 +216,8 @@ async def close(self) -> None:
214216
await self._connect_task
215217
except asyncio.CancelledError:
216218
pass
219+
if self.b01_q10_properties is not None:
220+
await self.b01_q10_properties.close()
217221
if self._unsub:
218222
self._unsub()
219223
self._unsub = None

roborock/devices/rpc/b01_q10_channel.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,39 @@
33
from __future__ import annotations
44

55
import logging
6+
from collections.abc import AsyncGenerator
7+
from typing import Any
68

79
from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
810
from roborock.devices.transport.mqtt_channel import MqttChannel
911
from roborock.exceptions import RoborockException
1012
from roborock.protocols.b01_q10_protocol import (
1113
ParamsType,
14+
decode_rpc_response,
1215
encode_mqtt_payload,
1316
)
1417

1518
_LOGGER = logging.getLogger(__name__)
1619

1720

21+
async def stream_decoded_responses(
22+
mqtt_channel: MqttChannel,
23+
) -> AsyncGenerator[dict[B01_Q10_DP, Any], None]:
24+
"""Stream decoded DPS messages received via MQTT."""
25+
26+
async for response_message in mqtt_channel.subscribe_stream():
27+
try:
28+
decoded_dps = decode_rpc_response(response_message)
29+
except RoborockException as ex:
30+
_LOGGER.debug(
31+
"Failed to decode B01 RPC response: %s: %s",
32+
response_message,
33+
ex,
34+
)
35+
continue
36+
yield decoded_dps
37+
38+
1839
async def send_command(
1940
mqtt_channel: MqttChannel,
2041
command: B01_Q10_DP,

roborock/devices/traits/b01/q10/__init__.py

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,72 @@
11
"""Traits for Q10 B01 devices."""
22

3+
import asyncio
4+
import logging
35
from typing import Any
46

5-
from roborock.devices.rpc.b01_q7_channel import send_decoded_command
7+
from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
8+
from roborock.devices.rpc.b01_q10_channel import stream_decoded_responses
69
from roborock.devices.traits import Trait
710
from roborock.devices.transport.mqtt_channel import MqttChannel
811

912
from .command import CommandTrait
13+
from .status import StatusTrait
1014

1115
__all__ = [
1216
"Q10PropertiesApi",
1317
]
1418

19+
_LOGGER = logging.getLogger(__name__)
20+
1521

1622
class Q10PropertiesApi(Trait):
1723
"""API for interacting with B01 devices."""
1824

1925
command: CommandTrait
2026
"""Trait for sending commands to Q10 devices."""
2127

28+
status: StatusTrait
29+
"""Trait for managing the status of Q10 devices."""
30+
2231
def __init__(self, channel: MqttChannel) -> None:
2332
"""Initialize the B01Props API."""
33+
self._channel = channel
2434
self.command = CommandTrait(channel)
35+
self.status = StatusTrait()
36+
self._subscribe_task: asyncio.Task[None] | None = None
37+
38+
async def start(self) -> None:
39+
"""Start any necessary subscriptions for the trait."""
40+
self._subscribe_task = asyncio.create_task(self._subscribe_loop())
41+
42+
async def close(self) -> None:
43+
"""Close any resources held by the trait."""
44+
if self._subscribe_task is not None:
45+
self._subscribe_task.cancel()
46+
try:
47+
await self._subscribe_task
48+
except asyncio.CancelledError:
49+
pass # ignore cancellation errors
50+
self._subscribe_task = None
51+
52+
async def refresh(self) -> None:
53+
"""Refresh all traits."""
54+
# Ask for updates to speific DPS values. Updates will be received
55+
# by the subscribe loop below.
56+
# For now we just ask for all DPS values that traits care about here
57+
# but this could be split out to give each trait its own refresh
58+
# method in the future if needed.
59+
await self.command.send(B01_Q10_DP.REQUETDPS, params={})
60+
61+
async def _subscribe_loop(self) -> None:
62+
"""Persistent loop to listen for status updates."""
63+
async for decoded_dps in stream_decoded_responses(self._channel):
64+
_LOGGER.debug("Received Q10 status update: %s", decoded_dps)
65+
66+
# Notify all traits about a new message and each trait will
67+
# only update what fields that it is responsible for.
68+
# More traits can be added here below.
69+
self.status.update_from_dps(decoded_dps)
2570

2671

2772
def create(channel: MqttChannel) -> Q10PropertiesApi:
Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
"""Common utilities for Q10 traits.
2+
3+
This module provides infrastructure for mapping Roborock Data Points (DPS) to
4+
Python dataclass fields and handling the lifecycle of data updates from the
5+
device.
6+
7+
### DPS Metadata Annotation
8+
9+
Classes extending `RoborockBase` can annotate their fields with DPS IDs using
10+
the `field(metadata={"dps": ...})` convention. This creates a declarative
11+
mapping that `DpsDataConverter` uses to automatically route incoming device
12+
data to the correct attribute.
13+
14+
Example:
15+
16+
```python
17+
@dataclass
18+
class MyStatus(RoborockBase):
19+
battery: int = field(metadata={"dps": B01_Q10_DP.BATTERY})
20+
```
21+
22+
### Update Lifecycle
23+
1. **Raw Data**: The device sends encoded DPS updates over MQTT.
24+
2. **Decoding**: The transport layer decodes these into a dictionary (e.g., `{"101": 80}`).
25+
3. **Conversion**: `DpsDataConverter` uses `RoborockBase.convert_dict` to transform
26+
raw values into appropriate Python types (e.g., Enums, ints) based on the
27+
dataclass field types.
28+
4. **Update**: `update_from_dps` maps these converted values to field names and
29+
updates the target object using `setattr`.
30+
31+
### Usage
32+
33+
Typically, a trait will instantiate a single `DpsDataConverter` for its status class
34+
and call `update_from_dps` whenever new data is received from the device stream.
35+
36+
"""
37+
38+
import dataclasses
39+
from typing import Any
40+
41+
from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
42+
from roborock.data.containers import RoborockBase
43+
44+
45+
class DpsDataConverter:
46+
"""Utility to handle the transformation and merging of DPS data into models.
47+
48+
This class pre-calculates the mapping between Data Point IDs and dataclass fields
49+
to optimize repeated updates from device streams.
50+
"""
51+
52+
def __init__(self, dps_type_map: dict[B01_Q10_DP, type], dps_field_map: dict[B01_Q10_DP, str]):
53+
"""Initialize the converter for a specific RoborockBase-derived class."""
54+
self._dps_type_map = dps_type_map
55+
self._dps_field_map = dps_field_map
56+
57+
@classmethod
58+
def from_dataclass(cls, borockBase: type[RoborockBase]):
59+
"""Initialize the converter for a specific RoborockBase-derived class."""
60+
dps_type_map: dict[B01_Q10_DP, type] = {}
61+
dps_field_map: dict[B01_Q10_DP, str] = {}
62+
for field_obj in dataclasses.fields(borockBase):
63+
if field_obj.metadata and "dps" in field_obj.metadata:
64+
dps_id = field_obj.metadata["dps"]
65+
dps_type_map[dps_id] = field_obj.type
66+
dps_field_map[dps_id] = field_obj.name
67+
return cls(dps_type_map, dps_field_map)
68+
69+
def update_from_dps(self, target: RoborockBase, decoded_dps: dict[B01_Q10_DP, Any]) -> None:
70+
"""Convert and merge raw DPS data into the target object.
71+
72+
Uses the pre-calculated type mapping to ensure values are converted to the
73+
correct Python types before being updated on the target.
74+
75+
Args:
76+
target: The target object to update.
77+
decoded_dps: The decoded DPS data to convert.
78+
"""
79+
conversions = RoborockBase.convert_dict(self._dps_type_map, decoded_dps)
80+
for dps_id, value in conversions.items():
81+
field_name = self._dps_field_map[dps_id]
82+
setattr(target, field_name, value)

0 commit comments

Comments
 (0)