🦎 q7: add B01 map decode/parse/render and map content trait

Author: arduano

roborock/devices/device_manager.py

@@ -251,7 +251,12 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                     trait = b01.q10.create(channel)
                 elif "sc" in model_part:
                     # Q7 devices start with 'sc' in their model naming.
-                    trait = b01.q7.create(channel)
+                    trait = b01.q7.create(
+                        channel,
+                        local_key=device.local_key,
+                        serial=device.sn,
+                        model=product.model,
+                    )
                 else:
                     raise UnsupportedDeviceError(f"Device {device.name} has unsupported B01 model: {product.model}")
             case _:

roborock/devices/rpc/b01_q7_channel.py

@@ -14,7 +14,7 @@
     decode_rpc_response,
     encode_mqtt_payload,
 )
-from roborock.roborock_message import RoborockMessage
+from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
 
 _LOGGER = logging.getLogger(__name__)
 _TIMEOUT = 10.0
@@ -99,3 +99,52 @@ def find_response(response_message: RoborockMessage) -> None:
         raise
     finally:
         unsub()
+
+
+async def send_map_command(mqtt_channel: MqttChannel, request_message: Q7RequestMessage) -> bytes:
+    """Send map upload command and wait for MAP_RESPONSE payload bytes."""
+
+    roborock_message = encode_mqtt_payload(request_message)
+    future: asyncio.Future[bytes] = asyncio.get_running_loop().create_future()
+
+    def find_response(response_message: RoborockMessage) -> None:
+        if future.done():
+            return
+
+        if response_message.protocol == RoborockMessageProtocol.MAP_RESPONSE and response_message.payload:
+            future.set_result(response_message.payload)
+            return
+
+        try:
+            decoded_dps = decode_rpc_response(response_message)
+        except RoborockException:
+            return
+
+        for dps_value in decoded_dps.values():
+            if not isinstance(dps_value, str):
+                continue
+            try:
+                inner = json.loads(dps_value)
+            except (json.JSONDecodeError, TypeError):
+                continue
+            if not isinstance(inner, dict) or inner.get("msgId") != str(request_message.msg_id):
+                continue
+            code = inner.get("code", 0)
+            if code != 0:
+                future.set_exception(RoborockException(f"B01 command failed with code {code} ({request_message})"))
+                return
+            data = inner.get("data")
+            if isinstance(data, dict) and isinstance(data.get("payload"), str):
+                try:
+                    future.set_result(bytes.fromhex(data["payload"]))
+                except ValueError:
+                    pass
+
+    unsub = await mqtt_channel.subscribe(find_response)
+    try:
+        await mqtt_channel.publish(roborock_message)
+        return await asyncio.wait_for(future, timeout=_TIMEOUT)
+    except TimeoutError as ex:
+        raise RoborockException(f"B01 map command timed out after {_TIMEOUT}s ({request_message})") from ex
+    finally:
+        unsub()

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

@@ -21,10 +21,12 @@
 from roborock.roborock_typing import RoborockB01Q7Methods
 
 from .clean_summary import CleanSummaryTrait
+from .map_content import Q7MapContentTrait
 
 __all__ = [
     "Q7PropertiesApi",
     "CleanSummaryTrait",
+    "Q7MapContentTrait",
 ]
 
 
@@ -33,11 +35,13 @@ class Q7PropertiesApi(Trait):
 
     clean_summary: CleanSummaryTrait
     """Trait for clean records / clean summary (Q7 `service.get_record_list`)."""
+    map_content: Q7MapContentTrait
 
-    def __init__(self, channel: MqttChannel) -> None:
+    def __init__(self, channel: MqttChannel, *, local_key: str, serial: str, model: str) -> None:
         """Initialize the B01Props API."""
         self._channel = channel
         self.clean_summary = CleanSummaryTrait(channel)
+        self.map_content = Q7MapContentTrait(channel, local_key=local_key, serial=serial, model=model)
 
     async def query_values(self, props: list[RoborockB01Props]) -> B01Props | None:
         """Query the device for the values of the given Q7 properties."""
@@ -142,6 +146,6 @@ async def send(self, command: CommandType, params: ParamsType) -> Any:
         )
 
 
-def create(channel: MqttChannel) -> Q7PropertiesApi:
-    """Create traits for B01 devices."""
-    return Q7PropertiesApi(channel)
+def create(channel: MqttChannel, *, local_key: str, serial: str, model: str) -> Q7PropertiesApi:
+    """Create traits for B01 Q7 devices."""
+    return Q7PropertiesApi(channel, local_key=local_key, serial=serial, model=model)

roborock/devices/traits/b01/q7/map_content.py

@@ -0,0 +1,46 @@
+"""Map content trait for B01/Q7 devices."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from roborock.devices.rpc.b01_q7_channel import send_map_command
+from roborock.devices.transport.mqtt_channel import MqttChannel
+from roborock.devices.traits import Trait
+from roborock.devices.traits.v1.map_content import MapContent
+from roborock.map.b01_map_parser import decode_b01_map_payload, parse_scmap_payload, render_map_png
+from roborock.protocols.b01_q7_protocol import Q7RequestMessage
+from roborock.roborock_typing import RoborockB01Q7Methods
+
+
+@dataclass
+class B01MapContent(MapContent):
+    """B01 map content wrapper."""
+
+
+class Q7MapContentTrait(B01MapContent, Trait):
+    """Fetch and parse map content from B01/Q7 devices."""
+
+    def __init__(self, channel: MqttChannel, *, local_key: str, serial: str, model: str) -> None:
+        super().__init__()
+        self._channel = channel
+        self._local_key = local_key
+        self._serial = serial
+        self._model = model
+
+    async def refresh(self) -> B01MapContent:
+        raw_payload = await send_map_command(
+            self._channel,
+            Q7RequestMessage(dps=10000, command=RoborockB01Q7Methods.UPLOAD_BY_MAPTYPE, params={"maptype": 301}),
+        )
+        inflated = decode_b01_map_payload(
+            raw_payload,
+            local_key=self._local_key,
+            serial=self._serial,
+            model=self._model,
+        )
+        parsed = parse_scmap_payload(inflated)
+        self.raw_api_response = raw_payload
+        self.map_data = None
+        self.image_content = render_map_png(parsed)
+        return self

roborock/map/__init__.py

@@ -1,7 +1,13 @@
 """Module for Roborock map related data classes."""
 
+from .b01_map_parser import B01MapData, decode_b01_map_payload, parse_scmap_payload, render_map_png
 from .map_parser import MapParserConfig, ParsedMapData
 
 __all__ = [
+    "B01MapData",
     "MapParserConfig",
+    "ParsedMapData",
+    "decode_b01_map_payload",
+    "parse_scmap_payload",
+    "render_map_png",
 ]