🦎 roborock: address review items for q7 map + segment support

Author: arduano

pyproject.toml

@@ -29,6 +29,7 @@ dependencies = [
     "pyrate-limiter>=3.7.0,<4",
     "aiomqtt>=2.5.0,<3",
     "click-shell~=2.1",
+    "Pillow>=10,<12",
 ]
 
 [project.urls]

roborock/devices/rpc/b01_q7_channel.py

@@ -5,6 +5,7 @@
 import asyncio
 import json
 import logging
+import weakref
 from typing import Any
 
 from roborock.devices.transport.mqtt_channel import MqttChannel
@@ -18,6 +19,15 @@
 
 _LOGGER = logging.getLogger(__name__)
 _TIMEOUT = 10.0
+_map_command_locks: weakref.WeakKeyDictionary[MqttChannel, asyncio.Lock] = weakref.WeakKeyDictionary()
+
+
+def _get_map_command_lock(mqtt_channel: MqttChannel) -> asyncio.Lock:
+    lock = _map_command_locks.get(mqtt_channel)
+    if lock is None:
+        lock = asyncio.Lock()
+        _map_command_locks[mqtt_channel] = lock
+    return lock
 
 
 async def send_decoded_command(
@@ -102,49 +112,54 @@ def find_response(response_message: RoborockMessage) -> None:
 
 
 async def send_map_command(mqtt_channel: MqttChannel, request_message: Q7RequestMessage) -> bytes:
-    """Send map upload command and wait for MAP_RESPONSE payload 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()
+    Map requests are serialized per channel so concurrent map calls cannot
+    cross-wire responses between callers.
+    """
 
-    def find_response(response_message: RoborockMessage) -> None:
-        if future.done():
-            return
+    async with _get_map_command_lock(mqtt_channel):
+        roborock_message = encode_mqtt_payload(request_message)
+        future: asyncio.Future[bytes] = asyncio.get_running_loop().create_future()
 
-        if response_message.protocol == RoborockMessageProtocol.MAP_RESPONSE and response_message.payload:
-            future.set_result(response_message.payload)
-            return
+        def find_response(response_message: RoborockMessage) -> None:
+            if future.done():
+                return
 
-        try:
-            decoded_dps = decode_rpc_response(response_message)
-        except RoborockException:
-            return
+            if response_message.protocol == RoborockMessageProtocol.MAP_RESPONSE and response_message.payload:
+                future.set_result(response_message.payload)
+                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})"))
+                decoded_dps = decode_rpc_response(response_message)
+            except RoborockException:
                 return
-            data = inner.get("data")
-            if isinstance(data, dict) and isinstance(data.get("payload"), str):
+
+            for dps_value in decoded_dps.values():
+                if not isinstance(dps_value, str):
+                    continue
                 try:
-                    future.set_result(bytes.fromhex(data["payload"]))
-                except ValueError:
-                    pass
+                    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()
+        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

@@ -35,13 +35,24 @@ 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, *, local_key: str, serial: str, model: str) -> None:
+    map_content: Q7MapContentTrait | None
+
+    def __init__(
+        self,
+        channel: MqttChannel,
+        *,
+        local_key: str | None = None,
+        serial: str | None = None,
+        model: str | None = None,
+    ) -> 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)
+        if local_key and serial and model:
+            self.map_content = Q7MapContentTrait(channel, local_key=local_key, serial=serial, model=model)
+        else:
+            # Keep backwards compatibility for direct callers that only use command/query traits.
+            self.map_content = None
 
     async def query_values(self, props: list[RoborockB01Props]) -> B01Props | None:
         """Query the device for the values of the given Q7 properties."""
@@ -91,14 +102,14 @@ async def start_clean(self) -> None:
             },
         )
 
-    async def clean_segments(self, room_ids: list[int]) -> None:
-        """Start segment/room cleaning for the given room ids."""
+    async def clean_segments(self, segment_ids: list[int]) -> None:
+        """Start segment cleaning for the given ids (Q7 uses room ids)."""
         await self.send(
             command=RoborockB01Q7Methods.SET_ROOM_CLEAN,
             params={
                 "clean_type": CleanTaskTypeMapping.ROOM.code,
                 "ctrl_value": SCDeviceCleanParam.START.code,
-                "room_ids": room_ids,
+                "room_ids": segment_ids,
             },
         )
 
@@ -146,6 +157,12 @@ async def send(self, command: CommandType, params: ParamsType) -> Any:
         )
 
 
-def create(channel: MqttChannel, *, local_key: str, serial: str, model: str) -> Q7PropertiesApi:
+def create(
+    channel: MqttChannel,
+    *,
+    local_key: str | None = None,
+    serial: str | None = None,
+    model: str | None = None,
+) -> 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

@@ -10,11 +10,7 @@
 from roborock.devices.traits.v1.map_content import MapContent
 from roborock.devices.transport.mqtt_channel import MqttChannel
 from roborock.exceptions import RoborockException
