Skip to content

Commit 99a63e6

Browse files
committed
refactor(q7): use generated protobuf message types
1 parent 3259338 commit 99a63e6

File tree

2 files changed

+15
-79
lines changed

2 files changed

+15
-79
lines changed

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@
88
For B01/Q7 devices, the underlying raw map payload is retrieved via `MapTrait`.
99
"""
1010

11-
from __future__ import annotations
12-
1311
from dataclasses import dataclass
1412

1513
from vacuum_map_parser_base.map_data import MapData

roborock/map/b01_map_parser.py

Lines changed: 15 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,12 @@
1010
`roborock/map/proto/b01_scmap.proto`.
1111
"""
1212

13-
from __future__ import annotations
14-
1513
import base64
1614
import binascii
1715
import hashlib
1816
import io
1917
import zlib
2018
from dataclasses import dataclass, field
21-
from typing import Protocol, cast
2219

2320
from Crypto.Cipher import AES
2421
from Crypto.Util.Padding import pad, unpad
@@ -28,80 +25,21 @@
2825
from vacuum_map_parser_base.map_data import ImageData, MapData
2926

3027
from roborock.exceptions import RoborockException
31-
from roborock.map.proto import b01_scmap_pb2
28+
from roborock.map.proto.b01_scmap_pb2 import ( # type: ignore[attr-defined]
29+
DevicePointInfo,
30+
MapBoundaryInfo,
31+
MapExtInfo,
32+
MapHeadInfo,
33+
RobotMap,
34+
RoomDataInfo,
35+
)
3236

3337
from .map_parser import ParsedMapData
3438

3539
_B64_CHARS = set(b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=")
3640
_MAP_FILE_FORMAT = "PNG"
3741

3842

39-
class _ProtoMessage(Protocol):
40-
def HasField(self, field_name: str) -> bool: ...
41-
42-
43-
class _ScPointMessage(_ProtoMessage, Protocol):
44-
x: float
45-
y: float
46-
47-
48-
class _ScMapBoundaryInfoMessage(_ProtoMessage, Protocol):
49-
mapMd5: str
50-
vMinX: int
51-
vMaxX: int
52-
vMinY: int
53-
vMaxY: int
54-
55-
56-
class _ScMapExtInfoMessage(_ProtoMessage, Protocol):
57-
taskBeginDate: int
58-
mapUploadDate: int
59-
mapValid: int
60-
radian: int
61-
force: int
62-
cleanPath: int
63-
boudaryInfo: _ScMapBoundaryInfoMessage
64-
mapVersion: int
65-
mapValueType: int
66-
67-
68-
class _ScMapHeadMessage(_ProtoMessage, Protocol):
69-
mapHeadId: int
70-
sizeX: int
71-
sizeY: int
72-
minX: float
73-
minY: float
74-
maxX: float
75-
maxY: float
76-
resolution: float
77-
78-
79-
class _ScRoomDataMessage(_ProtoMessage, Protocol):
80-
roomId: int
81-
roomName: str
82-
roomTypeId: int
83-
meterialId: int
84-
cleanState: int
85-
roomClean: int
86-
roomCleanIndex: int
87-
roomNamePost: _ScPointMessage
88-
colorId: int
89-
floor_direction: int
90-
global_seq: int
91-
92-
93-
class _ScMapDataContainerMessage(_ProtoMessage, Protocol):
94-
mapData: bytes
95-
96-
97-
class _ScMapMessage(_ProtoMessage, Protocol):
98-
mapType: int
99-
mapExtInfo: _ScMapExtInfoMessage
100-
mapHead: _ScMapHeadMessage
101-
mapData: _ScMapDataContainerMessage
102-
roomDataInfo: list[_ScRoomDataMessage]
103-
104-
10543
@dataclass(frozen=True)
10644
class _ScPoint:
10745
x: float | None = None
@@ -265,14 +203,14 @@ def _decode_map_data_bytes(value: bytes) -> bytes:
265203
return value
266204

267205

268-
def _parse_sc_point(parsed: _ScPointMessage) -> _ScPoint:
206+
def _parse_sc_point(parsed: DevicePointInfo) -> _ScPoint:
269207
return _ScPoint(
270208
x=parsed.x if parsed.HasField("x") else None,
271209
y=parsed.y if parsed.HasField("y") else None,
272210
)
273211

274212

275-
def _parse_sc_map_boundary_info(parsed: _ScMapBoundaryInfoMessage) -> _ScMapBoundaryInfo:
213+
def _parse_sc_map_boundary_info(parsed: MapBoundaryInfo) -> _ScMapBoundaryInfo:
276214
return _ScMapBoundaryInfo(
277215
map_md5=parsed.mapMd5 if parsed.HasField("mapMd5") else None,
278216
v_min_x=parsed.vMinX if parsed.HasField("vMinX") else None,
@@ -282,7 +220,7 @@ def _parse_sc_map_boundary_info(parsed: _ScMapBoundaryInfoMessage) -> _ScMapBoun
282220
)
283221

284222

285-
def _parse_sc_map_ext_info(parsed: _ScMapExtInfoMessage) -> _ScMapExtInfo:
223+
def _parse_sc_map_ext_info(parsed: MapExtInfo) -> _ScMapExtInfo:
286224
return _ScMapExtInfo(
287225
task_begin_date=parsed.taskBeginDate if parsed.HasField("taskBeginDate") else None,
288226
map_upload_date=parsed.mapUploadDate if parsed.HasField("mapUploadDate") else None,
@@ -296,7 +234,7 @@ def _parse_sc_map_ext_info(parsed: _ScMapExtInfoMessage) -> _ScMapExtInfo:
296234
)
297235

298236

299-
def _parse_sc_map_head(parsed: _ScMapHeadMessage) -> _ScMapHead:
237+
def _parse_sc_map_head(parsed: MapHeadInfo) -> _ScMapHead:
300238
return _ScMapHead(
301239
map_head_id=parsed.mapHeadId if parsed.HasField("mapHeadId") else None,
302240
size_x=parsed.sizeX if parsed.HasField("sizeX") else None,
@@ -309,7 +247,7 @@ def _parse_sc_map_head(parsed: _ScMapHeadMessage) -> _ScMapHead:
309247
)
310248

311249

312-
def _parse_sc_room_data(parsed: _ScRoomDataMessage) -> _ScRoomData:
250+
def _parse_sc_room_data(parsed: RoomDataInfo) -> _ScRoomData:
313251
return _ScRoomData(
314252
room_id=parsed.roomId if parsed.HasField("roomId") else None,
315253
room_name=parsed.roomName if parsed.HasField("roomName") else None,
@@ -327,8 +265,8 @@ def _parse_sc_room_data(parsed: _ScRoomDataMessage) -> _ScRoomData:
327265

328266
def _parse_scmap_payload(payload: bytes) -> _ScMapPayload:
329267
"""Parse inflated SCMap bytes into typed map metadata."""
330-
parsed = cast(_ScMapMessage, getattr(b01_scmap_pb2, "RobotMap")())
331-
_parse_proto(payload, cast(Message, parsed), context="B01 SCMap")
268+
parsed = RobotMap()
269+
_parse_proto(payload, parsed, context="B01 SCMap")
332270

333271
map_data = None
334272
if parsed.HasField("mapData"):

0 commit comments

Comments
 (0)