feat: Add various optional features with support for checking device features

roborock/cli.py

@@ -41,7 +41,7 @@
 from pyshark.capture.live_capture import LiveCapture, UnknownInterfaceException  # type: ignore
 from pyshark.packet.packet import Packet  # type: ignore
 
-from roborock import SHORT_MODEL_TO_ENUM, RoborockCommand, RoborockException
+from roborock import SHORT_MODEL_TO_ENUM, RoborockCommand
 from roborock.containers import CombinedMapInfo, DeviceData, HomeData, NetworkInfo, RoborockBase, UserData
 from roborock.device_features import DeviceFeatures
 from roborock.devices.cache import Cache, CacheData
@@ -51,6 +51,7 @@
 from roborock.devices.traits.v1 import V1TraitMixin
 from roborock.devices.traits.v1.consumeable import ConsumableAttribute
 from roborock.devices.traits.v1.map_content import MapContentTrait
+from roborock.exceptions import RoborockException, RoborockUnsupportedFeature
 from roborock.protocol import MessageParser
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
 from roborock.web_api import RoborockApiClient
@@ -401,14 +402,21 @@ async def _v1_trait(context: RoborockContext, device_id: str, display_func: Call
     device = await device_manager.get_device(device_id)
     if device.v1_properties is None:
         raise RoborockException(f"Device {device.name} does not support V1 protocol")
-
+    await device.v1_properties.discover_features()
     trait = display_func(device.v1_properties)
     await trait.refresh()
     return trait
 
 
 async def _display_v1_trait(context: RoborockContext, device_id: str, display_func: Callable[[], Trait]) -> None:
-    trait = await _v1_trait(context, device_id, display_func)
+    try:
+        trait = await _v1_trait(context, device_id, display_func)
+    except RoborockUnsupportedFeature:
+        click.echo("Feature not supported by device")
+        return
+    except RoborockException as e:
+        click.echo(f"Error: {e}")
+        return
     click.echo(dump_json(trait.as_dict()))
 
 
@@ -532,6 +540,116 @@ async def reset_consumable(ctx, device_id: str, consumable: str):
     click.echo(f"Reset {consumable} for device {device_id}")
 
 
+@session.command()
+@click.option("--device_id", required=True)
+@click.option("--enabled", type=bool, help="Enable (True) or disable (False) the child lock.")
+@click.pass_context
+@async_command
+async def child_lock(ctx, device_id: str, enabled: bool | None):
+    """Get device child lock status."""
+    context: RoborockContext = ctx.obj
+    try:
+        trait = await _v1_trait(context, device_id, lambda v1: v1.child_lock)
+    except RoborockUnsupportedFeature:
+        click.echo("Feature not supported by device")
+        return
+    if enabled is not None:
+        if enabled:
+            await trait.enable()
+        else:
+            await trait.disable()
+        click.echo(f"Set child lock to {enabled} for device {device_id}")
+        await trait.refresh()
+
+    click.echo(dump_json(trait.as_dict()))
+
+
+@session.command()
+@click.option("--device_id", required=True)
+@click.option("--enabled", type=bool, help="Enable (True) or disable (False) the DND status.")
+@click.pass_context
+@async_command
+async def dnd(ctx, device_id: str, enabled: bool | None):
+    """Get Do Not Disturb Timer status."""
+    context: RoborockContext = ctx.obj
+    try:
+        trait = await _v1_trait(context, device_id, lambda v1: v1.dnd)
+    except RoborockUnsupportedFeature:
+        click.echo("Feature not supported by device")
+        return
+    if enabled is not None:
+        if enabled:
+            await trait.enable()
+        else:
+            await trait.disable()
+        click.echo(f"Set DND to {enabled} for device {device_id}")
+        await trait.refresh()
+
+    click.echo(dump_json(trait.as_dict()))
+
+
+@session.command()
+@click.option("--device_id", required=True)
+@click.option("--enabled", required=False, type=bool, help="Enable (True) or disable (False) the Flow LED.")
+@click.pass_context
+@async_command
+async def flow_led_status(ctx, device_id: str, enabled: bool | None):
+    """Get device Flow LED status."""
+    context: RoborockContext = ctx.obj
+    try:
+        trait = await _v1_trait(context, device_id, lambda v1: v1.flow_led_status)
+    except RoborockUnsupportedFeature:
+        click.echo("Feature not supported by device")
+        return
+    if enabled is not None:
+        if enabled:
+            await trait.enable()
+        else:
+            await trait.disable()
+        click.echo(f"Set Flow LED to {enabled} for device {device_id}")
+        await trait.refresh()
+
+    click.echo(dump_json(trait.as_dict()))
+
+
+@session.command()
+@click.option("--device_id", required=True)
+@click.option("--enabled", required=False, type=bool, help="Enable (True) or disable (False) the LED.")
+@click.pass_context
+@async_command
+async def led_status(ctx, device_id: str, enabled: bool | None):
+    """Get device LED status."""
+    context: RoborockContext = ctx.obj
+    try:
+        trait = await _v1_trait(context, device_id, lambda v1: v1.led_status)
+    except RoborockUnsupportedFeature:
+        click.echo("Feature not supported by device")
+        return
+    if enabled is not None:
+        if enabled:
+            await trait.enable()
+        else:
+            await trait.disable()
+        click.echo(f"Set LED Status to {enabled} for device {device_id}")
+        await trait.refresh()
+
+    click.echo(dump_json(trait.as_dict()))
+
+
+@session.command()
+@click.option("--device_id", required=True)
+@click.option("--enabled", required=True, type=bool, help="Enable (True) or disable (False) the child lock.")
+@click.pass_context
+@async_command
+async def set_child_lock(ctx, device_id: str, enabled: bool):
+    """Set the child lock status."""
+    context: RoborockContext = ctx.obj
+    trait = await _v1_trait(context, device_id, lambda v1: v1.child_lock)
+    await trait.set_child_lock(enabled)
+    status = "enabled" if enabled else "disabled"
+    click.echo(f"Child lock {status} for device {device_id}")
+
+
 @session.command()
 @click.option("--device_id", required=True)
 @click.pass_context
@@ -837,6 +955,10 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur
 cli.add_command(rooms)
 cli.add_command(home)
 cli.add_command(features)
+cli.add_command(child_lock)
+cli.add_command(dnd)
+cli.add_command(flow_led_status)
+cli.add_command(led_status)
 
 
 def main():

roborock/containers.py

@@ -913,12 +913,17 @@ class CombinedMapInfo(RoborockBase):
 
 @dataclass
 class ChildLockStatus(RoborockBase):
-    lock_status: int
+    lock_status: int = 0
 
 
 @dataclass
 class FlowLedStatus(RoborockBase):
-    status: int
+    status: int = 0
+
+
+@dataclass
+class LedStatus(RoborockBase):
+    status: int = 0
 
 
 @dataclass

roborock/devices/cache.py

@@ -9,6 +9,7 @@
 from typing import Protocol
 
 from roborock.containers import CombinedMapInfo, HomeData, NetworkInfo
+from roborock.device_features import DeviceFeatures
 
 
 @dataclass
@@ -24,6 +25,9 @@ class CacheData:
     home_cache: dict[int, CombinedMapInfo] = field(default_factory=dict)
     """Home cache information indexed by map_flag."""
 
+    device_features: DeviceFeatures | None = None
+    """Device features information."""
+
 
 class Cache(Protocol):
     """Protocol for a cache that can store and retrieve values."""

roborock/devices/traits/v1/__init__.py

@@ -2,19 +2,23 @@
 
 import logging
 from dataclasses import dataclass, field, fields
+from typing import get_args
 
 from roborock.containers import HomeData, HomeDataProduct
 from roborock.devices.cache import Cache
 from roborock.devices.traits import Trait
 from roborock.devices.v1_rpc_channel import V1RpcChannel
 from roborock.map.map_parser import MapParserConfig
 
+from .child_lock import ChildLockTrait
 from .clean_summary import CleanSummaryTrait
 from .common import V1TraitMixin
 from .consumeable import ConsumableTrait
 from .device_features import DeviceFeaturesTrait
 from .do_not_disturb import DoNotDisturbTrait
+from .flow_led_status import FlowLedStatusTrait
 from .home import HomeTrait
+from .led_status import LedStatusTrait
 from .map_content import MapContentTrait
 from .maps import MapsTrait
 from .rooms import RoomsTrait
@@ -35,6 +39,9 @@
     "ConsumableTrait",
     "HomeTrait",
     "DeviceFeaturesTrait",
+    "ChildLockTrait",
+    "FlowLedStatusTrait",
+    "LedStatusTrait",
 ]
 
 