-from roborock.map.b01_map_parser import (
-    decode_b01_map_payload,
-    parse_scmap_payload,
-    render_map_png,
-)
+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
 
@@ -23,6 +19,8 @@
 class B01MapContent(MapContent):
     """B01 map content wrapper."""
 
+    rooms: dict[int, str] | None = None
+
 
 def _extract_current_map_id(map_list_response: dict[str, Any] | None) -> int | None:
     if not isinstance(map_list_response, dict):
@@ -77,5 +75,6 @@ async def refresh(self) -> B01MapContent:
         parsed = parse_scmap_payload(inflated)
         self.raw_api_response = raw_payload
         self.map_data = None
+        self.rooms = parsed.rooms
         self.image_content = render_map_png(parsed)
         return self

roborock/map/__init__.py

@@ -1,4 +1,4 @@
-"""Module for Roborock map related data classes."""
+"""Utilities and data classes for working with Roborock maps."""
 
 from .b01_map_parser import B01MapData, decode_b01_map_payload, parse_scmap_payload, render_map_png
 from .map_parser import MapParserConfig, ParsedMapData

roborock/map/b01_map_parser.py

@@ -24,6 +24,7 @@ class B01MapData:
     size_x: int
     size_y: int
     map_data: bytes
+    rooms: dict[int, str] | None = None
 
 
 def _read_varint(buf: bytes, idx: int) -> tuple[int, int]:
@@ -72,12 +73,36 @@ def _parse_map_data_info(blob: bytes) -> bytes:
     raise RoborockException("B01 map payload missing mapData")
 
 
+def _parse_room_data_info(blob: bytes) -> tuple[int | None, str | None]:
+    room_id: int | None = None
+    room_name: str | None = None
+    idx = 0
+    while idx < len(blob):
+        key, idx = _read_varint(blob, idx)
+        field_no = key >> 3
+        wire = key & 0x07
+        if wire == 0:
+            value, idx = _read_varint(blob, idx)
+            if field_no == 1:
+                room_id = int(value)
+        elif wire == 2:
+            value, idx = _read_len_delimited(blob, idx)
+            if field_no == 2:
+                room_name = value.decode("utf-8", errors="replace")
+        elif wire == 5:
+            idx += 4
+        else:
+            raise RoborockException(f"Unsupported wire type {wire} in B01 room data info")
+    return room_id, room_name
+
+
 def parse_scmap_payload(payload: bytes) -> B01MapData:
-    """Parse SCMap protobuf payload and extract occupancy grid bytes."""
+    """Parse SCMap protobuf payload and extract occupancy grid bytes and room names."""
 
     size_x = 0
     size_y = 0
     grid = b""
+    rooms: dict[int, str] = {}
     idx = 0
     while idx < len(payload):
         key, idx = _read_varint(payload, idx)
@@ -112,12 +137,16 @@ def parse_scmap_payload(payload: bytes) -> B01MapData:
                     raise RoborockException(f"Unsupported wire type {hwire} in B01 map header")
         elif field_no == 4:  # mapDataInfo
             grid = _parse_map_data_info(value)
+        elif field_no == 12:  # roomDataInfo (repeated)
+            room_id, room_name = _parse_room_data_info(value)
+            if room_id is not None:
+                rooms[room_id] = room_name or f"Room {room_id}"
 
     if not size_x or not size_y or not grid:
         raise RoborockException("Failed to parse B01 map header/grid")
     if len(grid) < size_x * size_y:
         raise RoborockException("B01 map data shorter than expected dimensions")
-    return B01MapData(size_x=size_x, size_y=size_y, map_data=grid)
+    return B01MapData(size_x=size_x, size_y=size_y, map_data=grid, rooms=rooms or None)
 
 
 def _derive_b01_iv(iv_seed: int) -> bytes:

tests/devices/traits/b01/q7/test_init.py

@@ -378,3 +378,9 @@ async def test_q7_map_content_refresh_errors_without_map_list(
 
     with pytest.raises(RoborockException, match="Unable to determine map_id"):
         await q7_api.map_content.refresh()
+
+
+async def test_q7_api_constructor_backwards_compatible_without_map_context(fake_channel: FakeChannel):
+    """Direct API construction without map context should still work."""
+    api = Q7PropertiesApi(fake_channel)  # type: ignore[arg-type]
+    assert api.map_content is None

tests/map/test_b01_map_parser.py

@@ -20,6 +20,8 @@ def test_parse_scmap_payload_fixture() -> None:
     assert parsed.size_x == 340
     assert parsed.size_y == 300
     assert len(parsed.map_data) >= parsed.size_x * parsed.size_y
+    assert parsed.rooms is not None
+    assert parsed.rooms.get(10) == "room1"
 
 
 def test_render_map_png_fixture() -> None: