roborock-cli 0.1.1__py3-none-any.whl

roborock_cli/_vendor/roborock/devices/traits/v1/maps.py

@@ -0,0 +1,80 @@
+"""Trait for managing maps and room mappings on Roborock devices.
+
+New datatypes are introduced here to manage the additional information associated
+with maps and rooms, such as map names and room names. These override the
+base container datatypes to add additional fields.
+"""
+
+import logging
+from typing import Self
+
+from roborock_cli._vendor.roborock.data import MultiMapsList, MultiMapsListMapInfo
+from roborock_cli._vendor.roborock.devices.traits.v1 import common
+from roborock_cli._vendor.roborock.roborock_typing import RoborockCommand
+
+from .status import StatusTrait
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@common.mqtt_rpc_channel
+class MapsTrait(MultiMapsList, common.V1TraitMixin):
+    """Trait for managing the maps of Roborock devices.
+
+    A device may have multiple maps, each identified by a unique map_flag.
+    Each map can have multiple rooms associated with it, in a `RoomMapping`.
+
+    The MapsTrait depends on the StatusTrait to determine the currently active
+    map. It is the responsibility of the caller to ensure that the StatusTrait
+    is up to date before using this trait. However, there is a possibility of
+    races if another client changes the current map between the time the
+    StatusTrait is refreshed and when the MapsTrait is used. This is mitigated
+    by the fact that the map list is unlikely to change frequently, and the
+    current map is only changed when the user explicitly switches maps.
+    """
+
+    command = RoborockCommand.GET_MULTI_MAPS_LIST
+
+    def __init__(self, status_trait: StatusTrait) -> None:
+        """Initialize the MapsTrait.
+
+        We keep track of the StatusTrait to ensure we have the latest
+        status information when dealing with maps.
+        """
+        super().__init__()
+        self._status_trait = status_trait
+
+    @property
+    def current_map(self) -> int | None:
+        """Returns the currently active map (map_flag), if available."""
+        return self._status_trait.current_map
+
+    @property
+    def current_map_info(self) -> MultiMapsListMapInfo | None:
+        """Returns the currently active map info, if available."""
+        if (current_map := self.current_map) is None or self.map_info is None:
+            return None
+        for map_info in self.map_info:
+            if map_info.map_flag == current_map:
+                return map_info
+        return None
+
+    async def set_current_map(self, map_flag: int) -> None:
+        """Update the current map of the device by it's map_flag id."""
+        await self.rpc_channel.send_command(RoborockCommand.LOAD_MULTI_MAP, params=[map_flag])
+        # Refresh our status to make sure it reflects the new map
+        await self._status_trait.refresh()
+
+    def _parse_response(self, response: common.V1ResponseData) -> Self:
+        """Parse the response from the device into a MapsTrait instance.
+
+        This overrides the base implementation to handle the specific
+        response format for the multi maps list. This is needed because we have
+        a custom constructor that requires the StatusTrait.
+        """
+        if not isinstance(response, list):
+            raise ValueError(f"Unexpected MapsTrait response format: {response!r}")
+        response = response[0]
+        if not isinstance(response, dict):
+            raise ValueError(f"Unexpected MapsTrait response format: {response!r}")
+        return MultiMapsList.from_dict(response)

roborock_cli/_vendor/roborock/devices/traits/v1/network_info.py

