Merge branch 'main' into dynamic_clean_modes

Author: Lash-L

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "python-roborock"
-version = "2.33.0"
+version = "2.35.0"
 description = "A package to control Roborock vacuums."
 authors = ["humbertogontijo <humbertogontijo@users.noreply.github.com>"]
 license = "GPL-3.0-only"

pyproject.toml

@@ -3,9 +3,7 @@
 from __future__ import annotations
 
 import asyncio
-import base64
 import logging
-import secrets
 import time
 from abc import ABC, abstractmethod
 from typing import Any
@@ -37,14 +35,11 @@ class RoborockClient(ABC):
     def __init__(self, device_info: DeviceData) -> None:
         """Initialize RoborockClient."""
         self.device_info = device_info
-        self._nonce = secrets.token_bytes(16)
         self._waiting_queue: dict[int, RoborockFuture] = {}
         self._last_device_msg_in = time.monotonic()
         self._last_disconnection = time.monotonic()
         self.keep_alive = KEEPALIVE
-        self._diagnostic_data: dict[str, dict[str, Any]] = {
-            "misc_info": {"Nonce": base64.b64encode(self._nonce).decode("utf-8")}
-        }
+        self._diagnostic_data: dict[str, dict[str, Any]] = {}
         self.is_available: bool = True
 
     async def async_release(self) -> None:
@@ -116,7 +111,3 @@ def _async_response(self, request_id: int, protocol_id: int = 0) -> Any:
             request_id = new_id
         self._waiting_queue[request_id] = queue
         return asyncio.ensure_future(self._wait_response(request_id, queue))
-
-    @abstractmethod
-    async def send_message(self, roborock_message: RoborockMessage):
-        """Send a message to the Roborock device."""

roborock/api.py

@@ -96,6 +96,18 @@ def get_clean_modes(features: DeviceFeatures) -> list[VacuumModes]:
     if features.is_customized_clean_supported:
         modes.append(VacuumModes.CUSTOMIZED)
 
+def get_clean_modes(features: DeviceFeatures) -> list[CleanModes]:
+    """Get the valid clean modes for the device - also known as 'fan power' or 'suction mode'"""
+    modes = [CleanModes.QUIET, CleanModes.BALANCED, CleanModes.TURBO, CleanModes.MAX]
+    if features.is_max_plus_mode_supported or features.is_none_pure_clean_mop_with_max_plus:
+        # If the vacuum has max plus mode supported
+        modes.append(CleanModes.MAX_PLUS)
+    if features.is_pure_clean_mop_supported:
+        # If the vacuum is capable of 'pure mop clean' aka no vacuum
+        modes.append(CleanModes.OFF)
+    else:
+        # If not, we can add gentle
+        modes.append(CleanModes.GENTLE)
     return modes
 
 

roborock/clean_modes.py

@@ -163,7 +163,7 @@ async def _discover(ctx):
     if not cache_data:
         raise Exception("You need to login first")
     client = RoborockApiClient(cache_data.email)
-    home_data = await client.get_home_data(cache_data.user_data)
+    home_data = await client.get_home_data_v3(cache_data.user_data)
     cache_data.home_data = home_data
     context.update(cache_data)
     click.echo(f"Discovered devices {', '.join([device.name for device in home_data.get_all_devices()])}")

roborock/cli.py

@@ -0,0 +1,30 @@
+"""Thin wrapper around the MQTT channel for Roborock B01 devices."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from roborock.protocols.b01_protocol import (
+    CommandType,
+    ParamsType,
+    decode_rpc_response,
+    encode_mqtt_payload,
+)
+
+from .mqtt_channel import MqttChannel
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def send_decoded_command(
+    mqtt_channel: MqttChannel,
+    dps: int,
+    command: CommandType,
+    params: ParamsType,
+) -> dict[int, Any]:
+    """Send a command on the MQTT channel and get a decoded response."""
+    _LOGGER.debug("Sending MQTT command: %s", params)
+    roborock_message = encode_mqtt_payload(dps, command, params)
+    response = await mqtt_channel.send_message(roborock_message)
+    return decode_rpc_response(response)  # type: ignore[return-value]

roborock/devices/b01_channel.py

@@ -21,6 +21,7 @@
 from .cache import Cache, NoCache
 from .channel import Channel
 from .mqtt_channel import create_mqtt_channel
+from .traits.b01.props import B01PropsApi
 from .traits.dyad import DyadApi
 from .traits.status import StatusTrait
 from .traits.trait import Trait
@@ -45,6 +46,7 @@ class DeviceVersion(enum.StrEnum):
 
     V1 = "1.0"
     A01 = "A01"
+    B01 = "B01"
     UNKNOWN = "unknown"
 
 
@@ -120,7 +122,7 @@ def create_home_data_api(email: str, user_data: UserData) -> HomeDataApi:
     client = RoborockApiClient(email)
 
     async def home_data_api() -> HomeData:
-        return await client.get_home_data(user_data)
+        return await client.get_home_data_v3(user_data)
 
     return home_data_api
 
@@ -159,6 +161,9 @@ def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> Roborock
                         traits.append(ZeoApi(mqtt_channel))
                     case _:
                         raise NotImplementedError(f"Device {device.name} has unsupported category {product.category}")
+            case DeviceVersion.B01:
+                mqtt_channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
+                traits.append(B01PropsApi(mqtt_channel))
             case _:
                 raise NotImplementedError(f"Device {device.name} has unsupported version {device.pv}")
         return RoborockDevice(device, channel, traits)

roborock/devices/device_manager.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from roborock import RoborockB01Methods
+from roborock.roborock_message import RoborockB01Props
+
+from ...b01_channel import send_decoded_command
+from ...mqtt_channel import MqttChannel
+from ..trait import Trait
+
+_LOGGER = logging.getLogger(__name__)
+
+__all__ = [
+    "B01PropsApi",
+]
+
+
+class B01PropsApi(Trait):
+    """API for interacting with B01 devices."""
+
+    name = "B01_props"
+
+    def __init__(self, channel: MqttChannel) -> None:
+        """Initialize the B01Props API."""
+        self._channel = channel
+
+    async def query_values(self, props: list[RoborockB01Props]) -> dict[int, Any]:
+        """Query the device for the values of the given Dyad protocols."""
+        return await send_decoded_command(self._channel, dps=10000, command=RoborockB01Methods.GET_PROP, params=props)