feat(q10): enhance Q10 device traits with status management and data conversion utilities

Author: lboue

roborock/data/b01_q10/b01_q10_containers.py

@@ -1,53 +1,102 @@
-from ..containers import RoborockBase
+"""Data container classes for Q10 B01 devices.
+
+Many of these classes use the `field(metadata={"dps": ...})` convention to map
+fields to device Data Points (DPS). This metadata is utilized by the
+`update_from_dps` helper in `roborock.devices.traits.b01.q10.common` to
+automatically update objects from raw device responses.
+"""
 
+from dataclasses import dataclass, field
 
+from ..containers import RoborockBase
+from .b01_q10_code_mappings import (
+    B01_Q10_DP,
+    YXBackType,
+    YXDeviceCleanTask,
+    YXDeviceState,
+    YXDeviceWorkMode,
+    YXFanLevel,
+    YXWaterLevel,
+)
+
+
+@dataclass
 class dpCleanRecord(RoborockBase):
     op: str
     result: int
     id: str
     data: list
 
 
+@dataclass
 class dpMultiMap(RoborockBase):
     op: str
     result: int
     data: list
 
 
+@dataclass
 class dpGetCarpet(RoborockBase):
     op: str
     result: int
     data: str
 
 
+@dataclass
 class dpSelfIdentifyingCarpet(RoborockBase):
     op: str
     result: int
     data: str
 
 
+@dataclass
 class dpNetInfo(RoborockBase):
     wifiName: str
     ipAdress: str
     mac: str
     signal: int
 
 
+@dataclass
 class dpNotDisturbExpand(RoborockBase):
     disturb_dust_enable: int
     disturb_light: int
     disturb_resume_clean: int
     disturb_voice: int
 
 
+@dataclass
 class dpCurrentCleanRoomIds(RoborockBase):
     room_id_list: list
 
 
+@dataclass
 class dpVoiceVersion(RoborockBase):
     version: int
 
 
+@dataclass
 class dpTimeZone(RoborockBase):
     timeZoneCity: str
     timeZoneSec: int
+
+
+@dataclass
+class Q10Status(RoborockBase):
+    """Status for Q10 devices.
+
+    Fields are mapped to DPS values using metadata. Objects of this class can be
+    automatically updated using the `update_from_dps` helper.
+    """
+
+    clean_time: int | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_TIME})
+    clean_area: int | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_AREA})
+    battery: int | None = field(default=None, metadata={"dps": B01_Q10_DP.BATTERY})
+    status: YXDeviceState | None = field(default=None, metadata={"dps": B01_Q10_DP.STATUS})
+    fan_level: YXFanLevel | None = field(default=None, metadata={"dps": B01_Q10_DP.FAN_LEVEL})
+    water_level: YXWaterLevel | None = field(default=None, metadata={"dps": B01_Q10_DP.WATER_LEVEL})
+    clean_count: int | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_COUNT})
+    clean_mode: YXDeviceWorkMode | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_MODE})
+    clean_task_type: YXDeviceCleanTask | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEAN_TASK_TYPE})
+    back_type: YXBackType | None = field(default=None, metadata={"dps": B01_Q10_DP.BACK_TYPE})
+    cleaning_progress: int | None = field(default=None, metadata={"dps": B01_Q10_DP.CLEANING_PROGRESS})

roborock/data/containers.py

@@ -91,10 +91,10 @@ def from_dict(cls, data: dict[str, Any]):
         if not isinstance(data, dict):
             return None
         field_types = {field.name: field.type for field in dataclasses.fields(cls)}
-        result: dict[str, Any] = {}
+        normalized_data: dict[str, Any] = {}
         for orig_key, value in data.items():
             key = _decamelize(orig_key)
-            if (field_type := field_types.get(key)) is None:
+            if field_types.get(key) is None:
                 if (log_key := f"{cls.__name__}.{key}") not in RoborockBase._missing_logged:
                     _LOGGER.debug(
                         "Key '%s' (decamelized: '%s') not found in %s fields, skipping",
@@ -104,6 +104,19 @@ def from_dict(cls, data: dict[str, Any]):
                     )
                     RoborockBase._missing_logged.add(log_key)
                 continue
