chore: some misc changes

Author: Lash-L

roborock/device_trait.py

@@ -0,0 +1,79 @@
+import datetime
+from abc import ABC, abstractmethod
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+
+from . import RoborockCommand
+from .containers import Consumable, DeviceFeatures, DnDTimer, RoborockBase
+
+
+@dataclass
+class DeviceTrait(ABC):
+    handle_command: RoborockCommand
+    _status_type: type[RoborockBase] = RoborockBase
+
+    def __init__(self, send_command: Callable[..., Awaitable[None]]):
+        self.send_command = send_command
+        self.status: RoborockBase | None = None
+        self.subscriptions = []
+
+    @classmethod
+    @abstractmethod
+    def supported(cls, features: DeviceFeatures) -> bool:
+        raise NotImplementedError
+
+    def on_message(self, data: dict) -> None:
+        self.status = self._status_type.from_dict(data)
+        for callback in self.subscriptions:
+            callback(self.status)
+
+    def subscribe(self, callable: Callable):
+        # Maybe needs to handle async too?
+        self.subscriptions.append(callable)
+
+    @abstractmethod
+    def get(self):
+        raise NotImplementedError
+
+
+class DndTrait(DeviceTrait):
+    handle_command: RoborockCommand = RoborockCommand.GET_DND_TIMER
+    _status_type: type[DnDTimer] = DnDTimer
+    status: DnDTimer
+
+    def __init__(self, send_command: Callable[..., Awaitable[None]]):
+        super().__init__(send_command)
+
+    @classmethod
+    def supported(cls, features: DeviceFeatures) -> bool:
+        return features.is_support_custom_dnd
+
+    async def update_dnd(self, enabled: bool, start_time: datetime.time, end_time: datetime.time) -> None:
+        if self.status.enabled and not enabled:
+            await self.send_command(RoborockCommand.CLOSE_DND_TIMER)
+        else:
+            start = start_time if start_time is not None else self.status.start_time
+            end = end_time if end_time is not None else self.status.end_time
+            await self.send_command(RoborockCommand.SET_DND_TIMER, [start.hour, start.minute, end.hour, end.minute])
+
+    async def get(self) -> None:
+        await self.send_command(RoborockCommand.GET_DND_TIMER)
+
+
+class ConsumableTrait(DeviceTrait):
+    handle_command = RoborockCommand.GET_CONSUMABLE
+    _status_type: type[Consumable] = DnDTimer
+    status: Consumable
+
+    def __init__(self, send_command: Callable[..., Awaitable[None]]):
+        super().__init__(send_command)
+
+    @classmethod
+    def supported(cls, features: DeviceFeatures) -> bool:
+        return True
+
+    async def reset_consumable(self, consumable: str) -> None:
+        await self.send_command(RoborockCommand.RESET_CONSUMABLE, [consumable])
+
+    async def get(self) -> None:
+        await self.send_command(RoborockCommand.GET_CONSUMABLE)

roborock/device_traits/__init__.py

@@ -89,7 +89,8 @@ async def unsubscribe(device: DeviceData):
                     device_id = message.topic.value.split("/")[-1]
                     device = device_map[device_id]
                     message = MessageParser.parse(message.payload, device.device.local_key)
-                    callbacks[device_id](message)
+                    for m in message[0]:
+                        callbacks[device_id](m)
                 except Exception:
                     ...
 

roborock/mqtt_manager.py

@@ -6,7 +6,8 @@
 import time
 
 from . import RoborockCommand
-from .containers import DeviceData, UserData
+from .containers import DeviceData, ModelStatus, S7MaxVStatus, Status, UserData
+from .device_trait import ConsumableTrait, DeviceTrait, DndTrait
 from .mqtt_manager import RoborockMqttManager
 from .protocol import MessageParser, Utils
 from .roborock_message import RoborockMessage, RoborockMessageProtocol
@@ -25,35 +26,34 @@ def __init__(self, user_data: UserData, device_info: DeviceData):
         self._local_endpoint = "abc"
         self._nonce = secrets.token_bytes(16)
         self.manager = RoborockMqttManager()
