Draft to set trait values

Author: allenporter

roborock/devices/traits/v1/child_lock.py

@@ -9,6 +9,7 @@ class ChildLockTrait(ChildLockStatus, common.V1TraitMixin, common.RoborockSwitch
     """Trait for controlling the child lock of a Roborock device."""
 
     command = RoborockCommand.GET_CHILD_LOCK_STATUS
+    converter = common.V1TraitDataConverter(ChildLockStatus)
     requires_feature = "is_set_child_supported"
 
     @property

roborock/devices/traits/v1/clean_summary.py

@@ -9,26 +9,12 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-class CleanSummaryTrait(CleanSummaryWithDetail, common.V1TraitMixin):
-    """Trait for managing the clean summary of Roborock devices."""
+class CleanSummaryConverter(common.V1TraitDataConverter):
 
-    command = RoborockCommand.GET_CLEAN_SUMMARY
 
-    async def refresh(self) -> None:
-        """Refresh the clean summary data and last clean record.
-
-        Assumes that the clean summary has already been fetched.
-        """
-        await super().refresh()
-        if not self.records:
-            _LOGGER.debug("No clean records available in clean summary.")
-            self.last_clean_record = None
-            return
-        last_record_id = self.records[0]
-        self.last_clean_record = await self.get_clean_record(last_record_id)
 
     @classmethod
-    def _parse_type_response(cls, response: common.V1ResponseData) -> Self:
+    def _convert_to_dict(cls, response: common.V1ResponseData) -> Self:
         """Parse the response from the device into a CleanSummary."""
         if isinstance(response, dict):
             return cls.from_dict(response)
@@ -44,6 +30,26 @@ def _parse_type_response(cls, response: common.V1ResponseData) -> Self:
             return cls(clean_time=response)
         raise ValueError(f"Unexpected clean summary format: {response!r}")
 
+
+class CleanSummaryTrait(CleanSummaryWithDetail, common.V1TraitMixin):
+    """Trait for managing the clean summary of Roborock devices."""
+
+    command = RoborockCommand.GET_CLEAN_SUMMARY
+    converter = common.V1TraitDataConverter(CleanSummaryWithDetail)
+
+    async def refresh(self) -> None:
+        """Refresh the clean summary data and last clean record.
+
+        Assumes that the clean summary has already been fetched.
+        """
+        await super().refresh()
+        if not self.records:
+            _LOGGER.debug("No clean records available in clean summary.")
+            self.last_clean_record = None
+            return
+        last_record_id = self.records[0]
+        self.last_clean_record = await self.get_clean_record(last_record_id)
+
     async def get_clean_record(self, record_id: int) -> CleanRecord:
         """Load a specific clean record by ID."""
         response = await self.rpc_channel.send_command(RoborockCommand.GET_CLEAN_RECORD, params=[record_id])

roborock/devices/traits/v1/common.py

@@ -6,18 +6,50 @@
 import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, fields
-from typing import ClassVar, Self
+from typing import Any, ClassVar, Self
 
+from roborock.callbacks import CallbackList
 from roborock.data import RoborockBase
 from roborock.protocols.v1_protocol import V1RpcChannel
 from roborock.roborock_typing import RoborockCommand
+from roborock.devices.traits.common import TraitDataConverter, TraitUpdateListener
 
 _LOGGER = logging.getLogger(__name__)
 
 V1ResponseData = dict | list | int | str
 
+class V1TraitDataConverter:
+    """Utility to handle the transformation and merging of data into models.
+
+    This has some special handling for v1 responses, then uses the common data
+    converters for updating the target objects.
+    """
+    
+    def __init__(self, dataclass_type: type[RoborockBase]):
+        """Initialize the converter for a specific RoborockBase-derived class."""
+        self._dataclass_type = dataclass_type
+        self._converter = TraitDataConverter(dataclass_type)
+
+    def update_from_data(self, target: RoborockBase, data: V1ResponseData) -> None:
+        """Update the target object from a dictionary of raw values.
+
+        Returns True if any values were updated.
+        """
+        response = self._convert_to_dict(data)
+        return self._converter.update_from_dict(target, response)
+
+    def _convert_to_dict(self, response: V1ResponseData) -> None:
+        """Convert the values to a dict that can be parsed as a RoborockBase.
+        
+        Subclasses can override to implement custom parsing logic
+        """
+        if isinstance(response, list):
+            response = response[0]
+        if not isinstance(response, dict):
+            raise ValueError(f"Unexpected {self._dataclass_type.__name__} response format: {response!r}")
+        return response
+
 
-@dataclass
 class V1TraitMixin(ABC):
     """Base model that supports v1 traits.
 
@@ -42,38 +74,13 @@ class V1TraitMixin(ABC):
     """
 
     command: ClassVar[RoborockCommand]
+    converter: ClassVar[V1TraitDataConverter]
 
-    @classmethod
-    def _parse_type_response(cls, response: V1ResponseData) -> RoborockBase:
-        """Parse the response from the device into a a RoborockBase.
-
-        Subclasses should override this method to implement custom parsing
-        logic as needed.
-        """
-        if not issubclass(cls, RoborockBase):
-            raise NotImplementedError(f"Trait {cls} does not implement RoborockBase")
-        # Subclasses can override to implement custom parsing logic
-        if isinstance(response, list):
-            response = response[0]
-        if not isinstance(response, dict):
-            raise ValueError(f"Unexpected {cls} response format: {response!r}")
-        return cls.from_dict(response)
-
-    def _parse_response(self, response: V1ResponseData) -> RoborockBase:
-        """Parse the response from the device into a a RoborockBase.
-
-        This is used by subclasses that want to override the class
-        behavior with instance-specific data.
-        """
-        return self._parse_type_response(response)
-
-    def __post_init__(self) -> None:
-        """Post-initialization to set up the RPC channel.
-
-        This is called automatically after the dataclass is initialized by the
-        device setup code.
-        """
+    def __init__(self) -> None:
+        """Initialize the V1TraitMixin."""
         self._rpc_channel = None
+        self._update_listener = TraitUpdateListener(_LOGGER)
+
 
     @property
     def rpc_channel(self) -> V1RpcChannel:
@@ -85,17 +92,23 @@ def rpc_channel(self) -> V1RpcChannel:
     async def refresh(self) -> None:
         """Refresh the contents of this trait."""
         response = await self.rpc_channel.send_command(self.command)
-        new_data = self._parse_response(response)
-        if not isinstance(new_data, RoborockBase):
-            raise ValueError(f"Internal error, unexpected response type: {new_data!r}")
-        _LOGGER.debug("Refreshed %s: %s", self.__class__.__name__, new_data)
-        self._update_trait_values(new_data)
+        self._update_data(response)
 
-    def _update_trait_values(self, new_data: RoborockBase) -> None:
-        """Update the values of this trait from another instance."""
-        for field in fields(new_data):
-            new_value = getattr(new_data, field.name, None)
-            setattr(self, field.name, new_value)
+    def _update_data(self, response: Any) -> None:
+        """Applies the new data to the object and notifies listeners.
+        
+        This can be overridden by subclasses to change the behavior, but it is
+        preferred to just use the `converter`.
+        """
+        if self.converter.update_from_data(self, response):
+            self._update_listener._notify_update()
+
+    def add_update_listener(self, callback: Callable[[], None]) -> Callable[[], None]:
+        """Register a callback when the trait has been updated.
+
+        Returns a callable to remove the listener.
+        """
+        self._update_listener.add_update_listener(callback)
 
 
 def _get_value_field(clazz: type[V1TraitMixin]) -> str:
@@ -108,9 +121,7 @@ def _get_value_field(clazz: type[V1TraitMixin]) -> str:
         )
     return value_fields[0]
 