+            normalized_data[key] = value
+
+        result = RoborockBase.convert_dict(field_types, normalized_data)
+        return cls(**result)
+
+    @staticmethod
+    def convert_dict(types_map: dict[Any, type], data: dict[Any, Any]) -> dict[Any, Any]:
+        """Convert a dictionary of values based on a schema map of types."""
+        result: dict[Any, Any] = {}
+        for key, value in data.items():
+            if key not in types_map:
+                continue
+            field_type = types_map[key]
             if value == "None" or value is None:
                 result[key] = None
                 continue
@@ -124,7 +137,7 @@ def from_dict(cls, data: dict[str, Any]):
                     _LOGGER.exception(f"Failed to convert {key} with value {value} to type {field_type}")
                     continue
 
-        return cls(**result)
+        return result
 
     def as_dict(self) -> dict:
         return asdict(

roborock/devices/device.py

@@ -197,12 +197,14 @@ async def connect(self) -> None:
         if self._unsub:
             raise ValueError("Already connected to the device")
         unsub = await self._channel.subscribe(self._on_message)
-        if self.v1_properties is not None:
-            try:
+        try:
+            if self.v1_properties is not None:
                 await self.v1_properties.discover_features()
-            except RoborockException:
-                unsub()
-                raise
+            elif self.b01_q10_properties is not None:
+                await self.b01_q10_properties.start()
+        except RoborockException:
+            unsub()
+            raise
         self._logger.info("Connected to device")
         self._unsub = unsub
 
@@ -214,6 +216,8 @@ async def close(self) -> None:
                 await self._connect_task
             except asyncio.CancelledError:
                 pass
+        if self.b01_q10_properties is not None:
+            await self.b01_q10_properties.close()
         if self._unsub:
             self._unsub()
             self._unsub = None

roborock/devices/rpc/b01_q10_channel.py

@@ -2,99 +2,53 @@
 
 from __future__ import annotations
 
-import asyncio
 import logging
-from collections.abc import Iterable
+from collections.abc import AsyncGenerator
 from typing import Any
 
 from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
 from roborock.devices.transport.mqtt_channel import MqttChannel
 from roborock.exceptions import RoborockException
-from roborock.protocols.b01_q10_protocol import (
-    ParamsType,
-    decode_rpc_response,
-    encode_mqtt_payload,
-)
-from roborock.roborock_message import RoborockMessage
+from roborock.protocols.b01_q10_protocol import ParamsType, encode_mqtt_payload
+from roborock.protocols.b01_q10_protocol import decode_rpc_response
 
 _LOGGER = logging.getLogger(__name__)
-_TIMEOUT = 10.0
 
 
-async def send_command(
+async def stream_decoded_responses(
     mqtt_channel: MqttChannel,
-    command: B01_Q10_DP,
-    params: ParamsType,
-) -> None:
-    """Send a command on the MQTT channel, without waiting for a response"""
-    _LOGGER.debug("Sending B01 MQTT command: cmd=%s params=%s", command, params)
-    roborock_message = encode_mqtt_payload(command, params)
-    _LOGGER.debug("Sending MQTT message: %s", roborock_message)
-    try:
-        await mqtt_channel.publish(roborock_message)
-    except RoborockException as ex:
-        _LOGGER.debug(
-            "Error sending B01 decoded command (method=%s params=%s): %s",
-            command,
-            params,
-            ex,
-        )
-        raise
-
+) -> AsyncGenerator[dict[B01_Q10_DP, Any], None]:
+    """Stream decoded DPS messages received via MQTT."""
 
-async def send_decoded_command(
-    mqtt_channel: MqttChannel,
-    command: B01_Q10_DP,
-    params: ParamsType,
-    expected_dps: Iterable[B01_Q10_DP] | None = None,
-) -> dict[B01_Q10_DP, Any]:
-    """Send a command and await the first decoded response.
-
-    Q10 responses are not correlated with a message id, so we filter on
-    expected datapoints when provided.
-    """
-    roborock_message = encode_mqtt_payload(command, params)
-    future: asyncio.Future[dict[B01_Q10_DP, Any]] = asyncio.get_running_loop().create_future()
-
-    expected_set = set(expected_dps) if expected_dps is not None else None
-
-    def find_response(response_message: RoborockMessage) -> None:
+    async for response_message in mqtt_channel.subscribe_stream():
         try:
             decoded_dps = decode_rpc_response(response_message)
         except RoborockException as ex:
             _LOGGER.debug(
-                "Failed to decode B01 Q10 RPC response (expecting %s): %s: %s",
-                command,
+                "Failed to decode B01 RPC response: %s: %s",
                 response_message,
                 ex,
             )
-            return
-        if expected_set and not any(dps in decoded_dps for dps in expected_set):
-            return
-        if not future.done():
-            future.set_result(decoded_dps)
+            continue
+        yield decoded_dps
 
-    unsub = await mqtt_channel.subscribe(find_response)
 
+async def send_command(
+    mqtt_channel: MqttChannel,
+    command: B01_Q10_DP,
+    params: ParamsType,
+) -> None:
+    """Send a command on the MQTT channel, without waiting for a response."""
+    _LOGGER.debug("Sending B01 MQTT command: cmd=%s params=%s", command, params)
+    roborock_message = encode_mqtt_payload(command, params)
     _LOGGER.debug("Sending MQTT message: %s", roborock_message)
     try:
         await mqtt_channel.publish(roborock_message)
-        return await asyncio.wait_for(future, timeout=_TIMEOUT)
-    except TimeoutError as ex:
-        raise RoborockException(f"B01 Q10 command timed out after {_TIMEOUT}s ({command})") from ex
     except RoborockException as ex:
-        _LOGGER.warning(
-            "Error sending B01 Q10 decoded command (%s): %s",
-            command,
-            ex,
-        )
-        raise
-    except Exception as ex:
-        _LOGGER.exception(
-            "Error sending B01 Q10 decoded command (%s): %s",
+        _LOGGER.debug(
+            "Error sending B01 decoded command (method=%s params=%s): %s",
             command,
+            params,
             ex,
         )
         raise
-    finally:
-        unsub()

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

@@ -1,5 +1,10 @@
 """Traits for Q10 B01 devices."""
 
+import asyncio
+import logging
+
+from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
+from roborock.devices.rpc.b01_q10_channel import stream_decoded_responses
 from roborock.devices.traits import Trait
 from roborock.devices.transport.mqtt_channel import MqttChannel
 
@@ -9,9 +14,10 @@
 
 __all__ = [
     "Q10PropertiesApi",
-    "StatusTrait",
 ]
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class Q10PropertiesApi(Trait):
     """API for interacting with B01 devices."""
@@ -20,16 +26,42 @@ class Q10PropertiesApi(Trait):
     """Trait for sending commands to Q10 devices."""
 
     status: StatusTrait
-    """Trait for querying Q10 device status."""
+    """Trait for managing the status of Q10 devices."""
 
     vacuum: VacuumTrait
     """Trait for sending vacuum related commands to Q10 devices."""
 
     def __init__(self, channel: MqttChannel) -> None:
         """Initialize the B01Props API."""
+        self._channel = channel
         self.command = CommandTrait(channel)
-        self.status = StatusTrait(channel)
         self.vacuum = VacuumTrait(self.command)
+        self.status = StatusTrait()
+        self._subscribe_task: asyncio.Task[None] | None = None
+
+    async def start(self) -> None:
+        """Start any necessary subscriptions for the trait."""
+        self._subscribe_task = asyncio.create_task(self._subscribe_loop())
+
+    async def close(self) -> None:
+        """Close any resources held by the trait."""
+        if self._subscribe_task is not None:
+            self._subscribe_task.cancel()
+            try:
+                await self._subscribe_task
+            except asyncio.CancelledError:
+                pass
+            self._subscribe_task = None
+
+    async def refresh(self) -> None:
+        """Refresh all traits."""
+        await self.command.send(B01_Q10_DP.REQUEST_DPS, params={})
+
+    async def _subscribe_loop(self) -> None:
+        """Persistent loop to listen for status updates."""
+        async for decoded_dps in stream_decoded_responses(self._channel):
+            _LOGGER.debug("Received Q10 status update: %s", decoded_dps)
+            self.status.update_from_dps(decoded_dps)
 
 
 def create(channel: MqttChannel) -> Q10PropertiesApi: