Add a v1 device trait for map contents (#517)

* feat: Add a v1 device trait for map contents

This uses the map parser library to return an image and the parsed data from a call.

* chore: Refactor to reuse the same payload functions

* chore: remove duplicate code

* chore: fix lint and typing errors

* chore: fix test wording

* chore: fix lint error

* chore: fix formatting

Author: allenporter

roborock/cli.py

@@ -48,6 +48,7 @@
 from roborock.devices.traits import Trait
 from roborock.devices.traits.v1 import V1TraitMixin
 from roborock.devices.traits.v1.consumeable import ConsumableAttribute
+from roborock.devices.traits.v1.map_content import MapContentTrait
 from roborock.protocol import MessageParser
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
 from roborock.web_api import RoborockApiClient
@@ -451,6 +452,49 @@ async def maps(ctx, device_id: str):
     await _display_v1_trait(context, device_id, lambda v1: v1.maps)
 
 
+@session.command()
+@click.option("--device_id", required=True)
+@click.option("--output-file", required=True, help="Path to save the map image.")
+@click.pass_context
+@async_command
+async def map_image(ctx, device_id: str, output_file: str):
+    """Get device map image and save it to a file."""
+    context: RoborockContext = ctx.obj
+    trait: MapContentTrait = await _v1_trait(context, device_id, lambda v1: v1.map_content)
+    if trait.image_content:
+        with open(output_file, "wb") as f:
+            f.write(trait.image_content)
+        click.echo(f"Map image saved to {output_file}")
+    else:
+        click.echo("No map image content available.")
+
+
+@session.command()
+@click.option("--device_id", required=True)
+@click.option("--include_path", is_flag=True, default=False, help="Include path data in the output.")
+@click.pass_context
+@async_command
+async def map_data(ctx, device_id: str, include_path: bool):
+    """Get parsed map data as JSON."""
+    context: RoborockContext = ctx.obj
+    trait: MapContentTrait = await _v1_trait(context, device_id, lambda v1: v1.map_content)
+    if not trait.map_data:
+        click.echo("No parsed map data available.")
+        return
+
+    # Pick some parts of the map data to display.
+    data_summary = {
+        "charger": trait.map_data.charger.as_dict() if trait.map_data.charger else None,
+        "image_size": trait.map_data.image.data.size if trait.map_data.image else None,
+        "vacuum_position": trait.map_data.vacuum_position.as_dict() if trait.map_data.vacuum_position else None,
+        "calibration": trait.map_data.calibration(),
+        "zones": [z.as_dict() for z in trait.map_data.zones or ()],
+    }
+    if include_path and trait.map_data.path:
+        data_summary["path"] = trait.map_data.path.as_dict()
+    click.echo(dump_json(data_summary))
+
+
 @session.command()
 @click.option("--device_id", required=True)
 @click.pass_context
@@ -727,6 +771,8 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur
 cli.add_command(volume)
 cli.add_command(set_volume)
 cli.add_command(maps)
+cli.add_command(map_image)
+cli.add_command(map_data)
 cli.add_command(consumables)
 cli.add_command(reset_consumable)
 cli.add_command(rooms)

roborock/devices/device_manager.py

@@ -14,6 +14,7 @@
     UserData,
 )
 from roborock.devices.device import RoborockDevice
+from roborock.map.map_parser import MapParserConfig
 from roborock.mqtt.roborock_session import create_lazy_mqtt_session
 from roborock.mqtt.session import MqttSession
 from roborock.protocol import create_mqtt_params
