🦎 q7: add b01 map content trait

Author: arduano

roborock/devices/device_manager.py

@@ -251,7 +251,7 @@ 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(product, device, channel)
                 else:
                     raise UnsupportedDeviceError(f"Device {device.name} has unsupported B01 model: {product.model}")
             case _:

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

@@ -1,10 +1,14 @@
 """Traits for Q7 B01 devices.
-Potentially other devices may fall into this category in the future."""
+
+Potentially other devices may fall into this category in the future.
+"""
+
+from __future__ import annotations
 
 from typing import Any
 
 from roborock import B01Props
-from roborock.data import Q7MapList, Q7MapListEntry
+from roborock.data import HomeDataDevice, HomeDataProduct, Q7MapList, Q7MapListEntry
 from roborock.data.b01_q7.b01_q7_code_mappings import (
     CleanPathPreferenceMapping,
     CleanRepeatMapping,
@@ -17,36 +21,56 @@
 from roborock.devices.rpc.b01_q7_channel import send_decoded_command
 from roborock.devices.traits import Trait
 from roborock.devices.transport.mqtt_channel import MqttChannel
+from roborock.exceptions import RoborockException
 from roborock.protocols.b01_q7_protocol import B01_Q7_DPS, CommandType, ParamsType, Q7RequestMessage
 from roborock.roborock_message import RoborockB01Props
 from roborock.roborock_typing import RoborockB01Q7Methods
 
 from .clean_summary import CleanSummaryTrait
 from .map import MapTrait
+from .map_content import MapContentTrait
 
 __all__ = [
     "Q7PropertiesApi",
     "CleanSummaryTrait",
     "MapTrait",
+    "MapContentTrait",
     "Q7MapList",
     "Q7MapListEntry",
 ]
 
 
 class Q7PropertiesApi(Trait):
-    """API for interacting with B01 devices."""
+    """API for interacting with B01 Q7 devices."""
 
     clean_summary: CleanSummaryTrait
     """Trait for clean records / clean summary (Q7 `service.get_record_list`)."""
 
     map: MapTrait
     """Trait for map list metadata + raw map payload retrieval."""
 
-    def __init__(self, channel: MqttChannel) -> None:
-        """Initialize the B01Props API."""
+    map_content: MapContentTrait
+    """Trait for fetching parsed current map content."""
+
+    def __init__(self, channel: MqttChannel, *, device: HomeDataDevice, product: HomeDataProduct) -> None:
+        """Initialize the Q7 API."""
         self._channel = channel
+        self._device = device
+        self._product = product
+
+        if not device.sn:
+            raise RoborockException(
+                "B01/Q7 map parsing requires the device serial number (HomeDataDevice.sn), but it was missing"
+            )
+
         self.clean_summary = CleanSummaryTrait(channel)
         self.map = MapTrait(channel)
+        self.map_content = MapContentTrait(
+            self.map,
+            local_key=device.local_key,
+            serial=device.sn,
+            model=product.model,
+        )
 
     async def query_values(self, props: list[RoborockB01Props]) -> B01Props | None:
         """Query the device for the values of the given Q7 properties."""
@@ -151,6 +175,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(product: HomeDataProduct, device: HomeDataDevice, channel: MqttChannel) -> Q7PropertiesApi:
+    """Create traits for B01 Q7 devices."""
+    return Q7PropertiesApi(channel, device=device, product=product)

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

@@ -0,0 +1,100 @@
+"""Trait for fetching parsed map content from B01/Q7 devices.
+
+This follows the same basic pattern as the v1 `MapContentTrait`:
+- `refresh()` performs I/O and populates cached fields
+- fields `image_content`, `map_data`, and `raw_api_response` are then readable
+
+For B01/Q7 devices, the underlying raw map payload is retrieved via `MapTrait`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from vacuum_map_parser_base.map_data import MapData
+
+from roborock.data import RoborockBase
+from roborock.devices.traits import Trait
+from roborock.exceptions import RoborockException
+from roborock.map.b01_map_parser import B01MapParser, B01MapParserConfig
+
+from .map import MapTrait
+
+_TRUNCATE_LENGTH = 20
+
+
+@dataclass
+class MapContent(RoborockBase):
+    """Dataclass representing map content."""
+
+    image_content: bytes | None = None
+    """The rendered image of the map in PNG format."""
+
+    map_data: MapData | None = None
+    """Parsed map data (metadata for points on the map)."""
+
+    raw_api_response: bytes | None = None
+    """Raw bytes of the map payload from the device.
+
+    This should be treated as an opaque blob used only internally by this
+    library to re-parse the map data when needed.
+    """
+
+    def __repr__(self) -> str:
+        img = self.image_content
+        if img and len(img) > _TRUNCATE_LENGTH:
+            img = img[: _TRUNCATE_LENGTH - 3] + b"..."
+        return f"MapContent(image_content={img!r}, map_data={self.map_data!r})"
+
+
+class MapContentTrait(MapContent, Trait):
+    """Trait for fetching parsed map content for Q7 devices."""
+
+    def __init__(
+        self,
+        map_trait: MapTrait,
+        *,
+        local_key: str,
+        serial: str,
+        model: str,
+        map_parser_config: B01MapParserConfig | None = None,
+    ) -> None:
+        super().__init__()
+        self._map_trait = map_trait
+        self._local_key = local_key
+        self._serial = serial
+        self._model = model
+        self._map_parser = B01MapParser(map_parser_config)
+
+    async def refresh(self) -> None:
+        """Fetch, decode, and parse the current map payload."""
+        raw_payload = await self._map_trait.get_current_map_payload()
+        parsed = self.parse_map_content(raw_payload)
+        self.image_content = parsed.image_content
+        self.map_data = parsed.map_data
+        self.raw_api_response = parsed.raw_api_response
+
+    def parse_map_content(self, response: bytes) -> MapContent:
+        """Parse map content from raw bytes.
+
+        Exposed so callers can re-parse cached map payload bytes without
+        performing I/O.
+        """
+        try:
+            parsed_data = self._map_parser.parse(
+                response,
+                local_key=self._local_key,
+                serial=self._serial,
+                model=self._model,
+            )
+        except Exception as ex:
+            raise RoborockException("Failed to parse B01 map data") from ex
+
+        if parsed_data.image_content is None:
+            raise RoborockException("Failed to render B01 map image")
+
+        return MapContent(
+            image_content=parsed_data.image_content,
+            map_data=parsed_data.map_data,
+            raw_api_response=response,
+        )