chore: Overhaul new device trait interfaces

roborock/cli.py

@@ -45,6 +45,7 @@
 from roborock.devices.device import RoborockDevice
 from roborock.devices.device_manager import DeviceManager, create_device_manager, create_home_data_api
 from roborock.devices.traits import Trait
+from roborock.devices.traits.v1 import V1TraitMixin
 from roborock.protocol import MessageParser
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
 from roborock.web_api import RoborockApiClient
@@ -379,7 +380,7 @@ async def execute_scene(ctx, scene_id):
     await client.execute_scene(cache_data.user_data, scene_id)
 
 
-async def _v1_trait(context: RoborockContext, device_id: str, display_func: Callable[[], Trait]) -> Trait:
+async def _v1_trait(context: RoborockContext, device_id: str, display_func: Callable[[], V1TraitMixin]) -> Trait:
     device_manager = await context.get_device_manager()
     device = await device_manager.get_device(device_id)
     if device.v1_properties is None:

roborock/devices/device.py

@@ -38,15 +38,15 @@ def __init__(
         self,
         device_info: HomeDataDevice,
         channel: Channel,
-        traits: list[Trait],
+        trait: Trait,
     ) -> None:
         """Initialize the RoborockDevice.
 
         The device takes ownership of the channel for communication with the device.
         Use `connect()` to establish the connection, which will set up the appropriate
         protocol channel. Use `close()` to clean up all connections.
         """
-        TraitsMixin.__init__(self, traits)
+        TraitsMixin.__init__(self, trait)
         self._duid = device_info.duid
         self._name = device_info.name
         self._channel = channel

roborock/devices/device_manager.py