@@ -0,0 +1,55 @@
+"""Trait for device network information."""
+
+from __future__ import annotations
+
+import logging
+
+from roborock_cli._vendor.roborock.data import NetworkInfo
+from roborock_cli._vendor.roborock.devices.cache import DeviceCache
+from roborock_cli._vendor.roborock.devices.traits.v1 import common
+from roborock_cli._vendor.roborock.roborock_typing import RoborockCommand
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class NetworkInfoTrait(NetworkInfo, common.V1TraitMixin):
+    """Trait for device network information.
+
+    This trait will always prefer reading from the cache if available. This
+    information is usually already fetched when creating the device local
+    connection, so reading from the cache avoids an unnecessary RPC call.
+    However, we have the fallback to reading from the device if the cache is
+    not populated for some reason.
+    """
+
+    command = RoborockCommand.GET_NETWORK_INFO
+
+    def __init__(self, device_uid: str, device_cache: DeviceCache) -> None:  # pylint: disable=super-init-not-called
+        """Initialize the trait."""
+        self._device_uid = device_uid
+        self._device_cache = device_cache
+        self.ip = ""
+
+    async def refresh(self) -> None:
+        """Refresh the network info from the cache."""
+
+        device_cache_data = await self._device_cache.get()
+        if device_cache_data.network_info:
+            _LOGGER.debug("Using cached network info for device %s", self._device_uid)
+            self._update_trait_values(device_cache_data.network_info)
+            return
+
+        # Load from device if not in cache
+        _LOGGER.debug("No cached network info for device %s, fetching from device", self._device_uid)
+        await super().refresh()
+
+        # Update the cache with the new network info
+        device_cache_data = await self._device_cache.get()
+        device_cache_data.network_info = self
+        await self._device_cache.set(device_cache_data)
+
+    def _parse_response(self, response: common.V1ResponseData) -> NetworkInfo:
+        """Parse the response from the device into a NetworkInfo."""
+        if not isinstance(response, dict):
+            raise ValueError(f"Unexpected NetworkInfoTrait response format: {response!r}")
+        return NetworkInfo.from_dict(response)

roborock_cli/_vendor/roborock/devices/traits/v1/rooms.py

@@ -0,0 +1,105 @@
+"""Trait for managing room mappings on Roborock devices."""
+
+import logging
+from dataclasses import dataclass
+
+from roborock_cli._vendor.roborock.data import HomeData, HomeDataRoom, NamedRoomMapping, RoborockBase
+from roborock_cli._vendor.roborock.devices.traits.v1 import common
+from roborock_cli._vendor.roborock.roborock_typing import RoborockCommand
+from roborock_cli._vendor.roborock.web_api import UserWebApiClient
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class Rooms(RoborockBase):
+    """Dataclass representing a collection of room mappings."""
+
+    rooms: list[NamedRoomMapping] | None = None
+    """List of room mappings."""
+
+    @property
+    def room_map(self) -> dict[int, NamedRoomMapping]:
+        """Returns a mapping of segment_id to NamedRoomMapping."""
+        if self.rooms is None:
+            return {}
+        return {room.segment_id: room for room in self.rooms}
+
+
+class RoomsTrait(Rooms, common.V1TraitMixin):
+    """Trait for managing the room mappings of Roborock devices."""
+
+    command = RoborockCommand.GET_ROOM_MAPPING
+
+    def __init__(self, home_data: HomeData, web_api: UserWebApiClient) -> None:
+        """Initialize the RoomsTrait."""
+        super().__init__()
+        self._home_data = home_data
+        self._web_api = web_api
+        self._discovered_iot_ids: set[str] = set()
+
+    async def refresh(self) -> None:
+        """Refresh room mappings and backfill unknown room names from the web API."""
+        response = await self.rpc_channel.send_command(self.command)
+        if not isinstance(response, list):
+            raise ValueError(f"Unexpected RoomsTrait response format: {response!r}")
+
+        segment_map = _extract_segment_map(response)
+        # Track all iot ids seen before. Refresh the room list when new ids are found.
+        new_iot_ids = set(segment_map.values()) - set(self._home_data.rooms_map.keys())
+        if new_iot_ids - self._discovered_iot_ids:
+            _LOGGER.debug("Refreshing room list to discover new room names")
+            if updated_rooms := await self._refresh_rooms():
+                _LOGGER.debug("Updating rooms: %s", list(updated_rooms))
+                self._home_data.rooms = updated_rooms
+            self._discovered_iot_ids.update(new_iot_ids)
+
+        new_data = self._parse_rooms(segment_map, self._home_data.rooms_name_map)
+        self._update_trait_values(new_data)
+        _LOGGER.debug("Refreshed %s: %s", self.__class__.__name__, new_data)
+
+    @staticmethod
+    def _parse_rooms(
+        segment_map: dict[int, str],
+        name_map: dict[str, str],
+    ) -> Rooms:
+        """Parse the response from the device into a list of NamedRoomMapping."""
+        return Rooms(
+            rooms=[
+                NamedRoomMapping(segment_id=segment_id, iot_id=iot_id, raw_name=name_map.get(iot_id))
+                for segment_id, iot_id in segment_map.items()
+            ]
+        )
+
+    async def _refresh_rooms(self) -> list[HomeDataRoom]:
+        """Fetch the latest rooms from the web API."""
+        try:
+            return await self._web_api.get_rooms()
+        except Exception:
+            _LOGGER.debug("Failed to fetch rooms from web API", exc_info=True)
+            return []
+
+
+def _extract_segment_map(response: list) -> dict[int, str]:
+    """Extract a segment_id -> iot_id mapping from the response.
+
+    The response format can be either a flat list of [segment_id, iot_id] or a
+    list of lists, where each inner list is a pair of [segment_id, iot_id]. This
+    function normalizes the response into a dict of segment_id to iot_id.
+
+    NOTE: We currently only partial samples of the room mapping formats, so
+    improving test coverage with samples from a real device with this format
+    would be helpful.
+    """
+    if len(response) == 2 and not isinstance(response[0], list):
+        segment_id, iot_id = response[0], response[1]
+        return {segment_id: str(iot_id)}
+
+    segment_map: dict[int, str] = {}
+    for part in response:
+        if not isinstance(part, list) or len(part) < 2:
+            _LOGGER.warning("Unexpected room mapping entry format: %r", part)
+            continue
+        segment_id, iot_id = part[0], part[1]
+        segment_map[segment_id] = str(iot_id)
+    return segment_map

roborock_cli/_vendor/roborock/devices/traits/v1/routines.py

@@ -0,0 +1,26 @@
+"""Routines trait for V1 devices."""
+
+from roborock_cli._vendor.roborock.data.containers import HomeDataScene
+from roborock_cli._vendor.roborock.web_api import UserWebApiClient
+
+
+class RoutinesTrait:
+    """Trait for interacting with routines."""
+
+    def __init__(self, device_id: str, web_api: UserWebApiClient) -> None:
+        """Initialize the routines trait."""
+        self._device_id = device_id
+        self._web_api = web_api
+
+    async def get_routines(self) -> list[HomeDataScene]:
+        """Get available routines."""
+        return await self._web_api.get_routines(self._device_id)
+
+    async def execute_routine(self, routine_id: int) -> None:
+        """Execute a routine by its ID.
+
+        Technically, routines are per-device, but the API does not
+        require the device ID to execute them. This can execute a
+        routine for any device but it is exposed here for convenience.
+        """
+        await self._web_api.execute_routine(routine_id)

roborock_cli/_vendor/roborock/devices/traits/v1/smart_wash_params.py

@@ -0,0 +1,13 @@
+"""Trait for smart wash parameters."""
+
+from roborock_cli._vendor.roborock.data import SmartWashParams
+from roborock_cli._vendor.roborock.device_features import is_wash_n_fill_dock
+from roborock_cli._vendor.roborock.devices.traits.v1 import common
+from roborock_cli._vendor.roborock.roborock_typing import RoborockCommand
+
+
+class SmartWashParamsTrait(SmartWashParams, common.V1TraitMixin):
+    """Trait for smart wash parameters."""
+
+    command = RoborockCommand.GET_SMART_WASH_PARAMS
+    requires_dock_type = is_wash_n_fill_dock

roborock_cli/_vendor/roborock/devices/traits/v1/status.py