@@ -57,7 +64,10 @@ class PropertiesApi(Trait):
     home: HomeTrait
     device_features: DeviceFeaturesTrait
 
-    # In the future optional fields can be added below based on supported features
+    # Optional features that may not be supported on all devices
+    child_lock: ChildLockTrait | None = None
+    led_status: LedStatusTrait | None = None
+    flow_led_status: FlowLedStatusTrait | None = None
 
     def __init__(
         self,
@@ -70,28 +80,71 @@ def __init__(
         map_parser_config: MapParserConfig | None = None,
     ) -> None:
         """Initialize the V1TraitProps."""
+        self._rpc_channel = rpc_channel
+        self._mqtt_rpc_channel = mqtt_rpc_channel
+        self._map_rpc_channel = map_rpc_channel
+
         self.status = StatusTrait(product)
         self.rooms = RoomsTrait(home_data)
         self.maps = MapsTrait(self.status)
         self.map_content = MapContentTrait(map_parser_config)
         self.home = HomeTrait(self.status, self.maps, self.rooms, cache)
-        self.device_features = DeviceFeaturesTrait(product.product_nickname)
-        # 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.
+        self.device_features = DeviceFeaturesTrait(product.product_nickname, cache)
+
+        # Dynamically create any traits that need to be populated
         for item in fields(self):
             if (trait := getattr(self, item.name, None)) is None:
+                # We exclude optional features and them via discover_features
+                if (union_args := get_args(item.type)) is None or len(union_args) > 0:
+                    continue
+                _LOGGER.debug("Initializing trait %s", item.name)
                 trait = item.type()
                 setattr(self, item.name, trait)
-            # The decorator `@common.mqtt_rpc_channel` means that the trait needs
-            # to use the mqtt_rpc_channel (cloud only) instead of the rpc_channel (adaptive)
-            if hasattr(trait, "mqtt_rpc_channel"):
-                trait._rpc_channel = mqtt_rpc_channel
-            elif hasattr(trait, "map_rpc_channel"):
-                trait._rpc_channel = map_rpc_channel
-            else:
-                trait._rpc_channel = rpc_channel
+            # 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.
+            trait._rpc_channel = self._get_rpc_channel(trait)
+
+    def _get_rpc_channel(self, trait: V1TraitMixin) -> V1RpcChannel:
+        # The decorator `@common.mqtt_rpc_channel` means that the trait needs
+        # to use the mqtt_rpc_channel (cloud only) instead of the rpc_channel (adaptive)
+        if hasattr(trait, "mqtt_rpc_channel"):
+            return self._mqtt_rpc_channel
+        elif hasattr(trait, "map_rpc_channel"):
+            return self._map_rpc_channel
+        else:
+            return self._rpc_channel
+
+    async def discover_features(self) -> None:
+        """Populate any supported traits that were not initialized in __init__."""
+        await self.device_features.refresh()
+
+        for item in fields(self):
+            if (trait := getattr(self, item.name, None)) is not None:
+                continue
+            if (union_args := get_args(item.type)) is None:
+                raise ValueError(f"Unexpected non-union type for trait {item.name}: {item.type}")
+            if len(union_args) != 2 or type(None) not in union_args:
+                raise ValueError(f"Unexpected non-optional type for trait {item.name}: {item.type}")
+
+            # Union args may not be in declared order
+            item_type = union_args[0] if union_args[1] is type(None) else union_args[1]
+            trait = item_type()
+            if not hasattr(trait, "requires_feature"):
+                _LOGGER.debug("Trait missing required feature %s", item.name)
+                continue
+            _LOGGER.debug("Checking for feature %s", trait.requires_feature)
+            is_supported = getattr(self.device_features, trait.requires_feature)
+            # _LOGGER.debug("Device features: %s", self.device_features)
+            if is_supported is None:
+                raise ValueError(f"Device feature '{trait.requires_feature}' on trait '{item.name}' is unknown")
+            if not is_supported:
+                _LOGGER.debug("Disabling optional feature trait %s", item.name)
+                continue
+            _LOGGER.debug("Enabling optional feature trait %s", item.name)
+            setattr(self, item.name, trait)
+            trait._rpc_channel = self._get_rpc_channel(trait)
 
 
 def create(

roborock/devices/traits/v1/child_lock.py

@@ -0,0 +1,20 @@
+from roborock.containers import ChildLockStatus
+from roborock.devices.traits.v1 import common
+from roborock.roborock_typing import RoborockCommand
+
+_STATUS_PARAM = "lock_status"
+
+
+class ChildLockTrait(ChildLockStatus, common.V1TraitMixin):
+    """Trait for controlling the child lock of a Roborock device."""
+
+    command = RoborockCommand.GET_CHILD_LOCK_STATUS
+    requires_feature = "is_set_child_supported"
+
+    async def enable(self) -> None:
+        """Enable the child lock."""
+        await self.rpc_channel.send_command(RoborockCommand.SET_CHILD_LOCK_STATUS, params={_STATUS_PARAM: 1})
+
+    async def disable(self) -> None:
+        """Disable the child lock."""
+        await self.rpc_channel.send_command(RoborockCommand.SET_CHILD_LOCK_STATUS, params={_STATUS_PARAM: 0})

roborock/devices/traits/v1/common.py

@@ -82,10 +82,14 @@ async def refresh(self) -> Self:
         new_data = self._parse_response(response)
         if not isinstance(new_data, RoborockBase):
             raise ValueError(f"Internal error, unexpected response type: {new_data!r}")
+        self._update_trait_values(new_data)
+        return self
+
+    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)
-        return self
 
 
 def _get_value_field(clazz: type[V1TraitMixin]) -> str: