refactor: trim q7 map parser scope

Author: arduano

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

@@ -1,7 +1,8 @@
 """Trait for fetching parsed map content from B01/Q7 devices.
 
-This follows the same basic pattern as the v1 `MapContentTrait`:
+This intentionally mirrors the v1 `MapContentTrait` contract:
 - `refresh()` performs I/O and populates cached fields
+- `parse_map_content()` reparses cached raw bytes without I/O
 - 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`.
@@ -75,8 +76,8 @@ async def refresh(self) -> None:
     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.
+        This mirrors the v1 trait behavior so cached map payload bytes can be
+        reparsed without going back to the device.
         """
         if not self._serial or not self._model:
             raise RoborockException(

roborock/map/b01_map_parser.py

@@ -6,11 +6,9 @@
 - PKCS7 padded
 - ASCII hex for a zlib-compressed SCMap payload
 
-The inner SCMap blob is a protobuf-wire message. We know the app's field layout
-well enough to describe the fields we care about declaratively, but we still
-avoid shipping generated protobuf classes from a reverse-engineered schema.
-That keeps the runtime parser narrow without overstating certainty about the
-full message definition.
+The inner SCMap blob uses protobuf wire encoding. We keep the runtime parser
+minimal and declarative: only the fields needed to render the occupancy image
+and expose room names are described here.
 """
 
 from __future__ import annotations
@@ -19,10 +17,9 @@
 import binascii
 import hashlib
 import io
-import struct
 import zlib
 from collections.abc import Callable
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 
 from Crypto.Cipher import AES
 from Crypto.Util.Padding import pad, unpad
@@ -52,70 +49,6 @@ class _ProtoField:
     parser: Callable[[object], object] | None = None
 
 
-@dataclass(frozen=True)
-class _ScPoint:
-    x: float | None = None
-    y: float | None = None
-
-
-@dataclass(frozen=True)
-class _ScMapBoundaryInfo:
-    map_md5: str | None = None
-    v_min_x: int | None = None
-    v_max_x: int | None = None
-    v_min_y: int | None = None
-    v_max_y: int | None = None
-
-
-@dataclass(frozen=True)
-class _ScMapExtInfo:
-    task_begin_date: int | None = None
-    map_upload_date: int | None = None
-    map_valid: int | None = None
-    radian: int | None = None
-    force: int | None = None
-    clean_path: int | None = None
-    boundary_info: _ScMapBoundaryInfo | None = None
-    map_version: int | None = None
-    map_value_type: int | None = None
-
-
-@dataclass(frozen=True)
-class _ScMapHead:
-    map_head_id: int | None = None
-    size_x: int | None = None
-    size_y: int | None = None
-    min_x: float | None = None
-    min_y: float | None = None
-    max_x: float | None = None
-    max_y: float | None = None
-    resolution: float | None = None
-
-
-@dataclass(frozen=True)
-class _ScRoomData:
-    room_id: int | None = None
-    room_name: str | None = None
-    room_type_id: int | None = None
-    material_id: int | None = None
-    clean_state: int | None = None
-    room_clean: int | None = None
-    room_clean_index: int | None = None
-    room_name_post: _ScPoint | None = None
-    color_id: int | None = None
-    floor_direction: int | None = None
-    global_seq: int | None = None
-
-
-@dataclass(frozen=True)
-class _ScMapPayload:
-    map_type: int | None = None
-    map_ext_info: _ScMapExtInfo | None = None
-    map_head: _ScMapHead | None = None
-    map_data: bytes | None = None
-    room_data_info: tuple[_ScRoomData, ...] = field(default_factory=tuple)
-
-
 @dataclass
 class B01MapParserConfig:
     """Configuration for the B01/Q7 map parser."""
@@ -133,9 +66,7 @@ def __init__(self, config: B01MapParserConfig | None = None) -> None:
     def parse(self, raw_payload: bytes, *, serial: str, model: str) -> ParsedMapData:
         """Parse a raw MAP_RESPONSE payload and return a PNG + MapData."""
         inflated = _decode_b01_map_payload(raw_payload, serial=serial, model=model)
-        scmap = _parse_scmap_payload(inflated)
-        size_x, size_y, grid = _extract_grid(scmap)
-        room_names = _extract_room_names(scmap.room_data_info)
+        size_x, size_y, grid, room_names = _parse_scmap_payload(inflated)
 
         image = _render_occupancy_image(grid, size_x=size_x, size_y=size_y, scale=self._config.map_scale)
 
@@ -247,10 +178,6 @@ def _decode_utf8(value: bytes) -> str:
     return value.decode("utf-8", errors="replace")
 
 
-def _decode_float32(value: bytes) -> float:
-    return struct.unpack("<f", value)[0]
-
-
 def _decode_map_data_bytes(value: bytes) -> bytes:
     try:
         return zlib.decompress(value)
@@ -291,145 +218,47 @@ def _parse_proto_message(blob: bytes, schema: dict[int, _ProtoField], *, context
     return parsed
 
 
-def _parse_sc_point(blob: bytes) -> _ScPoint:
-    parsed = _parse_proto_message(blob, _DEVICE_POINT_INFO_SCHEMA, context="B01 DevicePointInfo")
-    return _ScPoint(x=parsed.get("x"), y=parsed.get("y"))
-
-
-def _parse_sc_map_boundary_info(blob: bytes) -> _ScMapBoundaryInfo:
-    parsed = _parse_proto_message(blob, _MAP_BOUNDARY_INFO_SCHEMA, context="B01 MapBoundaryInfo")
-    return _ScMapBoundaryInfo(
-        map_md5=parsed.get("map_md5"),
-        v_min_x=parsed.get("v_min_x"),
-        v_max_x=parsed.get("v_max_x"),
-        v_min_y=parsed.get("v_min_y"),
-        v_max_y=parsed.get("v_max_y"),
-    )
-
-
-def _parse_sc_map_ext_info(blob: bytes) -> _ScMapExtInfo:
-    parsed = _parse_proto_message(blob, _MAP_EXT_INFO_SCHEMA, context="B01 MapExtInfo")
-    return _ScMapExtInfo(
-        task_begin_date=parsed.get("task_begin_date"),
-        map_upload_date=parsed.get("map_upload_date"),
-        map_valid=parsed.get("map_valid"),
-        radian=parsed.get("radian"),
-        force=parsed.get("force"),
-        clean_path=parsed.get("clean_path"),
-        boundary_info=parsed.get("boundary_info"),
-        map_version=parsed.get("map_version"),
-        map_value_type=parsed.get("map_value_type"),
-    )
+def _parse_map_head_info(blob: bytes) -> tuple[int | None, int | None]:
+    parsed = _parse_proto_message(blob, _MAP_HEAD_INFO_SCHEMA, context="B01 MapHeadInfo")
+    return parsed.get("size_x"), parsed.get("size_y")
 
 
-def _parse_sc_map_head(blob: bytes) -> _ScMapHead:
-    parsed = _parse_proto_message(blob, _MAP_HEAD_INFO_SCHEMA, context="B01 MapHeadInfo")
-    return _ScMapHead(
-        map_head_id=parsed.get("map_head_id"),
-        size_x=parsed.get("size_x"),
-        size_y=parsed.get("size_y"),
-        min_x=parsed.get("min_x"),
-        min_y=parsed.get("min_y"),
-        max_x=parsed.get("max_x"),
-        max_y=parsed.get("max_y"),
-        resolution=parsed.get("resolution"),
-    )
-
-
-def _parse_sc_map_data_info(blob: bytes) -> bytes:
+def _parse_map_data_info(blob: bytes) -> bytes:
     parsed = _parse_proto_message(blob, _MAP_DATA_INFO_SCHEMA, context="B01 MapDataInfo")
     if (map_data := parsed.get("map_data")) is None:
         raise RoborockException("B01 map payload missing mapData")
     return map_data
 
 
-def _parse_sc_room_data(blob: bytes) -> _ScRoomData:
+def _parse_room_data_info(blob: bytes) -> tuple[int | None, str | None]:
     parsed = _parse_proto_message(blob, _ROOM_DATA_INFO_SCHEMA, context="B01 RoomDataInfo")
-    return _ScRoomData(
-        room_id=parsed.get("room_id"),
-        room_name=parsed.get("room_name"),
-        room_type_id=parsed.get("room_type_id"),
-        material_id=parsed.get("material_id"),
-        clean_state=parsed.get("clean_state"),
-        room_clean=parsed.get("room_clean"),
-        room_clean_index=parsed.get("room_clean_index"),
-        room_name_post=parsed.get("room_name_post"),
-        color_id=parsed.get("color_id"),
-        floor_direction=parsed.get("floor_direction"),
-        global_seq=parsed.get("global_seq"),
-    )
-
-
-def _parse_scmap_payload(payload: bytes) -> _ScMapPayload:
-    """Parse inflated SCMap bytes using the reverse-engineered app field layout."""
-    parsed = _parse_proto_message(payload, _ROBOT_MAP_SCHEMA, context="B01 SCMap")
-    return _ScMapPayload(
-        map_type=parsed.get("map_type"),
-        map_ext_info=parsed.get("map_ext_info"),
-        map_head=parsed.get("map_head"),
-        map_data=parsed.get("map_data"),
-        room_data_info=tuple(parsed.get("room_data_info", [])),
-    )
+    return parsed.get("room_id"), parsed.get("room_name")
 
 
-def _extract_grid(scmap: _ScMapPayload) -> tuple[int, int, bytes]:
-    if scmap.map_head is None or scmap.map_data is None:
-        raise RoborockException("Failed to parse B01 map header/grid")
+def _parse_scmap_payload(payload: bytes) -> tuple[int, int, bytes, dict[int, str]]:
+    """Parse inflated SCMap bytes using the reverse-engineered field layout we need."""
+    parsed = _parse_proto_message(payload, _ROBOT_MAP_SCHEMA, context="B01 SCMap")
 
-    size_x = scmap.map_head.size_x or 0
-    size_y = scmap.map_head.size_y or 0
-    if not size_x or not size_y or not scmap.map_data:
+    size_x, size_y = parsed.get("map_head", (None, None))
+    grid = parsed.get("map_data")
+    if not size_x or not size_y or not grid:
         raise RoborockException("Failed to parse B01 map header/grid")
 
     expected_len = size_x * size_y
-    if len(scmap.map_data) < expected_len:
+    if len(grid) < expected_len:
         raise RoborockException("B01 map data shorter than expected dimensions")
 
-    return size_x, size_y, scmap.map_data[:expected_len]
-
-
-def _extract_room_names(rooms: tuple[_ScRoomData, ...]) -> dict[int, str]:
     room_names: dict[int, str] = {}
-    for room in rooms:
-        if room.room_id is not None:
-            room_names[room.room_id] = room.room_name or f"Room {room.room_id}"
-    return room_names
+    for room_id, room_name in parsed.get("room_data_info", []):
+        if room_id is not None:
+            room_names[room_id] = room_name or f"Room {room_id}"
 
+    return size_x, size_y, grid[:expected_len], room_names
 
-_DEVICE_POINT_INFO_SCHEMA = {
-    1: _ProtoField("x", _WIRE_FIXED32, parser=_decode_float32),
-    2: _ProtoField("y", _WIRE_FIXED32, parser=_decode_float32),
-}
-
-_MAP_BOUNDARY_INFO_SCHEMA = {
-    1: _ProtoField("map_md5", _WIRE_LEN, parser=_decode_utf8),
-    2: _ProtoField("v_min_x", _WIRE_VARINT, parser=_decode_uint32),
-    3: _ProtoField("v_max_x", _WIRE_VARINT, parser=_decode_uint32),
-    4: _ProtoField("v_min_y", _WIRE_VARINT, parser=_decode_uint32),
-    5: _ProtoField("v_max_y", _WIRE_VARINT, parser=_decode_uint32),
-}
-
-_MAP_EXT_INFO_SCHEMA = {
-    1: _ProtoField("task_begin_date", _WIRE_VARINT, parser=_decode_uint32),
-    2: _ProtoField("map_upload_date", _WIRE_VARINT, parser=_decode_uint32),
-    3: _ProtoField("map_valid", _WIRE_VARINT, parser=_decode_uint32),
-    4: _ProtoField("radian", _WIRE_VARINT, parser=_decode_uint32),
-    5: _ProtoField("force", _WIRE_VARINT, parser=_decode_uint32),
-    6: _ProtoField("clean_path", _WIRE_VARINT, parser=_decode_uint32),
-    7: _ProtoField("boundary_info", _WIRE_LEN, parser=_parse_sc_map_boundary_info),
-    8: _ProtoField("map_version", _WIRE_VARINT, parser=_decode_uint32),
-    9: _ProtoField("map_value_type", _WIRE_VARINT, parser=_decode_uint32),
-}
 
 _MAP_HEAD_INFO_SCHEMA = {
-    1: _ProtoField("map_head_id", _WIRE_VARINT, parser=_decode_uint32),
     2: _ProtoField("size_x", _WIRE_VARINT, parser=_decode_uint32),
     3: _ProtoField("size_y", _WIRE_VARINT, parser=_decode_uint32),
-    4: _ProtoField("min_x", _WIRE_FIXED32, parser=_decode_float32),
-    5: _ProtoField("min_y", _WIRE_FIXED32, parser=_decode_float32),
-    6: _ProtoField("max_x", _WIRE_FIXED32, parser=_decode_float32),
-    7: _ProtoField("max_y", _WIRE_FIXED32, parser=_decode_float32),
-    8: _ProtoField("resolution", _WIRE_FIXED32, parser=_decode_float32),
 }
 
 _MAP_DATA_INFO_SCHEMA = {
@@ -439,23 +268,12 @@ def _extract_room_names(rooms: tuple[_ScRoomData, ...]) -> dict[int, str]:
 _ROOM_DATA_INFO_SCHEMA = {
     1: _ProtoField("room_id", _WIRE_VARINT, parser=_decode_uint32),
     2: _ProtoField("room_name", _WIRE_LEN, parser=_decode_utf8),
-    3: _ProtoField("room_type_id", _WIRE_VARINT, parser=_decode_uint32),
-    4: _ProtoField("material_id", _WIRE_VARINT, parser=_decode_uint32),
-    5: _ProtoField("clean_state", _WIRE_VARINT, parser=_decode_uint32),
-    6: _ProtoField("room_clean", _WIRE_VARINT, parser=_decode_uint32),
-    7: _ProtoField("room_clean_index", _WIRE_VARINT, parser=_decode_uint32),
-    8: _ProtoField("room_name_post", _WIRE_LEN, parser=_parse_sc_point),
-    10: _ProtoField("color_id", _WIRE_VARINT, parser=_decode_uint32),
-    11: _ProtoField("floor_direction", _WIRE_VARINT, parser=_decode_uint32),
-    12: _ProtoField("global_seq", _WIRE_VARINT, parser=_decode_uint32),
 }
 
 _ROBOT_MAP_SCHEMA = {
-    1: _ProtoField("map_type", _WIRE_VARINT, parser=_decode_uint32),
-    2: _ProtoField("map_ext_info", _WIRE_LEN, parser=_parse_sc_map_ext_info),
-    3: _ProtoField("map_head", _WIRE_LEN, parser=_parse_sc_map_head),
-    4: _ProtoField("map_data", _WIRE_LEN, parser=_parse_sc_map_data_info),
-    12: _ProtoField("room_data_info", _WIRE_LEN, repeated=True, parser=_parse_sc_room_data),
+    3: _ProtoField("map_head", _WIRE_LEN, parser=_parse_map_head_info),
+    4: _ProtoField("map_data", _WIRE_LEN, parser=_parse_map_data_info),
+    12: _ProtoField("room_data_info", _WIRE_LEN, repeated=True, parser=_parse_room_data_info),
 }