Add Q7 map_content support and room segment cleaning helper

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/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

@@ -5,6 +5,7 @@
 import asyncio
 import json
 import logging
+import weakref
 from typing import Any
 
 from roborock.devices.transport.mqtt_channel import MqttChannel
@@ -14,10 +15,19 @@
     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
+_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(
@@ -99,3 +109,62 @@ 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.
+
+    Map requests are serialized per channel so concurrent map calls cannot
+    cross-wire responses between callers.
+    """
+
+    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()
+
+        def find_response(response_message: RoborockMessage) -> None:
+            if future.done():
+                return
+
+            if response_message.protocol == RoborockMessageProtocol.MAP_RESPONSE and response_message.payload:
+                if not future.done():
+                    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:
+                    if not future.done():
+                        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:
+                        if not future.done():
+                            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,24 @@ class Q7PropertiesApi(Trait):
 
     clean_summary: CleanSummaryTrait
     """Trait for clean records / clean summary (Q7 `service.get_record_list`)."""
-
-    def __init__(self, channel: MqttChannel) -> 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)
+        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."""
@@ -87,6 +102,17 @@ async def start_clean(self) -> None:
             },
         )
 
+    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": segment_ids,
+            },
+        )
+
     async def pause_clean(self) -> None:
         """Pause cleaning."""
         await self.send(
@@ -131,6 +157,12 @@ 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 | 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

@@ -0,0 +1,80 @@
+"""Map content trait for B01/Q7 devices."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from roborock.devices.rpc.b01_q7_channel import send_decoded_command, send_map_command
+from roborock.devices.traits import Trait
+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.protocols.b01_q7_protocol import Q7RequestMessage
+from roborock.roborock_typing import RoborockB01Q7Methods
+
+
+@dataclass
+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):
+        return None
+    map_list = map_list_response.get("map_list")
+    if not isinstance(map_list, list) or not map_list:
+        return None
+
+    for entry in map_list:
+        if isinstance(entry, dict) and entry.get("cur") and isinstance(entry.get("id"), int):
+            return entry["id"]
+
+    first = map_list[0]
+    if isinstance(first, dict) and isinstance(first.get("id"), int):
+        return first["id"]
+    return None
+
+
+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:
+        map_list_response = await send_decoded_command(
+            self._channel,
+            Q7RequestMessage(dps=10000, command=RoborockB01Q7Methods.GET_MAP_LIST, params={}),
+        )
+        map_id = _extract_current_map_id(map_list_response)
+        if map_id is None:
+            raise RoborockException(f"Unable to determine map_id from map list response: {map_list_response!r}")
+
+        raw_payload = await send_map_command(
+            self._channel,
+            Q7RequestMessage(
+                dps=10000,
+                command=RoborockB01Q7Methods.UPLOAD_BY_MAPID,
+                params={"map_id": map_id},
+            ),
+        )
+        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.rooms = parsed.rooms
+        self.image_content = render_map_png(parsed)
+        return self

roborock/map/__init__.py

@@ -1,7 +1,13 @@
-"""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
 
 __all__ = [
+    "B01MapData",
     "MapParserConfig",
+    "ParsedMapData",
+    "decode_b01_map_payload",
+    "parse_scmap_payload",
+    "render_map_png",
 ]