feat: add a trait for device features (#534)

Author: Lash-L

pyproject.toml

@@ -81,8 +81,8 @@ major_tags= ["refactor"]
 
 [tool.ruff]
 lint.ignore = ["F403", "E741"]
-line-length = 120
 lint.select=["E", "F", "UP", "I"]
+line-length = 120
 
 [tool.ruff.lint.per-file-ignores]
 "*/__init__.py" = ["F401"]
@@ -92,3 +92,9 @@ asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
 timeout = 30
 log_format = "%(asctime)s.%(msecs)03d %(levelname)s (%(threadName)s) [%(name)s] %(message)s"
+
+[tool.uv]
+dev-dependencies = [
+    "pyyaml>=6.0.3",
+    "pyshark>=0.6",
+]

roborock/cli.py

@@ -41,8 +41,9 @@
 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, DeviceFeatures, RoborockCommand, RoborockException
+from roborock import SHORT_MODEL_TO_ENUM, RoborockCommand, RoborockException
 from roborock.containers import CombinedMapInfo, DeviceData, HomeData, NetworkInfo, RoborockBase, UserData
+from roborock.device_features import DeviceFeatures
 from roborock.devices.cache import Cache, CacheData
 from roborock.devices.device import RoborockDevice
 from roborock.devices.device_manager import DeviceManager, create_device_manager, create_home_data_api
@@ -541,6 +542,16 @@ async def rooms(ctx, device_id: str):
     await _display_v1_trait(context, device_id, lambda v1: v1.rooms)
 
 
+@session.command()
+@click.option("--device_id", required=True)
+@click.pass_context
+@async_command
+async def features(ctx, device_id: str):
+    """Get device room mapping info."""
+    context: RoborockContext = ctx.obj
+    await _display_v1_trait(context, device_id, lambda v1: v1.device_features)
+
+
 @session.command()
 @click.option("--device_id", required=True)
 @click.option("--refresh", is_flag=True, default=False, help="Refresh status before discovery.")
@@ -825,6 +836,7 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur
 cli.add_command(reset_consumable)
 cli.add_command(rooms)
 cli.add_command(home)
+cli.add_command(features)
 
 
 def main():

roborock/containers.py

@@ -93,7 +93,6 @@
     STRAINER_REPLACE_TIME,
     ROBOROCK_G20S_Ultra,
 )
-from .device_features import DeviceFeatures
 from .exceptions import RoborockException
 
 _LOGGER = logging.getLogger(__name__)
@@ -280,6 +279,10 @@ class HomeDataProduct(RoborockBase):
     capability: int | None = None
     schema: list[HomeDataProductSchema] | None = None
 
+    @property
+    def product_nickname(self) -> RoborockProductNickname:
+        return SHORT_MODEL_TO_ENUM.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS)
+
 
 @dataclass
 class HomeDataDevice(RoborockBase):
@@ -863,7 +866,6 @@ class DeviceData(RoborockBase):
     device: HomeDataDevice
     model: str
     host: str | None = None
-    device_features: DeviceFeatures | None = None
 
     @property
     def product_nickname(self) -> RoborockProductNickname:

roborock/device_features.py

@@ -5,6 +5,7 @@
 from typing import Any
 
 from .code_mappings import RoborockProductNickname
+from .containers import RoborockBase
 
 
 class NewFeatureStrBit(IntEnum):
@@ -246,7 +247,7 @@ class ProductFeatures(StrEnum):
 
 
 @dataclass
-class DeviceFeatures:
+class DeviceFeatures(RoborockBase):
     """Represents the features supported by a Roborock device."""
 
     # Features from robot_new_features (lower 32 bits)

roborock/devices/traits/v1/__init__.py

@@ -12,6 +12,7 @@
 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 .home import HomeTrait
 from .map_content import MapContentTrait
@@ -33,6 +34,7 @@
     "MapContentTrait",
     "ConsumableTrait",
     "HomeTrait",
+    "DeviceFeaturesTrait",
 ]
 
 
@@ -53,6 +55,7 @@ class PropertiesApi(Trait):
     map_content: MapContentTrait
     consumables: ConsumableTrait
     home: HomeTrait
+    device_features: DeviceFeaturesTrait
 
     # In the future optional fields can be added below based on supported features
 
@@ -72,6 +75,7 @@ def __init__(
         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

roborock/devices/traits/v1/device_features.py

@@ -0,0 +1,26 @@
+from roborock import AppInitStatus, RoborockProductNickname
+from roborock.device_features import DeviceFeatures
+from roborock.devices.traits.v1 import common
+from roborock.roborock_typing import RoborockCommand
+
+
+class DeviceFeaturesTrait(DeviceFeatures, common.V1TraitMixin):
+    """Trait for managing Do Not Disturb (DND) settings on Roborock devices."""
+
+    command = RoborockCommand.APP_GET_INIT_STATUS
+
+    def __init__(self, product_nickname: RoborockProductNickname) -> None:
+        """Initialize MapContentTrait."""
+        self._nickname = product_nickname
+
+    def _parse_response(self, response: common.V1ResponseData) -> DeviceFeatures:
+        """Parse the response from the device into a MapContentTrait instance."""
+        if not isinstance(response, list):
+            raise ValueError(f"Unexpected AppInitStatus response format: {type(response)}")
+        app_status = AppInitStatus.from_dict(response[0])
+        return DeviceFeatures.from_feature_flags(
+            new_feature_info=app_status.new_feature_info,
+            new_feature_info_str=app_status.new_feature_info_str,
+            feature_info=app_status.feature_info,
+            product_nickname=self._nickname,
+        )