@@ -130,6 +131,7 @@ async def create_device_manager(
     user_data: UserData,
     home_data_api: HomeDataApi,
     cache: Cache | None = None,
+    map_parser_config: MapParserConfig | None = None,
 ) -> DeviceManager:
     """Convenience function to create and initialize a DeviceManager.
 
@@ -149,7 +151,14 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
         match device.pv:
             case DeviceVersion.V1:
                 channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
-                trait = v1.create(product, home_data, channel.rpc_channel, channel.mqtt_rpc_channel)
+                trait = v1.create(
+                    product,
+                    home_data,
+                    channel.rpc_channel,
+                    channel.mqtt_rpc_channel,
+                    channel.map_rpc_channel,
+                    map_parser_config=map_parser_config,
+                )
             case DeviceVersion.A01:
                 channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
                 trait = a01.create(product, channel)

roborock/devices/traits/v1/__init__.py

@@ -6,11 +6,13 @@
 from roborock.containers import HomeData, HomeDataProduct
 from roborock.devices.traits import Trait
 from roborock.devices.v1_rpc_channel import V1RpcChannel
+from roborock.map.map_parser import MapParserConfig
 
 from .clean_summary import CleanSummaryTrait
 from .common import V1TraitMixin
 from .consumeable import ConsumableTrait
 from .do_not_disturb import DoNotDisturbTrait
+from .map_content import MapContentTrait
 from .maps import MapsTrait
 from .rooms import RoomsTrait
 from .status import StatusTrait
@@ -26,6 +28,7 @@
     "CleanSummaryTrait",
     "SoundVolumeTrait",
     "MapsTrait",
+    "MapContentTrait",
     "ConsumableTrait",
 ]
 
@@ -44,18 +47,25 @@ class PropertiesApi(Trait):
     sound_volume: SoundVolumeTrait
     rooms: RoomsTrait
     maps: MapsTrait
+    map_content: MapContentTrait
     consumables: ConsumableTrait
 
     # In the future optional fields can be added below based on supported features
 
     def __init__(
-        self, product: HomeDataProduct, home_data: HomeData, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel
+        self,
+        product: HomeDataProduct,
+        home_data: HomeData,
+        rpc_channel: V1RpcChannel,
+        mqtt_rpc_channel: V1RpcChannel,
+        map_rpc_channel: V1RpcChannel,
+        map_parser_config: MapParserConfig | None = None,
     ) -> None:
         """Initialize the V1TraitProps."""
         self.status = StatusTrait(product)
         self.rooms = RoomsTrait(home_data)
         self.maps = MapsTrait(self.status)
-
+        self.map_content = MapContentTrait(map_parser_config)
         # This is a hack to allow setting the rpc_channel on all traits. This is
         # used so we can preserve the dataclass behavior when the values in the
         # traits are updated, but still want to allow them to have a reference
@@ -68,12 +78,19 @@ def __init__(
             # to use the mqtt_rpc_channel (cloud only) instead of the rpc_channel (adaptive)
             if hasattr(trait, "mqtt_rpc_channel"):
                 trait._rpc_channel = mqtt_rpc_channel
+            elif hasattr(trait, "map_rpc_channel"):
+                trait._rpc_channel = map_rpc_channel
             else:
                 trait._rpc_channel = rpc_channel
 
 
 def create(
-    product: HomeDataProduct, home_data: HomeData, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel
+    product: HomeDataProduct,
+    home_data: HomeData,
+    rpc_channel: V1RpcChannel,
+    mqtt_rpc_channel: V1RpcChannel,
+    map_rpc_channel: V1RpcChannel,
+    map_parser_config: MapParserConfig | None = None,
 ) -> PropertiesApi:
     """Create traits for V1 devices."""
-    return PropertiesApi(product, home_data, rpc_channel, mqtt_rpc_channel)
+    return PropertiesApi(product, home_data, rpc_channel, mqtt_rpc_channel, map_rpc_channel, map_parser_config)

roborock/devices/traits/v1/common.py

@@ -133,3 +133,13 @@ def wrapper(*args, **kwargs):
 
     cls.mqtt_rpc_channel = True  # type: ignore[attr-defined]
     return wrapper
+
+
+def map_rpc_channel(cls):
+    """Decorator to mark a function as cloud only using the map rpc format."""
+
+    def wrapper(*args, **kwargs):
+        return cls(*args, **kwargs)
+
+    cls.map_rpc_channel = True  # type: ignore[attr-defined]
+    return wrapper

roborock/devices/traits/v1/map_content.py

@@ -0,0 +1,49 @@
+"""Trait for fetching the map content from Roborock devices."""
+import logging
+from dataclasses import dataclass
+
+from vacuum_map_parser_base.map_data import MapData
+
+from roborock.containers import RoborockBase
+from roborock.devices.traits.v1 import common
+from roborock.map.map_parser import MapParser, MapParserConfig
+from roborock.roborock_typing import RoborockCommand
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@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
+    """The parsed map data which contains metadata for points on the map."""
+
+
+@common.map_rpc_channel
+class MapContentTrait(MapContent, common.V1TraitMixin):
+    """Trait for fetching the map content."""
+
+    command = RoborockCommand.GET_MAP_V1
+
+    def __init__(self, map_parser_config: MapParserConfig | None = None) -> None:
+        """Initialize MapContentTrait."""
+        super().__init__()
+        self._map_parser = MapParser(map_parser_config or MapParserConfig())
+
+    def _parse_response(self, response: common.V1ResponseData) -> MapContent:
+        """Parse the response from the device into a MapContentTrait instance."""
+        if not isinstance(response, bytes):
+            raise ValueError(f"Unexpected MapContentTrait response format: {type(response)}")
+
+        parsed_data = self._map_parser.parse(response)
+        if parsed_data is None:
+            raise ValueError("Failed to parse map data")
+
+        return MapContent(
+            image_content=parsed_data.image_content,
+            map_data=parsed_data.map_data,
+        )

roborock/devices/v1_channel.py

@@ -27,6 +27,7 @@
     PickFirstAvailable,
     V1RpcChannel,
     create_local_rpc_channel,
+    create_map_rpc_channel,
     create_mqtt_rpc_channel,
 )
 
@@ -80,6 +81,7 @@ def __init__(
         self._combined_rpc_channel = PickFirstAvailable(
             [lambda: self._local_rpc_channel, lambda: self._mqtt_rpc_channel]
         )
+        self._map_rpc_channel = create_map_rpc_channel(mqtt_channel, security_data)
         self._mqtt_unsub: Callable[[], None] | None = None
         self._local_unsub: Callable[[], None] | None = None
         self._callback: Callable[[RoborockMessage], None] | None = None
@@ -112,6 +114,11 @@ def mqtt_rpc_channel(self) -> V1RpcChannel:
         """Return the MQTT RPC channel."""
         return self._mqtt_rpc_channel
 
+    @property
+    def map_rpc_channel(self) -> V1RpcChannel:
+        """Return the map RPC channel used for fetching map content."""
+        return self._map_rpc_channel
+
     async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callable[[], None]:
         """Subscribe to all messages from the device.
 
@@ -132,7 +139,6 @@ async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callab
 
         # Start a background task to manage the local connection health. This
         # happens independent of whether we were able to connect locally now.
-        _LOGGER.info("self._reconnect_task=%s", self._reconnect_task)
         if self._reconnect_task is None:
             loop = asyncio.get_running_loop()
             self._reconnect_task = loop.create_task(self._background_reconnect())

roborock/devices/v1_rpc_channel.py

@@ -6,6 +6,7 @@
 and local connections, preferring local when available.
 """
 
+
 import asyncio
 import logging
 from collections.abc import Callable
@@ -15,10 +16,13 @@
 from roborock.exceptions import RoborockException
 from roborock.protocols.v1_protocol import (
     CommandType,
+    MapResponse,
     ParamsType,
     RequestMessage,
     ResponseData,
+    ResponseMessage,
     SecurityData,
+    create_map_response_decoder,
     decode_rpc_response,
 )
 from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
@@ -31,6 +35,7 @@
 
 
 _T = TypeVar("_T", bound=RoborockBase)
+_V = TypeVar("_V")
 
 
 class V1RpcChannel(Protocol):
@@ -120,36 +125,40 @@ def __init__(
         name: str,
         channel: MqttChannel | LocalChannel,
         payload_encoder: Callable[[RequestMessage], RoborockMessage],
+        decoder: Callable[[RoborockMessage], ResponseMessage] | Callable[[RoborockMessage], MapResponse | None],
     ) -> None:
         """Initialize the channel with a raw channel and an encoder function."""
         self._name = name
         self._channel = channel
         self._payload_encoder = payload_encoder
+        self._decoder = decoder
 
     async def _send_raw_command(
         self,
         method: CommandType,
         *,
         params: ParamsType = None,
-    ) -> ResponseData:
+    ) -> ResponseData | bytes:
         """Send a command and return a parsed response RoborockBase type."""
         request_message = RequestMessage(method, params=params)
         _LOGGER.debug(
             "Sending command (%s, request_id=%s): %s, params=%s", self._name, request_message.request_id, method, params
         )
         message = self._payload_encoder(request_message)
 
-        future: asyncio.Future[ResponseData] = asyncio.Future()
+        future: asyncio.Future[ResponseData | bytes] = asyncio.Future()
 
         def find_response(response_message: RoborockMessage) -> None:
             try:
-                decoded = decode_rpc_response(response_message)
+                decoded = self._decoder(response_message)
             except RoborockException as ex:
                 _LOGGER.debug("Exception while decoding message (%s): %s", response_message, ex)
                 return
-            _LOGGER.debug("Received response (request_id=%s): %s", self._name, decoded.request_id)
+            if decoded is None:
+                return
+            _LOGGER.debug("Received response (%s, request_id=%s)", self._name, decoded.request_id)
             if decoded.request_id == request_message.request_id:
-                if decoded.api_error:
+                if isinstance(decoded, ResponseMessage) and decoded.api_error:
                     future.set_exception(decoded.api_error)
                 else:
                     future.set_result(decoded.data)
@@ -171,6 +180,7 @@ def create_mqtt_rpc_channel(mqtt_channel: MqttChannel, security_data: SecurityDa
         "mqtt",
         mqtt_channel,
         lambda x: x.encode_message(RoborockMessageProtocol.RPC_REQUEST, security_data=security_data),
+        decode_rpc_response,
     )
 
 
@@ -180,4 +190,23 @@ def create_local_rpc_channel(local_channel: LocalChannel) -> V1RpcChannel:
         "local",
         local_channel,
         lambda x: x.encode_message(RoborockMessageProtocol.GENERAL_REQUEST),
+        decode_rpc_response,
+    )
+
+
+def create_map_rpc_channel(
+    mqtt_channel: MqttChannel,
+    security_data: SecurityData,
+) -> V1RpcChannel:
+    """Create a V1 RPC channel that fetches map data.
+
+    This will prefer local channels when available, falling back to MQTT
+    channels if not. If neither is available, an exception will be raised
+    when trying to send a command.
+    """
+    return PayloadEncodedV1RpcChannel(
+        "map",
+        mqtt_channel,
+        lambda x: x.encode_message(RoborockMessageProtocol.RPC_REQUEST, security_data=security_data),
+        create_map_response_decoder(security_data=security_data),
     )

roborock/map/__init__.py

@@ -0,0 +1,7 @@
+"""Module for Roborock map related data classes."""
+
+from .map_parser import MapParserConfig, ParsedMapData
+
+__all__ = [
+    "MapParserConfig",
+]

roborock/protocols/v1_protocol.py

@@ -187,7 +187,7 @@ def _decode_map_response(message: RoborockMessage) -> MapResponse | None:
         header, body = message.payload[:24], message.payload[24:]
         [endpoint, _, request_id, _] = struct.unpack("<8s8sH6s", header)
         if not endpoint.decode().startswith(security_data.endpoint):
-            _LOGGER.debug("Received map response requested not made by this device, ignoring.")
+            _LOGGER.debug("Received map response not requested by this device, ignoring.")
             return None
         try:
             decrypted = Utils.decrypt_cbc(body, security_data.nonce)