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.

Author: allenporter

roborock/cli.py

@@ -46,6 +46,7 @@
 from roborock.devices.device_manager import DeviceManager, create_device_manager, create_home_data_api
 from roborock.devices.traits import Trait
 from roborock.devices.traits.v1 import V1TraitMixin
+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
@@ -449,6 +450,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))
+
+
 @click.command()
 @click.option("--device_id", required=True)
 @click.option("--cmd", required=True)
@@ -691,6 +735,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)
 
 
 def main():

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,13 @@ def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> Roborock
         match device.pv:
             case DeviceVersion.V1:
                 channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
-                trait = v1.create(product, channel.rpc_channel, channel.mqtt_rpc_channel)
+                trait = v1.create(
+                    product,
+                    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,10 +6,12 @@
 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 .do_not_disturb import DoNotDisturbTrait
+from .map_content import MapContentTrait
 from .maps import MapsTrait
 from .status import StatusTrait
 from .volume import SoundVolumeTrait
@@ -24,6 +26,7 @@
     "CleanSummaryTrait",
     "SoundVolumeTrait",
     "MapsTrait",
+    "MapContentTrait",
 ]
 
 
@@ -40,14 +43,22 @@ class PropertiesApi(Trait):
     clean_summary: CleanSummaryTrait
     sound_volume: SoundVolumeTrait
     maps: MapsTrait
+    map_content: MapContentTrait
 
     # In the future optional fields can be added below based on supported features
 
-    def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> None:
+    def __init__(
+        self,
+        product: HomeDataProduct,
+        rpc_channel: V1RpcChannel,
+        mqtt_rpc_channel: V1RpcChannel,
+        map_rpc_channel: V1RpcChannel,
+        map_parser_config: MapParserConfig | None = None,
+    ) -> None:
         """Initialize the V1TraitProps with None values."""
         self.status = StatusTrait(product)
         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
@@ -60,10 +71,18 @@ def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc
             # 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, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> PropertiesApi:
+def create(
+    product: HomeDataProduct,
+    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, rpc_channel, mqtt_rpc_channel)
+    return PropertiesApi(product, rpc_channel, mqtt_rpc_channel, map_rpc_channel, map_parser_config)

roborock/devices/traits/v1/common.py

@@ -38,7 +38,7 @@ class V1TraitMixin(ABC):
     command: ClassVar[RoborockCommand]
 
     @classmethod
-    def _parse_type_response(cls, response: V1ResponseData) -> Self:
+    def _parse_type_response(cls, response: V1ResponseData) -> RoborockBase:
         """Parse the response from the device into a a RoborockBase.
 
         Subclasses should override this method to implement custom parsing
@@ -53,7 +53,7 @@ def _parse_type_response(cls, response: V1ResponseData) -> Self:
             raise ValueError(f"Unexpected {cls} response format: {response!r}")
         return cls.from_dict(response)
 
-    def _parse_response(self, response: V1ResponseData) -> Self:
+    def _parse_response(self, response: V1ResponseData) -> RoborockBase:
         """Parse the response from the device into a a RoborockBase.
 
         This is used by subclasses that want to override the class
@@ -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())