Add an approach for determining if a dataclass field is supported

roborock/data/containers.py

@@ -232,6 +232,13 @@ def summary_info(self) -> str:
  """Return a string with key product information for logging purposes."""
  return f"{self.name} (model={self.model}, category={self.category})"
 
+ @cached_property
+ def supported_schema_codes(self) -> set[str]:
+ """Return a set of fields that are supported by the device."""
+ if self.schema is None:
+ return set()
+ return {schema.code for schema in self.schema if schema.code is not None}
+
 
 @dataclass
 class HomeDataDevice(RoborockBase):

roborock/data/v1/v1_containers.py

@@ -1,6 +1,7 @@
 import datetime
 import logging
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from enum import StrEnum
 from typing import Any
 
 from roborock.const import (
@@ -91,12 +92,39 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+class FieldNameBase(StrEnum):
+ """A base enum class that represents a field name in a RoborockBase dataclass."""
+
+
+class StatusField(FieldNameBase):
+ """An enum that represents a field in the `Status` class.
+
+ This is used with `roborock.devices.traits.v1.status.DeviceFeaturesTrait`
+ to understand if a feature is supported by the device using `is_field_supported`.
+
+ The enum values are names of fields in the `Status` class. Each field is
+ annotated with `requires_schema_code` metadata to map the field to a schema
+ code in the product schema, which may have a different name than the field/attribute name.
+ """
+
+ STATE = "state"
+ BATTERY = "battery"
+ FAN_POWER = "fan_power"
+ WATER_BOX_MODE = "water_box_mode"
+ CHARGE_STATUS = "charge_status"
+ DRY_STATUS = "dry_status"
+
+
+def _requires_schema_code(requires_schema_code: str, default=None) -> Any:
+ return field(metadata={"requires_schema_code": requires_schema_code}, default=default)
+
+
 @dataclass
 class Status(RoborockBase):
  msg_ver: int | None = None
  msg_seq: int | None = None
- state: RoborockStateCode | None = None
- battery: int | None = None
+ state: RoborockStateCode | None = _requires_schema_code("state", default=None)
+ battery: int | None = _requires_schema_code("battery", default=None)
  clean_time: int | None = None
  clean_area: int | None = None
  error_code: RoborockErrorCode | None = None
@@ -109,12 +137,12 @@ class Status(RoborockBase):
  back_type: int | None = None
  wash_phase: int | None = None
  wash_ready: int | None = None
- fan_power: RoborockFanPowerCode | None = None
+ fan_power: RoborockFanPowerCode | None = _requires_schema_code("fan_power", default=None)
  dnd_enabled: int | None = None
  map_status: int | None = None
  is_locating: int | None = None
  lock_status: int | None = None
- water_box_mode: RoborockMopIntensityCode | None = None
+ water_box_mode: RoborockMopIntensityCode | None = _requires_schema_code("water_box_mode", default=None)
  water_box_carriage_status: int | None = None
  mop_forbidden_enable: int | None = None
  camera_status: int | None = None
@@ -132,13 +160,13 @@ class Status(RoborockBase):
  collision_avoid_status: int | None = None
  switch_map_mode: int | None = None
  dock_error_status: RoborockDockErrorCode | None = None
- charge_status: int | None = None
+ charge_status: int | None = _requires_schema_code("charge_status", default=None)
  unsave_map_reason: int | None = None
  unsave_map_flag: int | None = None
  wash_status: int | None = None
  distance_off: int | None = None
  in_warmup: int | None = None
- dry_status: int | None = None
+ dry_status: int | None = _requires_schema_code("drying_status", default=None)
  rdt: int | None = None
  clean_percent: int | None = None
  rss: int | None = None

roborock/devices/traits/v1/__init__.py

@@ -44,6 +44,12 @@
  available features.
  - `requires_dock_type` - If set, this is a function that accepts a `RoborockDockTypeCode`
  and returns a boolean indicating whether the trait is supported for that dock type.
+
+Additionally, DeviceFeaturesTrait has a method `is_field_supported` that is used to
+check individual trait field values. This is a more fine grained version to allow
+optional fields in a dataclass, vs the above feature checks that apply to an entire
+trait. The `requires_schema_code` field metadata attribute is a string of the schema
+code in HomeDataProduct Schema that is required for the field to be supported.
 """
 
 import logging
@@ -189,7 +195,7 @@ def __init__(
  self.maps = MapsTrait(self.status)
  self.map_content = MapContentTrait(map_parser_config)
  self.home = HomeTrait(self.status, self.maps, self.map_content, self.rooms, self._device_cache)
- self.device_features = DeviceFeaturesTrait(product.product_nickname, self._device_cache)
+ self.device_features = DeviceFeaturesTrait(product, self._device_cache)
  self.network_info = NetworkInfoTrait(device_uid, self._device_cache)
  self.routines = RoutinesTrait(device_uid, web_api)
 

roborock/devices/traits/v1/device_features.py

@@ -1,26 +1,50 @@
-from dataclasses import fields
+from dataclasses import Field, fields
 
-from roborock.data import AppInitStatus, RoborockProductNickname
+from roborock.data import AppInitStatus, HomeDataProduct, RoborockBase
+from roborock.data.v1.v1_containers import FieldNameBase
 from roborock.device_features import DeviceFeatures
 from roborock.devices.cache import DeviceCache
 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."""
+ """Trait for managing supported features on Roborock devices."""
 
  command = RoborockCommand.APP_GET_INIT_STATUS
 
- def __init__(self, product_nickname: RoborockProductNickname, device_cache: DeviceCache) -> None: # pylint: disable=super-init-not-called
- """Initialize MapContentTrait."""
- self._nickname = product_nickname
+ def __init__(self, product: HomeDataProduct, device_cache: DeviceCache) -> None: # pylint: disable=super-init-not-called
+ """Initialize DeviceFeaturesTrait."""
+ self._product = product
+ self._nickname = product.product_nickname
  self._device_cache = device_cache
  # All fields of DeviceFeatures are required. Initialize them to False
  # so we have some known state.
  for field in fields(self):
  setattr(self, field.name, False)
 
+ def is_field_supported(self, cls: type[RoborockBase], field_name: FieldNameBase) -> bool:
+ """Determines if the specified field is supported by this device.
+
+ We use dataclass attributes on the field to specify the schema code that is required
+ for the field to be supported and it is compared against the list of
+ supported schema codes for the device returned in the product information.
+ """
+ dataclass_field: Field | None = None
+ for field in fields(cls):
+ if field.name == field_name:
+ dataclass_field = field
+ break
+ if dataclass_field is None:
+ raise ValueError(f"Field {field_name} not found in {cls}")
+
+ requires_schema_code = dataclass_field.metadata.get("requires_schema_code", None)
+ if requires_schema_code is None:
+ # We assume the field is supported
+ return True
+ # If the field requires a protocol that is not supported, we return False
+ return requires_schema_code in self._product.supported_schema_codes
+
  async def refresh(self) -> None:
  """Refresh the contents of this trait.
 

tests/data/test_containers.py

@@ -195,6 +195,25 @@ def test_home_data():
  assert schema[0].type == "RAW"
  assert schema[0].product_property is None
  assert schema[0].desc is None
+ assert product.supported_schema_codes == {
+ "additional_props",
+ "battery",
+ "charge_status",
+ "drying_status",
+ "error_code",
+ "fan_power",
+ "filter_life",
+ "main_brush_life",
+ "rpc_request_code",
+ "rpc_response",
+ "side_brush_life",
+ "state",
+ "task_cancel_in_motion",
+ "task_cancel_low_power",
+ "task_complete",
+ "water_box_mode",
+ }
+
  device = hd.devices[0]
  assert device.duid == "abc123"
  assert device.name == "Roborock S7 MaxV"

tests/devices/traits/v1/__snapshots__/test_device_features.ambr

@@ -0,0 +1,31 @@
+# serializer version: 1
+# name: test_is_attribute_supported[home_data_device_s5e.json]
+ dict({
+ 'battery': True,
+ 'charge_status': True,
+ 'dry_status': True,
+ 'fan_power': True,
+ 'state': True,
+ 'water_box_mode': True,
+ })
+# ---
+# name: test_is_attribute_supported[home_data_device_s7_maxv.json]
+ dict({
+ 'battery': True,
+ 'charge_status': True,
+ 'dry_status': True,
+ 'fan_power': True,
+ 'state': True,
+ 'water_box_mode': True,
+ })
+# ---
+# name: test_is_attribute_supported[home_data_device_saros_10r.json]
+ dict({
+ 'battery': True,
+ 'charge_status': True,
+ 'dry_status': True,
+ 'fan_power': True,
+ 'state': True,
+ 'water_box_mode': True,
+ })
+# ---

tests/devices/traits/v1/fixtures.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from roborock.data import HomeData, RoborockDockTypeCode, S7MaxVStatus, UserData
+from roborock.data import HomeData, HomeDataDevice, HomeDataProduct, RoborockDockTypeCode, S7MaxVStatus, UserData
 from roborock.devices.cache import Cache, DeviceCache, InMemoryCache
 from roborock.devices.device import RoborockDevice
 from roborock.devices.traits import v1
@@ -57,6 +57,18 @@ def device_cache_fixture(roborock_cache: Cache) -> DeviceCache:
  return DeviceCache(HOME_DATA.devices[0].duid, roborock_cache)
 
 
+@pytest.fixture(name="device_info")
+def device_info_fixture() -> HomeDataDevice:
+ """Fixture to provide a DeviceInfo instance for tests."""
+ return HOME_DATA.devices[0]
+
+
+@pytest.fixture(name="products")
+def products_fixture() -> list[HomeDataProduct]:
+ """Fixture to provide a Product instance for tests."""
+ return [HomeDataProduct.from_dict(product) for product in mock_data.PRODUCTS.values()]
+
+
 @pytest.fixture(autouse=True, name="device")
 def device_fixture(
  channel: AsyncMock,
@@ -65,15 +77,18 @@ def device_fixture(
  mock_map_rpc_channel: AsyncMock,
  web_api_client: AsyncMock,
  device_cache: DeviceCache,
+ device_info: HomeDataDevice,
+ products: list[HomeDataProduct],
 ) -> RoborockDevice:
  """Fixture to set up the device for tests."""
+ product = next(filter(lambda product: product.id == device_info.product_id, products))
  return RoborockDevice(
- device_info=HOME_DATA.devices[0],
- product=HOME_DATA.products[0],
+ device_info=device_info,
+ product=product,
  channel=channel,
  trait=v1.create(
-HOME_DATA.devices[0].duid,
-HOME_DATA.products[0],
+device_info.duid,
+product,
  HOME_DATA,
  mock_rpc_channel,
  mock_mqtt_rpc_channel,

tests/devices/traits/v1/test_device_features.py

@@ -0,0 +1,33 @@
+"""Tests for the DeviceFeaturesTrait related functionality."""
+
+import pytest
+from syrupy import SnapshotAssertion
+
+from roborock.data import HomeDataDevice
+from roborock.data.v1.v1_containers import StatusField
+from roborock.devices.device import RoborockDevice
+from roborock.devices.traits.v1.status import StatusTrait
+from tests import mock_data
+
+V1_DEVICES = {
+ k: HomeDataDevice.from_dict(device) for k, device in mock_data.DEVICES.items() if device.get("pv") == "1.0"
+}
+
+
+@pytest.mark.parametrize(
+ ("device_info"),
+ V1_DEVICES.values(),
+ ids=list(V1_DEVICES.keys()),
+)
+async def test_is_attribute_supported(
+ device_info: HomeDataDevice,
+ device: RoborockDevice,
+ snapshot: SnapshotAssertion,
+) -> None:
+ """Test if a field is supported."""
+ assert device.v1_properties is not None
+ assert device.v1_properties.device_features is not None
+ device_features_trait = device.v1_properties.device_features
+
+ is_supported = {field.value: device_features_trait.is_field_supported(StatusTrait, field) for field in StatusField}
+ assert is_supported == snapshot