@@ -145,23 +145,20 @@ async def create_device_manager(
 
     def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice:
         channel: Channel
-        traits: list[Trait] = []
+        trait: Trait
         match device.pv:
             case DeviceVersion.V1:
-                v1_channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
-                channel = v1_channel
-                traits.extend(v1.create_v1_traits(product, v1_channel.rpc_channel))
+                channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
+                trait = v1.create(product, channel.rpc_channel)
             case DeviceVersion.A01:
-                mqtt_channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
-                channel = mqtt_channel
-                traits.extend(a01.create_a01_traits(product, mqtt_channel))
+                channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
+                trait = a01.create(product, channel)
             case DeviceVersion.B01:
-                mqtt_channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
-                channel = mqtt_channel
-                traits.extend(b01.create_b01_traits(mqtt_channel))
+                channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
+                trait = b01.create(channel)
             case _:
                 raise NotImplementedError(f"Device {device.name} has unsupported version {device.pv}")
-        return RoborockDevice(device, channel, traits)
+        return RoborockDevice(device, channel, trait)
 
     manager = DeviceManager(home_data_api, device_creator, mqtt_session=mqtt_session, cache=cache)
     await manager.discover_devices()

roborock/devices/traits/a01/__init__.py

@@ -50,12 +50,12 @@ async def set_value(self, protocol: RoborockZeoProtocol, value: Any) -> dict[Rob
         return await send_decoded_command(self._channel, params)
 
 
-def create_a01_traits(product: HomeDataProduct, mqtt_channel: MqttChannel) -> list[Trait]:
+def create(product: HomeDataProduct, mqtt_channel: MqttChannel) -> DyadApi | ZeoApi:
     """Create traits for A01 devices."""
     match product.category:
         case RoborockCategory.WET_DRY_VAC:
-            return [DyadApi(mqtt_channel)]
+            return DyadApi(mqtt_channel)
         case RoborockCategory.WASHING_MACHINE:
-            return [ZeoApi(mqtt_channel)]
+            return ZeoApi(mqtt_channel)
         case _:
             raise NotImplementedError(f"Unsupported category {product.category}")

roborock/devices/traits/b01/__init__.py

@@ -1,17 +1,31 @@
 """Traits for B01 devices."""
 
+from roborock import RoborockB01Methods
 from roborock.devices.b01_channel import send_decoded_command
 from roborock.devices.mqtt_channel import MqttChannel
 from roborock.devices.traits import Trait
-
-from .props import B01PropsApi
+from roborock.roborock_message import RoborockB01Props
 
 __init__ = [
     "create_b01_traits",
     "B01PropsApi",
 ]
 
 
-def create_b01_traits(channel: MqttChannel) -> list[Trait]:
+class B01PropsApi(Trait):
+    """API for interacting with B01 devices."""
+
+    def __init__(self, channel: MqttChannel) -> None:
+        """Initialize the B01Props API."""
+        self._channel = channel
+
+    async def query_values(self, props: list[RoborockB01Props]) -> None:
+        """Query the device for the values of the given Dyad protocols."""
+        await send_decoded_command(
+            self._channel, dps=10000, command=RoborockB01Methods.GET_PROP, params={"property": props}
+        )
+
+
+def create(channel: MqttChannel) -> B01PropsApi:
     """Create traits for B01 devices."""
-    return [B01PropsApi(channel)]
+    return B01PropsApi(channel)

roborock/devices/traits/traits_mixin.py

@@ -22,7 +22,7 @@
 class TraitsMixin:
     """Mixin to provide trait accessors."""
 
-    v1_properties: v1.Properties | None = None
+    v1_properties: v1.PropertiesApi | None = None
     """V1 properties trait, if supported."""
 
     dyad: a01.DyadApi | None = None
@@ -31,20 +31,20 @@ class TraitsMixin:
     zeo: a01.ZeoApi | None = None
     """Zeo API, if supported."""
 
-    b01_properties: b01.B01PropsApi | None = None
+    b01_props_api: b01.B01PropsApi | None = None
     """B01 properties trait, if supported."""
 
-    def __init__(self, traits: list[Trait]) -> None:
-        """Initialize the TraitsMixin with the given traits list.
+    def __init__(self, trait: Trait) -> None:
+        """Initialize the TraitsMixin with the given trait.
 
         This will populate the appropriate trait attributes based on the types
         of the traits provided.
         """
-        trait_map: dict[type[Trait], Trait] = {type(item): item for item in traits}
         for item in fields(self):
             trait_type = _get_trait_type(item)
-            if (trait := trait_map.get(trait_type, None)) is not None:
+            if trait_type == type(trait):
                 setattr(self, item.name, trait)
+                break
 
 
 def _get_trait_type(item) -> type[Trait]:

roborock/devices/traits/v1/__init__.py

@@ -6,17 +6,24 @@
 from roborock.devices.traits import Trait
 from roborock.devices.v1_rpc_channel import V1RpcChannel
 
-from .properties import CleanSummaryTrait, DoNotDisturbTrait, SoundVolumeTrait, StatusTrait
+from .clean_summary import CleanSummaryTrait
+from .common import V1TraitMixin
+from .do_not_disturb import DoNotDisturbTrait
+from .status import StatusTrait
+from .volume import SoundVolumeTrait
 
 __all__ = [
-    "create_v1_traits",
-    "Properties",
-    "properties",
+    "create",
+    "PropertiesApi",
+    "StatusTrait",
+    "DoNotDisturbTrait",
+    "CleanSummaryTrait",
+    "SoundVolumeTrait",
 ]
 
 
 @dataclass
-class Properties(Trait):
+class PropertiesApi(Trait):
     """Common properties for V1 devices.
 
     This class holds all the traits that are common across all V1 devices.
@@ -34,16 +41,17 @@ def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel) -> None:
         """Initialize the V1TraitProps with None values."""
         self.status = StatusTrait(product)
 
+        # This is a hack to allow setting the rpc_channel on all traits. This is
+        # used so we can preserve the dataclass behavior when the values in the
+        # traits are updated, but still want to allow them to have a reference
+        # to the rpc channel for sending commands.
         for item in fields(self):
             if (trait := getattr(self, item.name, None)) is None:
                 trait = item.type()
                 setattr(self, item.name, trait)
             trait._rpc_channel = rpc_channel
 
 
-def create_v1_traits(product: HomeDataProduct, rpc_channel: V1RpcChannel) -> list[Trait]:
+def create(product: HomeDataProduct, rpc_channel: V1RpcChannel) -> PropertiesApi:
     """Create traits for V1 devices."""
-    return [
-        Properties(product, rpc_channel)
-        # Add optional traits here as needed in the future
-    ]
+    return PropertiesApi(product, rpc_channel)

roborock/devices/traits/v1/clean_summary.py

@@ -0,0 +1,29 @@
+from typing import Self
+
+from roborock.containers import CleanSummary
+from roborock.devices.traits.v1 import common
+from roborock.roborock_typing import RoborockCommand
+from roborock.util import unpack_list
+
+
+class CleanSummaryTrait(CleanSummary, common.V1TraitMixin):
+    """Trait for managing the clean summary of Roborock devices."""
+
+    command = RoborockCommand.GET_CLEAN_SUMMARY
+
+    @classmethod
+    def _parse_type_response(cls, response: common.V1ResponseData) -> Self:
+        """Parse the response from the device into a CleanSummary."""
+        if isinstance(response, dict):
+            return CleanSummaryTrait.from_dict(response)  # type: ignore[return-value]
+        elif isinstance(response, list):
+            clean_time, clean_area, clean_count, records = unpack_list(response, 4)
+            return CleanSummaryTrait(  # type: ignore[return-value]
+                clean_time=clean_time,
+                clean_area=clean_area,
+                clean_count=clean_count,
+                records=records,
+            )
+        elif isinstance(response, int):
+            return CleanSummaryTrait(clean_time=response)  # type: ignore[return-value]
+        raise ValueError(f"Unexpected clean summary format: {response!r}")

roborock/devices/traits/v1/do_not_disturb.py

@@ -0,0 +1,17 @@
+from roborock.containers import DnDTimer
+from roborock.devices.traits.v1 import common
+from roborock.roborock_typing import RoborockCommand
+
+
+class DoNotDisturbTrait(DnDTimer, common.V1TraitMixin):
+    """Trait for managing Do Not Disturb (DND) settings on Roborock devices."""
+
+    command = RoborockCommand.GET_DND_TIMER
+
+    async def set_dnd_timer(self, dnd_timer: DnDTimer) -> None:
+        """Set the Do Not Disturb (DND) timer settings of the device."""
+        await self.rpc_channel.send_command(RoborockCommand.SET_DND_TIMER, params=dnd_timer.as_dict())
+
+    async def clear_dnd_timer(self) -> None:
+        """Clear the Do Not Disturb (DND) timer settings of the device."""
+        await self.rpc_channel.send_command(RoborockCommand.CLOSE_DND_TIMER)