@@ -0,0 +1,101 @@
+from functools import cached_property
+from typing import Self
+
+from roborock_cli._vendor.roborock import (
+    CleanRoutes,
+    StatusV2,
+    VacuumModes,
+    WaterModes,
+    get_clean_modes,
+    get_clean_routes,
+    get_water_mode_mapping,
+    get_water_modes,
+)
+from roborock_cli._vendor.roborock.roborock_typing import RoborockCommand
+
+from . import common
+from .device_features import DeviceFeaturesTrait
+
+
+class StatusTrait(StatusV2, common.V1TraitMixin):
+    """Trait for managing the status of Roborock devices.
+
+    The StatusTrait gives you the access to the state of a Roborock vacuum.
+    The various attribute options on state change per each device.
+    Values like fan speed, mop mode, etc. have different options for every device
+    and are dynamically determined.
+
+    Usage:
+        Before accessing status properties, you should call `refresh()` to fetch
+        the latest data from the device. You must pass in the device feature trait
+        to this trait so that the dynamic attributes can be pre-determined.
+
+    The current dynamic attributes are:
+    - Fan Speed
+    - Water Mode
+    - Mop Route
+
+    You should call the _options() version of the attribute to know which are supported for your device
+    (i.e. fan_speed_options())
+    Then you can call the _mapping to convert an int value to the actual Enum. (i.e. fan_speed_mapping())
+    You can call the _name property to get the str value of the enum. (i.e. fan_speed_name)
+
+    """
+
+    command = RoborockCommand.GET_STATUS
+
+    def __init__(self, device_feature_trait: DeviceFeaturesTrait, region: str | None = None) -> None:
+        """Initialize the StatusTrait."""
+        super().__init__()
+        self._device_features_trait = device_feature_trait
+        self._region = region
+
+    @cached_property
+    def fan_speed_options(self) -> list[VacuumModes]:
+        return get_clean_modes(self._device_features_trait)
+
+    @cached_property
+    def fan_speed_mapping(self) -> dict[int, str]:
+        return {fan.code: fan.value for fan in self.fan_speed_options}
+
+    @cached_property
+    def water_mode_options(self) -> list[WaterModes]:
+        return get_water_modes(self._device_features_trait)
+
+    @cached_property
+    def water_mode_mapping(self) -> dict[int, str]:
+        return get_water_mode_mapping(self._device_features_trait)
+
+    @cached_property
+    def mop_route_options(self) -> list[CleanRoutes]:
+        return get_clean_routes(self._device_features_trait, self._region or "us")
+
+    @cached_property
+    def mop_route_mapping(self) -> dict[int, str]:
+        return {route.code: route.value for route in self.mop_route_options}
+
+    @property
+    def fan_speed_name(self) -> str | None:
+        if self.fan_power is None:
+            return None
+        return self.fan_speed_mapping.get(self.fan_power)
+
+    @property
+    def water_mode_name(self) -> str | None:
+        if self.water_box_mode is None:
+            return None
+        return self.water_mode_mapping.get(self.water_box_mode)
+
+    @property
+    def mop_route_name(self) -> str | None:
+        if self.mop_mode is None:
+            return None
+        return self.mop_route_mapping.get(self.mop_mode)
+
+    def _parse_response(self, response: common.V1ResponseData) -> Self:
+        """Parse the response from the device into a StatusV2-based status object."""
+        if isinstance(response, list):
+            response = response[0]
+        if isinstance(response, dict):
+            return StatusV2.from_dict(response)
+        raise ValueError(f"Unexpected status format: {response!r}")

roborock_cli/_vendor/roborock/devices/traits/v1/valley_electricity_timer.py

@@ -0,0 +1,44 @@
+from roborock_cli._vendor.roborock.data import ValleyElectricityTimer
+from roborock_cli._vendor.roborock.devices.traits.v1 import common
+from roborock_cli._vendor.roborock.roborock_typing import RoborockCommand
+
+_ENABLED_PARAM = "enabled"
+
+
+class ValleyElectricityTimerTrait(ValleyElectricityTimer, common.V1TraitMixin, common.RoborockSwitchBase):
+    """Trait for managing Valley Electricity Timer settings on Roborock devices."""
+
+    command = RoborockCommand.GET_VALLEY_ELECTRICITY_TIMER
+    requires_feature = "is_supported_valley_electricity"
+
+    @property
+    def is_on(self) -> bool:
+        """Return whether the Valley Electricity Timer is enabled."""
+        return self.enabled == 1
+
+    async def set_timer(self, timer: ValleyElectricityTimer) -> None:
+        """Set the Valley Electricity Timer settings of the device."""
+        await self.rpc_channel.send_command(RoborockCommand.SET_VALLEY_ELECTRICITY_TIMER, params=timer.as_list())
+        await self.refresh()
+
+    async def clear_timer(self) -> None:
+        """Clear the Valley Electricity Timer settings of the device."""
+        await self.rpc_channel.send_command(RoborockCommand.CLOSE_VALLEY_ELECTRICITY_TIMER)
+        await self.refresh()
+
+    async def enable(self) -> None:
+        """Enable the Valley Electricity Timer settings of the device."""
+        await self.rpc_channel.send_command(
+            RoborockCommand.SET_VALLEY_ELECTRICITY_TIMER,
+            params=self.as_list(),
+        )
+        # Optimistic update to avoid an extra refresh
+        self.enabled = 1
+
+    async def disable(self) -> None:
+        """Disable the Valley Electricity Timer settings of the device."""
+        await self.rpc_channel.send_command(
+            RoborockCommand.CLOSE_VALLEY_ELECTRICITY_TIMER,
+        )
+        # Optimistic update to avoid an extra refresh
+        self.enabled = 0

