refactor(map): address Q10 decoding review feedback

roborock/map/b01_grid_layers.py

@@ -15,6 +15,7 @@ class (background / wall / per-room floor / ...). This module turns such a grid
 import io
 from collections.abc import Callable, Iterable
 from dataclasses import dataclass, field
+from enum import IntEnum
 from math import ceil
 
 from PIL import Image
@@ -28,6 +29,14 @@ class (background / wall / per-room floor / ...). This module turns such a grid
 _PNG = "PNG"
 
 
+class _CalibrationCell(IntEnum):
+    """Cell categories used while scoring calibration candidates."""
+
+    OTHER = 0
+    FLOOR = 1
+    BLOCKED = 2
+
+
 @dataclass
 class RoomLayer:
     """A single room (segment) and where its pixels sit in the grid."""
@@ -140,6 +149,21 @@ def pixel_to_world(self, px: float, py: float) -> tuple[float, float]:
         return ((px - self.origin_x) * self.resolution, self.y_sign * (self.origin_y - py) * self.resolution)
 
 
+def _calibration_cells(layers: GridLayers) -> bytes:
+    """Classify the flat, row-major grid for calibration scoring."""
+    cells = bytearray()
+    for value in layers.grid:
+        layer = layers.cell_class(value)
+        if layer == LAYER_FLOOR:
+            cell = _CalibrationCell.FLOOR
+        elif layer in (LAYER_WALL, LAYER_BACKGROUND):
+            cell = _CalibrationCell.BLOCKED
+        else:
+            cell = _CalibrationCell.OTHER
+        cells.append(cell)
+    return bytes(cells)
+
+
 def solve_calibration(
     layers: GridLayers,
     points: list[tuple[float, float]],
@@ -162,11 +186,7 @@ def solve_calibration(
     if not points:
         return None
     w, h = layers.width, layers.height
-    classify = layers.classifier
-    # 1 = floor, 2 = wall/background (blocked), 0 = other. Index by cell.
-    klass = bytes(
-        1 if (c := classify(v)) == LAYER_FLOOR else 2 if c in (LAYER_WALL, LAYER_BACKGROUND) else 0 for v in layers.grid
-    )
+    cells = _calibration_cells(layers)
 
     best: tuple[float, GridCalibration] | None = None
     for resolution in resolutions:
@@ -186,10 +206,10 @@ def solve_calibration(
                     blocked = 0
                     for px_f, py_f in pts:
                         cell = int(oy - py_f) * w + int(px_f + ox)
-                        k = klass[cell]
-                        if k == 1:
+                        cell_type = cells[cell]
+                        if cell_type == _CalibrationCell.FLOOR:
                             on_floor += 1
-                        elif k == 2:
+                        elif cell_type == _CalibrationCell.BLOCKED:
                             blocked += 1
                     score = on_floor - 1.5 * blocked
                     if best is None or score > best[0]:
@@ -223,11 +243,7 @@ def solve_calibration_with_origin(
         return None
     w, h = layers.width, layers.height
     ox, oy = origin
-    classify = layers.classifier
-    # 1 = floor, 2 = wall/background (blocked), 0 = other. Index by cell.
-    klass = bytes(
-        1 if (c := classify(v)) == LAYER_FLOOR else 2 if c in (LAYER_WALL, LAYER_BACKGROUND) else 0 for v in layers.grid
-    )
+    cells = _calibration_cells(layers)
 
     best: tuple[float, GridCalibration] | None = None
     for resolution in resolutions:
@@ -242,10 +258,10 @@ def solve_calibration_with_origin(
                 if not (0 <= px < w and 0 <= py < h):
                     blocked += 1
                     continue
-                k = klass[py * w + px]
-                if k == 1:
+                cell_type = cells[py * w + px]
+                if cell_type == _CalibrationCell.FLOOR:
                     on_floor += 1
-                elif k == 2:
+                elif cell_type == _CalibrationCell.BLOCKED:
                     blocked += 1
             score = on_floor - 1.5 * blocked
             if best is None or score > best[0]:

roborock/map/b01_q10_map_parser.py

@@ -64,13 +64,6 @@ def classify_q10_cell(value: int) -> str:
     return LAYER_FLOOR
 
 
-def decompose_layers(packet: "Q10MapPacket") -> GridLayers:
-    """Split a parsed Q10 map packet into separable grid-pixel layers."""
-    rooms = [(room.id, room.name, room.pixel_value, room.pixel_count) for room in packet.rooms]
-    # The ss07 grid is stored top-down (row 0 = top), so no display flip is applied.
-    return decompose_grid(packet.width, packet.height, packet.grid, rooms, classify_q10_cell, flip=False)
-
-
 MAP_PACKET_MARKER = b"\x01\x01"
 TRACE_PACKET_MARKER = b"\x02\x01"
 
@@ -202,6 +195,13 @@ class Q10MapPacket:
     the same (top-down) pixel space as :attr:`grid`, where a non-zero cell is
     carpet (the value is the carpet kind). ``None`` if the packet carried none."""
 
+    @property
+    def layers(self) -> GridLayers:
+        """Split the occupancy grid into separable grid-pixel layers."""
+        rooms = [(room.id, room.name, room.pixel_value, room.pixel_count) for room in self.rooms]
+        # The ss07 grid is stored top-down (row 0 = top), so no display flip is applied.
+        return decompose_grid(self.width, self.height, self.grid, rooms, classify_q10_cell, flip=False)
+
 
 @dataclass
 class Q10Point:

roborock/map/b01_q10_overlays.py

@@ -27,6 +27,7 @@
 """
 
 import base64
+import binascii
 from dataclasses import dataclass, field
 
 _DEFAULT_RECORD_SIZE = 38  # 2-byte record header + up to 9 (x, y) int16 pairs
@@ -40,25 +41,23 @@ class Q10Zone:
     vertices: list[tuple[int, int]] = field(default_factory=list)
 
 
-def _as_bytes(data: bytes | str | None) -> bytes:
+def _decode_blob(data: str | None) -> bytes:
     if data is None:
         return b""
-    if isinstance(data, bytes):
-        return data
     try:
         return base64.b64decode(data + "=" * (-len(data) % 4))
-    except (ValueError, base64.binascii.Error):  # type: ignore[attr-defined]
+    except (ValueError, binascii.Error):
         return b""
 
 
-def parse_zone_blob(data: bytes | str | None) -> list[Q10Zone]:
+def parse_zone_blob(data: str | None) -> list[Q10Zone]:
     """Decode a Q10 zone/wall overlay blob into a list of :class:`Q10Zone`.
 
-    Accepts the raw bytes or the base64 string straight from the data point.
-    Returns ``[]`` for empty/absent/unparsable blobs (the device sends a single
-    ``0x00`` byte when there are none).
+    Accepts the base64 string straight from the data point. Returns ``[]`` for
+    empty/absent/unparsable blobs (the device sends a single ``0x00`` byte when
+    there are none).
     """
-    raw = _as_bytes(data)
+    raw = _decode_blob(data)
     if len(raw) < 2:
         return []
     count = raw[1]
@@ -94,7 +93,7 @@ def parse_zone_blob(data: bytes | str | None) -> list[Q10Zone]:
 _WALL_RECORD_SIZE = 8  # two (x, y) int16-BE endpoints
 
 
-def parse_virtual_wall_blob(data: bytes | str | None) -> list[Q10Zone]:
+def parse_virtual_wall_blob(data: str | None) -> list[Q10Zone]:
     """Decode a Q10 virtual-wall overlay blob (``dpVirtualWallUp`` 57).
 
     Virtual walls use a *different framing* from the restricted-zone DPs handled
@@ -109,17 +108,17 @@ def parse_virtual_wall_blob(data: bytes | str | None) -> list[Q10Zone]:
     can place them onto the map through the same
     :class:`~roborock.map.b01_grid_layers.GridCalibration` as the zones.
 
-    Accepts raw bytes or the base64 string straight from the data point. Returns
-    ``[]`` for empty/absent/unparsable blobs (the device sends a single ``0x00``
-    byte -- base64 ``AA==`` -- when there are none).
+    Accepts the base64 string straight from the data point. Returns ``[]`` for
+    empty/absent/unparsable blobs (the device sends a single ``0x00`` byte --
+    base64 ``AA==`` -- when there are none).
 
     The axis order was confirmed against the app: a horizontal wall drawn below
     a room reads back with x varying and y constant (and the wide RDC no-go zone
     reads back wide), so DP 57 shares DP 55's order. An earlier revision swapped
     the wall axes to ``(y, x)`` -- following a misreading of PR #850's notes --
     which placed every wall transposed 90 degrees from where it was drawn.
     """
-    raw = _as_bytes(data)
+    raw = _decode_blob(data)
     if len(raw) < 1:
         return []
     count = raw[0]

tests/map/test_b01_grid_layers.py

@@ -1,7 +1,6 @@
-"""Tests for the device-agnostic grid->layers decomposition + Q10 classifier."""
+"""Tests for the device-agnostic grid-to-layers decomposition."""
 
 import io
-from pathlib import Path
 
 import pytest
 from PIL import Image
@@ -15,28 +14,6 @@
     solve_calibration,
     solve_calibration_with_origin,
 )
-from roborock.map.b01_q10_map_parser import (
-    classify_q10_cell,
-    decompose_layers,
-    parse_map_packet,
-)
-
-FIXTURE = Path(__file__).resolve().parent / "testdata" / "b01_q10_map.bin"
-
-
-@pytest.mark.parametrize(
-    ("value", "expected"),
-    [
-        (0, "unknown"),
-        (8, LAYER_FLOOR),
-        (12, LAYER_FLOOR),
-        (240, LAYER_FLOOR),
-        (243, LAYER_BACKGROUND),
-        (249, LAYER_WALL),
-    ],
-)
-def test_classify_q10_cell(value: int, expected: str) -> None:
-    assert classify_q10_cell(value) == expected
 
 
 def test_decompose_grid_generic_classifier_and_bbox() -> None:
@@ -97,21 +74,8 @@ def test_render_scale_upsamples() -> None:
     assert Image.open(io.BytesIO(png)).size == (6, 3)
 
 
-def test_decompose_layers_on_q10_fixture() -> None:
-    """The Q10 synthetic fixture splits into floor + per-room layers."""
-    layers = decompose_layers(parse_map_packet(FIXTURE.read_bytes()))
-    assert layers.class_counts.get(LAYER_FLOOR) == 26
-    names = {room.id: room.name for room in layers.rooms}
-    assert names == {2: "Living Room", 3: "Bedroom"}
-    # Each room renders to a valid PNG and only its own pixels are opaque.
-    living = layers.render_room(2, (255, 0, 0, 255))
-    img = Image.open(io.BytesIO(living))
-    opaque = sum(1 for *_rgb, a in img.getdata() if a > 0)
-    assert opaque == next(r.pixel_count for r in layers.rooms if r.id == 2)
-
-
 def test_render_room_unknown_id_raises() -> None:
-    layers = decompose_layers(parse_map_packet(FIXTURE.read_bytes()))
+    layers = decompose_grid(1, 1, b"\x01", [(1, "Room", 1, 1)], lambda _: LAYER_FLOOR)
     with pytest.raises(KeyError):
         layers.render_room(999, (0, 0, 0, 255))
 

tests/map/test_b01_q10_map_parser.py

@@ -1,13 +1,17 @@
 """Tests for the Roborock Q10 (B01/ss07) map parser."""
 
+import io
 from pathlib import Path
 
 import pytest
+from PIL import Image
 
 from roborock.exceptions import RoborockException
+from roborock.map.b01_grid_layers import LAYER_BACKGROUND, LAYER_FLOOR, LAYER_WALL
 from roborock.map.b01_q10_map_parser import (
     B01Q10MapParser,
     Q10Room,
+    classify_q10_cell,
     is_map_packet,
     is_trace_packet,
     lz4_block_decompress,
@@ -118,6 +122,33 @@ def test_parse_map_packet() -> None:
     assert [(r.id, r.raw_name) for r in packet.rooms] == [(2, "rr_living_room"), (3, "bedroom")]
 
 
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [
+        (0, "unknown"),
+        (8, LAYER_FLOOR),
+        (12, LAYER_FLOOR),
+        (240, LAYER_FLOOR),
+        (243, LAYER_BACKGROUND),
+        (249, LAYER_WALL),
+    ],
+)
+def test_classify_q10_cell(value: int, expected: str) -> None:
+    assert classify_q10_cell(value) == expected
+
+
+def test_packet_layers_decompose_q10_fixture() -> None:
+    """The Q10 synthetic fixture splits into floor + per-room layers."""
+    layers = parse_map_packet(_payload()).layers
+    assert layers.class_counts.get(LAYER_FLOOR) == 26
+    assert {room.id: room.name for room in layers.rooms} == {2: "Living Room", 3: "Bedroom"}
+
+    living = layers.render_room(2, (255, 0, 0, 255))
+    image = Image.open(io.BytesIO(living))
+    opaque = sum(1 for *_rgb, alpha in image.getdata() if alpha > 0)
+    assert opaque == next(room.pixel_count for room in layers.rooms if room.id == 2)
+
+
 def test_parse_map_packet_allows_zero_room_metadata() -> None:
     """A map can be present before the robot has room segmentation records."""
     grid = bytes([240, 240, 249, 243, 240, 240])

tests/map/test_b01_q10_overlays.py

@@ -24,6 +24,10 @@ def _rect(zone_type: int, corners: list[tuple[int, int]]) -> bytes:
     return out
 
 
+def _encoded(blob: bytes) -> str:
+    return base64.b64encode(blob).decode()
+
+
 def test_zone_type_constants() -> None:
     """ss07 + ioBroker: 0 no-go, 1 virtual wall, 2 no-mop, 3 threshold."""
     assert (ZONE_TYPE_NO_GO, ZONE_TYPE_VIRTUAL_WALL, ZONE_TYPE_NO_MOP, ZONE_TYPE_THRESHOLD) == (0, 1, 2, 3)
@@ -33,44 +37,43 @@ def test_parse_zone_blob_distinguishes_no_mop_and_threshold() -> None:
     """A no-mop (2) and a door-threshold (3) zone keep distinct types."""
     no_mop = _rect(ZONE_TYPE_NO_MOP, [(0, 0), (10, 0), (10, 10), (0, 10)])
     threshold = _rect(ZONE_TYPE_THRESHOLD, [(20, 20), (30, 20), (30, 22), (20, 22)])
-    zones = parse_zone_blob(_blob(1, [no_mop, threshold]))
+    zones = parse_zone_blob(_encoded(_blob(1, [no_mop, threshold])))
     assert [z.type for z in zones] == [ZONE_TYPE_NO_MOP, ZONE_TYPE_THRESHOLD]
 
 
 def test_parse_zone_blob_two_typed_rectangles() -> None:
     rect_a = _rect(ZONE_TYPE_NO_GO, [(0, 0), (10, 0), (10, 10), (0, 10)])
     rect_b = _rect(ZONE_TYPE_NO_MOP, [(-5, -5), (5, -5), (5, 5), (-5, 5)])
-    zones = parse_zone_blob(_blob(1, [rect_a, rect_b]))
+    zones = parse_zone_blob(_encoded(_blob(1, [rect_a, rect_b])))
     assert [z.type for z in zones] == [ZONE_TYPE_NO_GO, ZONE_TYPE_NO_MOP]
     assert zones[0].vertices == [(0, 0), (10, 0), (10, 10), (0, 10)]
     assert zones[1].vertices == [(-5, -5), (5, -5), (5, 5), (-5, 5)]  # signed coords
 
 
 def test_parse_zone_blob_accepts_base64() -> None:
     blob = _blob(1, [_rect(ZONE_TYPE_NO_GO, [(1, 2), (3, 4), (5, 6), (7, 8)])])
-    zones = parse_zone_blob(base64.b64encode(blob).decode())
+    zones = parse_zone_blob(_encoded(blob))
     assert len(zones) == 1 and zones[0].vertices[2] == (5, 6)
 
 
 def test_parse_zone_blob_empty_variants() -> None:
     assert parse_zone_blob(None) == []
-    assert parse_zone_blob(b"\x00") == []  # device's "no zones" sentinel
     assert parse_zone_blob("AA==") == []  # base64 of 0x00
-    assert parse_zone_blob(bytes([1, 0, 0])) == []  # version=1, count=0
+    assert parse_zone_blob("AQAA") == []  # version=1, count=0
 
 
 def test_parse_zone_blob_skips_malformed_record() -> None:
     # vertex_count claims 9 verts (needs 38 bytes) but record is only 18 -> skipped.
     bad = bytes([ZONE_TYPE_NO_GO, 9]) + b"\x00" * 16
     good = _rect(ZONE_TYPE_NO_GO, [(1, 1), (2, 2), (3, 3), (4, 4)])
-    zones = parse_zone_blob(_blob(1, [bad, good]))
+    zones = parse_zone_blob(_encoded(_blob(1, [bad, good])))
     assert len(zones) == 1 and zones[0].vertices[0] == (1, 1)
 
 
 def test_parse_zone_blob_real_record_size_inferred() -> None:
     """Record size is inferred from total/count (real device uses 38)."""
     rect = _rect(ZONE_TYPE_NO_GO, [(100, 200), (300, 200), (300, 50), (100, 50)])
-    zones = parse_zone_blob(_blob(1, [rect], record_size=38))
+    zones = parse_zone_blob(_encoded(_blob(1, [rect], record_size=38)))
     assert len(zones) == 1 and zones[0].vertices[0] == (100, 200)
 
 
@@ -148,22 +151,21 @@ def test_parse_virtual_wall_blob_real_rdc_two_walls() -> None:
 
 def test_parse_virtual_wall_blob_empty_variants() -> None:
     assert parse_virtual_wall_blob(None) == []
-    assert parse_virtual_wall_blob(b"\x00") == []  # device's "no walls" sentinel
     assert parse_virtual_wall_blob("AA==") == []  # base64 of 0x00
 
 
 def test_parse_virtual_wall_blob_multiple_walls() -> None:
     """Two walls back-to-back; each is a separate 8-byte (x, y) record."""
     wall_a = bytes.fromhex("000a0014001e0028")  # (x,y)=(10,20)->(30,40)
     wall_b = bytes.fromhex("fffb0005fff6000a")  # (x,y)=(-5,5)->(-10,10)
-    walls = parse_virtual_wall_blob(bytes([2]) + wall_a + wall_b)
+    walls = parse_virtual_wall_blob(_encoded(bytes([2]) + wall_a + wall_b))
     assert [w.vertices for w in walls] == [[(10, 20), (30, 40)], [(-5, 5), (-10, 10)]]
 
 
 def test_parse_virtual_wall_blob_truncated_record_dropped() -> None:
     """A trailing record shorter than 8 bytes is dropped, not misread."""
     blob = bytes([2]) + bytes([0x00, 0x0A, 0x00, 0x14, 0x00, 0x1E, 0x00, 0x28]) + b"\x00\x00"
-    walls = parse_virtual_wall_blob(blob)
+    walls = parse_virtual_wall_blob(_encoded(blob))
     assert [w.vertices for w in walls] == [[(10, 20), (30, 40)]]