roborock.devices.traits.v1.map_content

 1"""Trait for fetching the map content from Roborock devices."""
 2
 3import logging
 4from dataclasses import dataclass
 5
 6from vacuum_map_parser_base.map_data import MapData
 7
 8from roborock.data import RoborockBase
 9from roborock.devices.traits.v1 import common
10from roborock.map.map_parser import MapParser, MapParserConfig
11from roborock.roborock_typing import RoborockCommand
12
13_LOGGER = logging.getLogger(__name__)
14
15_TRUNCATE_LENGTH = 20
16
17
18@dataclass
19class MapContent(RoborockBase):
20    """Dataclass representing map content."""
21
22    image_content: bytes | None = None
23    """The rendered image of the map in PNG format."""
24
25    map_data: MapData | None = None
26    """The parsed map data which contains metadata for points on the map."""
27
28    raw_api_response: bytes | None = None
29    """The raw bytes of the map data from the API for caching for future use.
30
31    This should be treated as an opaque blob used only internally by this library
32    to re-parse the map data when needed.
33    """
34
35    def __repr__(self) -> str:
36        """Return a string representation of the MapContent."""
37        img = self.image_content
38        if self.image_content and len(self.image_content) > _TRUNCATE_LENGTH:
39            img = self.image_content[: _TRUNCATE_LENGTH - 3] + b"..."
40        return f"MapContent(image_content={img!r}, map_data={self.map_data!r})"
41
42
43class MapContentConverter(common.V1TraitDataConverter):
44    """Convert map response data to MapContent."""
45
46    def __init__(self, map_parser: MapParser) -> None:
47        """Initialize MapContentConverter."""
48        self._map_parser = map_parser
49
50    def convert(self, response: common.V1ResponseData) -> MapContent:
51        """Parse the response from the device into a MapContent instance."""
52        if not isinstance(response, bytes):
53            raise ValueError(f"Unexpected MapContentTrait response format: {type(response)}")
54        return self.parse_map_content(response)
55
56    def parse_map_content(self, response: bytes) -> MapContent:
57        """Parse the map content from raw bytes.
58
59        This method is exposed so that cached map data can be parsed without
60        needing to go through the RPC channel.
61
62        Args:
63            response: The raw bytes of the map data from the API.
64
65        Returns:
66            MapContent: The parsed map content.
67
68        Raises:
69            RoborockException: If the map data cannot be parsed.
70        """
71        parsed_data = self._map_parser.parse(response)
72        if parsed_data is None:
73            raise ValueError("Failed to parse map data")
74
75        return MapContent(
76            image_content=parsed_data.image_content,
77            map_data=parsed_data.map_data,
78            raw_api_response=response,
79        )
80
81
82@common.map_rpc_channel
83class MapContentTrait(MapContent, common.V1TraitMixin):
84    """Trait for fetching the map content."""
85
86    command = RoborockCommand.GET_MAP_V1
87    converter: MapContentConverter
88
89    def __init__(self, map_parser_config: MapParserConfig | None = None) -> None:
90        """Initialize MapContentTrait."""
91        super().__init__()
92        self.converter = MapContentConverter(MapParser(map_parser_config or MapParserConfig()))