roborock_cli/_vendor/roborock/devices/traits/v1/volume.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass, field
+
+from roborock_cli._vendor.roborock.devices.traits.v1 import common
+from roborock_cli._vendor.roborock.roborock_typing import RoborockCommand
+
+# TODO: This is currently the pattern for holding all the commands that hold a
+# single value, but it still seems too verbose. Maybe we can generate these
+# dynamically or somehow make them less code.
+
+
+@dataclass
+class SoundVolume(common.RoborockValueBase):
+    """Dataclass for sound volume."""
+
+    volume: int | None = field(default=None, metadata={"roborock_value": True})
+    """Sound volume level (0-100)."""
+
+
+class SoundVolumeTrait(SoundVolume, common.V1TraitMixin):
+    """Trait for controlling the sound volume of a Roborock device."""
+
+    command = RoborockCommand.GET_SOUND_VOLUME
+
+    async def set_volume(self, volume: int) -> None:
+        """Set the sound volume of the device."""
+        await self.rpc_channel.send_command(RoborockCommand.CHANGE_SOUND_VOLUME, params=[volume])
+        self.volume = volume

roborock_cli/_vendor/roborock/devices/traits/v1/wash_towel_mode.py

@@ -0,0 +1,13 @@
+"""Trait for wash towel mode."""
+
+from roborock_cli._vendor.roborock.data import WashTowelMode
+from roborock_cli._vendor.roborock.device_features import is_wash_n_fill_dock
+from roborock_cli._vendor.roborock.devices.traits.v1 import common
+from roborock_cli._vendor.roborock.roborock_typing import RoborockCommand
+
+
+class WashTowelModeTrait(WashTowelMode, common.V1TraitMixin):
+    """Trait for wash towel mode."""
+
+    command = RoborockCommand.GET_WASH_TOWEL_MODE
+    requires_dock_type = is_wash_n_fill_dock

roborock_cli/_vendor/roborock/devices/transport/__init__.py

@@ -0,0 +1,8 @@
+"""Module for handling network connections to Roborock devices.
+
+This is used internally by the device manager for creating connections to
+Roborock devices. These modules contain common code, not specific to a
+particular device or application level protocol.
+"""
+
+__all__: list[str] = []

roborock_cli/_vendor/roborock/devices/transport/channel.py

@@ -0,0 +1,32 @@
+"""Low-level interface for connections to Roborock devices."""
+
+import logging
+from collections.abc import Callable
+from typing import Protocol
+
+from roborock_cli._vendor.roborock.roborock_message import RoborockMessage
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class Channel(Protocol):
+    """A generic channel for establishing a connection with a Roborock device.
+
+    Individual channel implementations have their own methods for speaking to
+    the device that hide some of the protocol specific complexity, but they
+    are still specialized for the device type and protocol.
+    """
+
+    @property
+    def is_connected(self) -> bool:
+        """Return true if the channel is connected."""
+        ...
+
+    @property
+    def is_local_connected(self) -> bool:
+        """Return true if the channel is connected locally."""
+        ...
+
+    async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callable[[], None]:
+        """Subscribe to messages from the device."""
+        ...