-
-@dataclass(init=False, kw_only=True)
-class RoborockValueBase(V1TraitMixin, RoborockBase):
+class SingleValueConverter(V1TraitDataConverter):
     """Base class for traits that represent a single value.
 
     This class is intended to be subclassed by traits that represent a single
@@ -119,6 +130,11 @@ class RoborockValueBase(V1TraitMixin, RoborockBase):
     represents the main value of the trait.
     """
 
+    def __init__(self, dataclass_type: type[RoborockBase], value_field: str) -> None:
+        """Initialize the converter."""
+        super().__init__(dataclass_type)
+        self._value_field = value_field
+
     @classmethod
     def _parse_response(cls, response: V1ResponseData) -> Self:
         """Parse the response from the device into a RoborockValueBase."""
@@ -127,8 +143,7 @@ def _parse_response(cls, response: V1ResponseData) -> Self:
         if not isinstance(response, int):
             raise ValueError(f"Unexpected response format: {response!r}")
         value_field = _get_value_field(cls)
-        return cls(**{value_field: response})
-
+        return cls(**{self.__value_field: response})
 
 class RoborockSwitchBase(ABC):
     """Base class for traits that represent a boolean switch."""

roborock/devices/traits/v1/consumeable.py

@@ -41,6 +41,10 @@ class ConsumableTrait(Consumable, common.V1TraitMixin):
     """
 
     command = RoborockCommand.GET_CONSUMABLE
+    converter = common.V1TraitDataConverter(Consumable)
+
+    def __init__(self) -> None:
+        common.V1TraitMixin.__init__(self)
 
     async def reset_consumable(self, consumable: ConsumableAttribute) -> None:
         """Reset a specific consumable attribute on the device."""

roborock/devices/traits/v1/device_features.py

@@ -12,9 +12,11 @@ class DeviceFeaturesTrait(DeviceFeatures, common.V1TraitMixin):
     """Trait for managing supported features on Roborock devices."""
 
     command = RoborockCommand.APP_GET_INIT_STATUS
+    converter = common.V1TraitDataConverter(DeviceFeatures)
 
     def __init__(self, product: HomeDataProduct, device_cache: DeviceCache) -> None:  # pylint: disable=super-init-not-called
         """Initialize DeviceFeaturesTrait."""
+        common.V1TraitMixin.__init__(self)
         self._product = product
         self._nickname = product.product_nickname
         self._device_cache = device_cache

roborock/devices/traits/v1/do_not_disturb.py

@@ -9,6 +9,7 @@ class DoNotDisturbTrait(DnDTimer, common.V1TraitMixin, common.RoborockSwitchBase
     """Trait for managing Do Not Disturb (DND) settings on Roborock devices."""
 
     command = RoborockCommand.GET_DND_TIMER
+    converter = common.V1TraitDataConverter(DnDTimer)
 
     @property
     def is_on(self) -> bool:

roborock/devices/traits/v1/dust_collection_mode.py

@@ -10,4 +10,5 @@ class DustCollectionModeTrait(DustCollectionMode, common.V1TraitMixin):
     """Trait for dust collection mode."""
 
     command = RoborockCommand.GET_DUST_COLLECTION_MODE
+    converter = common.V1TraitDataConverter(DustCollectionMode)
     requires_dock_type = is_valid_dock

roborock/devices/traits/v1/flow_led_status.py

@@ -9,6 +9,7 @@ class FlowLedStatusTrait(FlowLedStatus, common.V1TraitMixin, common.RoborockSwit
     """Trait for controlling the Flow LED status of a Roborock device."""
 
     command = RoborockCommand.GET_FLOW_LED_STATUS
+    converter = common.V1TraitDataConverter(FlowLedStatus)
     requires_feature = "is_flow_led_setting_supported"
 
     @property

roborock/devices/traits/v1/home.py

@@ -41,6 +41,7 @@ class HomeTrait(RoborockBase, common.V1TraitMixin):
     """Trait that represents a full view of the home layout."""
 
     command = RoborockCommand.GET_MAP_V1  # This is not used
+    converter = common.V1TraitDataConverter(RoborockBase)
 
     def __init__(
         self,

roborock/devices/traits/v1/led_status.py

@@ -9,6 +9,7 @@ class LedStatusTrait(LedStatus, common.V1TraitMixin, common.RoborockSwitchBase):
     """Trait for controlling the LED status of a Roborock device."""
 
     command = RoborockCommand.GET_LED_STATUS
+    converter = common.V1TraitDataConverter(LedStatus)
     requires_feature = "is_led_status_switch_supported"
 
     @property