chore: change a bit of how traits work

Author: Lash-L

roborock/containers.py

@@ -1184,3 +1184,12 @@ class DyadSndState(RoborockBase):
 @dataclass
 class DyadOtaNfo(RoborockBase):
     mqttOtaData: dict
+
+
+@dataclass
+class DndActions(RoborockBase):
+    dry: int | None = None
+    dust: int | None = None
+    led: int | None = None
+    resume: int | None = None
+    vol: int | None = None

roborock/device_trait.py

@@ -1,16 +1,14 @@
-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
+from .containers import DeviceFeatures, 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
@@ -22,8 +20,12 @@ def __init__(self, send_command: Callable[..., Awaitable[None]]):
     def supported(cls, features: DeviceFeatures) -> bool:
         raise NotImplementedError
 
+    @abstractmethod
+    def from_dict(cls, data: dict) -> bool:
+        raise NotImplementedError
+
     def on_message(self, data: dict) -> None:
-        self.status = self._status_type.from_dict(data)
+        self.status = self.from_dict(data)
         for callback in self.subscriptions:
             callback(self.status)
 
@@ -36,44 +38,20 @@ 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)
+# 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

@@ -0,0 +1 @@
+from .dnd import Dnd

roborock/device_traits/dnd.py

@@ -0,0 +1,56 @@
+import datetime
+from collections.abc import Awaitable, Callable
+
+from roborock import DeviceFeatures, DndActions, RoborockCommand
+from roborock.device_trait import DeviceTrait
+
+
+class Dnd(DeviceTrait):
+    handle_command: RoborockCommand = RoborockCommand.GET_DND_TIMER
+
+    def __init__(self, send_command: Callable[..., Awaitable[None]]):
+        self.start_hour: int | None = None
+        self.start_minute: int | None = None
+        self.end_hour: int | None = None
+        self.end_minute: int | None = None
+        self.enabled: bool | None = None
+        self.start_time: datetime.time | None = None
+        self.end_time: datetime.time | None = None
+        self.actions: DndActions | None = None
+        super().__init__(send_command)
+
+    def from_dict(self, dnd_dict: dict):
+        self.start_hour = dnd_dict.get("start_hour")
+        self.start_minute = dnd_dict.get("start_minute")
+        self.end_hour = dnd_dict.get("end_hour")
+        self.end_minute = dnd_dict.get("end_minute")
+        self.enabled = bool(dnd_dict.get("enabled"))
+        self.actions = DndActions.from_dict(dnd_dict.get("actions"))
+        self.start_time = (
+            datetime.time(hour=self.start_hour, minute=self.start_minute)
+            if self.start_hour is not None and self.start_minute is not None
+            else None
+        )
+        self.end_time = (
+            datetime.time(hour=self.end_hour, minute=self.end_minute)
+            if self.end_hour is not None and self.end_minute is not None
+            else None
+        )
+
+    def to_dict(self) -> dict:
+        return {}
+
+    @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.enabled and not enabled:
+            await self.send_command(RoborockCommand.CLOSE_DND_TIMER)
+        else:
+            start = start_time if start_time is not None else self.start_time
+            end = end_time if end_time is not None else self.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)

roborock/roborock_device.py

@@ -8,7 +8,8 @@
 
 from . import RoborockCommand
 from .containers import DeviceData, ModelStatus, S7MaxVStatus, Status, UserData
-from .device_trait import ConsumableTrait, DeviceTrait, DndTrait
+from .device_trait import DeviceTrait
+from .device_traits import Dnd
 from .mqtt.roborock_session import MqttParams, RoborockMqttSession
 from .protocol import MessageParser, Utils, md5hex
 from .roborock_message import RoborockMessage, RoborockMessageProtocol
@@ -36,8 +37,8 @@ def __init__(self, user_data: UserData, device_info: DeviceData):
         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._dnd_trait: Dnd | None = self.determine_supported_traits(Dnd)
+        # self._consumable_trait: ConsumableTrait | None = self.determine_supported_traits(ConsumableTrait)
         self._status_type: type[Status] = ModelStatus.get(device_info.model, S7MaxVStatus)
         # TODO: One per client EVER
         self.session = RoborockMqttSession(
@@ -153,11 +154,3 @@ def on_message(self, message_bytes: bytes):
 
             # 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