-        self.update_commands = self.determine_supported_commands()
-
-    def determine_supported_commands(self):
-        # All devices support these
-        supported_commands = {
-            RoborockCommand.GET_CONSUMABLE,
-            RoborockCommand.GET_STATUS,
-            RoborockCommand.GET_CLEAN_SUMMARY,
-        }
-        # Get what features we can from the feature_set info.
-
-        # If a command is not described in feature_set, we should just add it anyways and then let it fail on the first call and remove it.
-        robot_new_features = int(self.device_info.device.feature_set)
-        new_feature_info_str = self.device_info.device.new_feature_set
-        if 33554432 & int(robot_new_features):
-            supported_commands.add(RoborockCommand.GET_DUST_COLLECTION_MODE)
-        if 2 & int(new_feature_info_str[-8:], 16):
-            # TODO: May not be needed as i think this can just be found in Status, but just POC
-            supported_commands.add(RoborockCommand.APP_GET_CLEAN_ESTIMATE_INFO)
-        return supported_commands
+        self._message_id_types: dict[int, DeviceTrait] = {}
+        self._command_to_trait = {}
+        self._all_supported_traits = []
+        self._dnd_trait: DndTrait | None = self.determine_supported_traits(DndTrait)
+        self._consumable_trait: ConsumableTrait | None = self.determine_supported_traits(ConsumableTrait)
+        self._status_type: type[Status] = ModelStatus.get(device_info.model, S7MaxVStatus)
+
+    def determine_supported_traits(self, trait: type[DeviceTrait]):
+        def _send_command(
+            method: RoborockCommand | str, params: list | dict | int | None = None, use_cloud: bool = True
+        ):
+            return self.send_message(method, params, use_cloud)
+
+        if trait.supported(self.device_info.device_features):
+            trait_instance = trait(_send_command)
+            self._all_supported_traits.append(trait(_send_command))
+            self._command_to_trait[trait.handle_command] = trait_instance
+            return trait_instance
+        return None
 
     async def connect(self):
         """Connect via MQTT and Local if possible."""
         await self.manager.subscribe(self.user_data, self.device_info, self.on_message)
         await self.update()
 
     async def update(self):
-        for cmd in self.update_commands:
-            await self.send_message(method=cmd)
+        for trait in self._all_supported_traits:
+            await trait.get()
 
     def _get_payload(
         self,
@@ -91,7 +91,9 @@ async def send_message(
         request_id, timestamp, payload = self._get_payload(method, params, True, use_cloud)
         request_protocol = RoborockMessageProtocol.RPC_REQUEST
         roborock_message = RoborockMessage(timestamp=timestamp, protocol=request_protocol, payload=payload)
-
+        if request_id in self._message_id_types:
+            raise Exception("Duplicate id!")
+        self._message_id_types[request_id] = self._command_to_trait[method]
         local_key = self.device_info.device.local_key
         msg = MessageParser.build(roborock_message, local_key, False)
         if use_cloud:
@@ -101,6 +103,19 @@ async def send_message(
             pass
 
     def on_message(self, message: RoborockMessage):
+        message_payload = message.get_payload()
+        message_id = message.get_request_id()
+        for data_point_number, data_point in message_payload.get("dps").items():
+            if data_point_number == "102":
+                data_point_response = json.loads(data_point)
+                result = data_point_response.get("result")
+                if isinstance(result, list) and len(result) == 1:
+                    result = result[0]
+                if result and (trait := self._message_id_types.get(message_id)) is not None:
+                    trait.on_message(result)
+                if (error := result.get("error")) is not None:
+                    print(error)
+        print()
         # If message is command not supported - remove from self.update_commands
 
         # If message is an error - log it?
@@ -115,3 +130,11 @@ def on_message(self, message: RoborockMessage):
 
         # This should also probably be split with on_cloud_message and on_local_message.
         print(message)
+
+    @property
+    def dnd(self) -> DndTrait | None:
+        return self._dnd_trait
+
+    @property
+    def consumable(self) -> ConsumableTrait | None:
+        return self._consumable_trait

roborock/roborock_device.py

@@ -161,11 +161,16 @@ class RoborockMessage:
     random: int = field(default_factory=lambda: get_next_int(10000, 99999))
     timestamp: int = field(default_factory=lambda: math.floor(time.time()))
     message_retry: MessageRetry | None = None
+    _parsed_payload: dict | None = None
+
+    def get_payload(self) -> dict | None:
+        if self.payload and not self._parsed_payload:
+            self._parsed_payload = json.loads(self.payload.decode())
+        return self._parsed_payload
 
     def get_request_id(self) -> int | None:
-        if self.payload:
-            payload = json.loads(self.payload.decode())
-            for data_point_number, data_point in payload.get("dps").items():
+        if self._parsed_payload:
+            for data_point_number, data_point in self._parsed_payload.get("dps").items():
                 if data_point_number in ["101", "102"]:
                     data_point_response = json.loads(data_point)
                     return data_point_response.get("id")
@@ -180,19 +185,17 @@ def get_method(self) -> str | None:
         if self.message_retry:
             return self.message_retry.method
         protocol = self.protocol
-        if self.payload and protocol in [4, 5, 101, 102]:
-            payload = json.loads(self.payload.decode())
-            for data_point_number, data_point in payload.get("dps").items():
+        if self._parsed_payload and protocol in [4, 5, 101, 102]:
+            for data_point_number, data_point in self._parsed_payload.get("dps").items():
                 if data_point_number in ["101", "102"]:
                     data_point_response = json.loads(data_point)
                     return data_point_response.get("method")
         return None
 
     def get_params(self) -> list | dict | None:
         protocol = self.protocol
-        if self.payload and protocol in [4, 101, 102]:
-            payload = json.loads(self.payload.decode())
-            for data_point_number, data_point in payload.get("dps").items():
+        if self._parsed_payload and protocol in [4, 101, 102]:
+            for data_point_number, data_point in self._parsed_payload.get("dps").items():
                 if data_point_number in ["101", "102"]:
                     data_point_response = json.loads(data_point)
                     return data_point_response.get("params")