feat: add better dock feature control

roborock/data/v1/v1_code_mappings.py

@@ -650,22 +650,31 @@ class RoborockDockErrorCode(RoborockEnum):
 
 class RoborockDockTypeCode(RoborockEnum):
  unknown = -9999
- no_dock = 0
- auto_empty_dock = 1
- empty_wash_fill_dock = 3
- auto_empty_dock_pure = 5
- s7_max_ultra_dock = 6
- s8_dock = 7
- p10_dock = 8
- p10_pro_dock = 9
- s8_maxv_ultra_dock = 10
- qrevo_master_dock = 14
- qrevo_s_dock = 15
- saros_r10_dock = 16
- qrevo_curv_dock = 17
- saros_10_dock = 18
- qrevo_s5v_dock = 22
- saros_20_dock = 27
+ o0_dock = 0
+ o1_dock = 1
+ o2_dock = 2
+ o3_dock = 3
+ oc_dock = 5
+ o3_plus_dock = 6
+ o4_dock = 7
+ pearl_dock = 8
+ pearl_plus_dock = 9
+ o5_dock = 10
+ shell_2s_dock = 11
+ couple_dock = 13
+ shell_3_dock = 14
+ shell_2c_dock = 15
+ shell_3s_dock = 16
+ k1_dock = 17
+ o6_dock = 18
+ k1c_dock = 19
+ k1s_dock = 20
+ shell_e_dock = 21
+ shell_2e_dock = 22
+ shell_3c_dock = 23
+ type_27_dock = 27
+ k1c_lite_dock = 28
+ shell_2e_lite_dock = 30
 
 
 class RoborockDockDustCollectionModeCode(RoborockEnum):

roborock/data/v1/v1_containers.py

@@ -239,6 +239,12 @@ def current_map(self) -> int | None:
  return map_flag
  return None
 
+ @property
+ def has_am(self) -> bool | None:
+ if self.dss is None:
+ return None
+ return (self.dss & 3) == 2
+
  @property
  def clear_water_box_status(self) -> ClearWaterBoxStatus | None:
  if self.dss:
@@ -376,6 +382,12 @@ def current_map(self) -> int | None:
  return map_flag
  return None
 
+ @property
+ def has_am(self) -> bool | None:
+ if self.dss is None:
+ return None
+ return (self.dss & 3) == 2
+
  @property
  def clear_water_box_status(self) -> ClearWaterBoxStatus | None:
  if self.dss:

roborock/device_features.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, field, fields
 from enum import IntEnum, StrEnum
+from functools import cache
 from typing import Any, Self
 
 from roborock.data.code_mappings import RoborockProductNickname
@@ -644,25 +645,219 @@ def get_supported_features(self) -> list[str]:
  return [k for k, v in vars(self).items() if v]
 
 
+_NO_DOCK_TYPES = {
+ RoborockDockTypeCode.unknown,
+ RoborockDockTypeCode.o0_dock,
+}
+
+_PURE_COLLECT_DOCK_TYPES = {
+ RoborockDockTypeCode.o1_dock,
+ RoborockDockTypeCode.oc_dock,
+}
+
+_PURE_WASH_DOCK_TYPES = {
+ RoborockDockTypeCode.o2_dock,
+}
+
+_COLLECT_WASH_DOCK_TYPES = {
+ RoborockDockTypeCode.o3_dock,
+}
+
+_ONLY_X_SERIES_DOCK_TYPES = {
+ RoborockDockTypeCode.o2_dock,
+ RoborockDockTypeCode.o3_dock,
+ RoborockDockTypeCode.o3_plus_dock,
+ RoborockDockTypeCode.o4_dock,
+ RoborockDockTypeCode.o5_dock,
+ RoborockDockTypeCode.o6_dock,
+}
+
+_CLEAN_FLUID_AUTO_DELIVERY_DOCK_TYPES = {
+ RoborockDockTypeCode.o4_dock,
+ RoborockDockTypeCode.pearl_plus_dock,
+ RoborockDockTypeCode.shell_3s_dock,
+ RoborockDockTypeCode.o6_dock,
+ RoborockDockTypeCode.couple_dock,
+}
+
+_CLEAN_FLUID_AUTO_DELIVERY_AM_DOCK_TYPES = {
+ RoborockDockTypeCode.o5_dock,
+ RoborockDockTypeCode.shell_2s_dock,
+ RoborockDockTypeCode.shell_3_dock,
+ RoborockDockTypeCode.shell_2c_dock,
+ RoborockDockTypeCode.k1_dock,
+ RoborockDockTypeCode.k1s_dock,
+ RoborockDockTypeCode.shell_e_dock,
+ RoborockDockTypeCode.shell_2e_dock,
+ RoborockDockTypeCode.shell_3c_dock,
+}
+
+_AUTO_STERILIZE_DOCK_TYPES = {
+ RoborockDockTypeCode.shell_3s_dock,
+}
+
+_INNER_CHARGE_HEAD_DOCK_TYPES = {
+ RoborockDockTypeCode.o0_dock,
+ RoborockDockTypeCode.o1_dock,
+ RoborockDockTypeCode.oc_dock,
+ RoborockDockTypeCode.o2_dock,
+ RoborockDockTypeCode.o3_dock,
+ RoborockDockTypeCode.o3_plus_dock,
+ RoborockDockTypeCode.o4_dock,
+}
+
+_DUST_BUCKET_DOCK_TYPES = {
+ RoborockDockTypeCode.o3_dock,
+ RoborockDockTypeCode.o3_plus_dock,
+ RoborockDockTypeCode.o4_dock,
+}
+
+_CLEAN_CAROUSEL_SELF_CLEAN_DOCK_TYPES = {
+ RoborockDockTypeCode.shell_3_dock,
+ RoborockDockTypeCode.k1_dock,
+ RoborockDockTypeCode.couple_dock,
+ RoborockDockTypeCode.shell_3s_dock,
+ RoborockDockTypeCode.k1s_dock,
+ RoborockDockTypeCode.shell_3c_dock,
+}
+
+_WATER_UPDOWN_DRAIN_DOCK_TYPES = {
+ RoborockDockTypeCode.o4_dock,
+ RoborockDockTypeCode.pearl_plus_dock,
+ RoborockDockTypeCode.couple_dock,
+}
+
+_WATER_UPDOWN_DRAIN_AM_DOCK_TYPES = {
+ RoborockDockTypeCode.o5_dock,
+ RoborockDockTypeCode.shell_2s_dock,
+ RoborockDockTypeCode.shell_3_dock,
+ RoborockDockTypeCode.shell_2c_dock,
+ RoborockDockTypeCode.shell_3s_dock,
+ RoborockDockTypeCode.k1_dock,
+ RoborockDockTypeCode.o6_dock,
+ RoborockDockTypeCode.shell_e_dock,
+ RoborockDockTypeCode.shell_2e_dock,
+}
+
+_HATCH_DOOR_DOCK_COOL_FAN_DOCK_TYPES = {
+ RoborockDockTypeCode.couple_dock,
+}
+
+_SPECIAL_WASH_TEMP_DOCK_TYPES = {
+ RoborockDockTypeCode.shell_2e_lite_dock,
+}
+
+
+@dataclass(frozen=True)
+class RoborockDockFeatures:
+ """Capabilities derived from a V1 dock type.
+
+ This mirrors the Roborock app's DK capability model: feature availability
+ should be attached to the dock family instead of repeated per trait.
+ """
+
+ dock_type: RoborockDockTypeCode
+ has_am: bool | None = None
+
+ @classmethod
+ @cache
+ def from_dock_type(cls, dock_type: RoborockDockTypeCode | None, has_am: bool | None = None) -> Self:
+ """Return cached capabilities for a V1 dock type."""
+ return cls(dock_type or RoborockDockTypeCode.o0_dock, has_am=has_am)
+
+ def _matches_am_variant(self, dock_types: set[RoborockDockTypeCode]) -> bool:
+ return self.has_am is True and self.dock_type in dock_types
+
+ def _matches_non_am_variant(self, dock_types: set[RoborockDockTypeCode]) -> bool:
+ return self.has_am is not True and self.dock_type in dock_types
+
+ @property
+ def has_dock(self) -> bool:
+ return self.dock_type not in _NO_DOCK_TYPES
+
+ @property
+ def is_pure_collect(self) -> bool:
+ return self.dock_type in _PURE_COLLECT_DOCK_TYPES
+
+ @property
+ def is_pure_wash(self) -> bool:
+ return self.dock_type in _PURE_WASH_DOCK_TYPES
+
+ @property
+ def is_collect_wash(self) -> bool:
+ return self.dock_type in _COLLECT_WASH_DOCK_TYPES
+
+ @property
+ def is_collect_wash_dry(self) -> bool:
+ return self.has_dock and not (self.is_pure_collect or self.is_pure_wash or self.is_collect_wash)
+
+ @property
+ def is_collectable(self) -> bool:
+ return self.is_collect_wash or self.is_collect_wash_dry or self.is_pure_collect
+
+ @property
+ def is_washable(self) -> bool:
+ return self.is_pure_wash or self.is_collect_wash or self.is_collect_wash_dry
+
+ @property
+ def is_dryable(self) -> bool:
+ return self.is_collect_wash_dry
+
+ @property
+ def is_only_x_series(self) -> bool:
+ return self.dock_type in _ONLY_X_SERIES_DOCK_TYPES
+
+ @property
+ def is_cleaning_brush_supported(self) -> bool:
+ return self.is_only_x_series
+
+ @property
+ def is_clean_fluid_auto_delivery_supported(self) -> bool:
+ return self.dock_type in _CLEAN_FLUID_AUTO_DELIVERY_DOCK_TYPES or self._matches_am_variant(
+ _CLEAN_FLUID_AUTO_DELIVERY_AM_DOCK_TYPES
+ )
+
+ @property
+ def is_auto_sterilize_supported(self) -> bool:
+ return self.dock_type in _AUTO_STERILIZE_DOCK_TYPES
+
+ @property
+ def is_inner_charge_head(self) -> bool:
+ return self.dock_type in _INNER_CHARGE_HEAD_DOCK_TYPES
+
+ @property
+ def is_dust_bucket_supported(self) -> bool:
+ return self.dock_type in _DUST_BUCKET_DOCK_TYPES
+
+ @property
+ def is_clean_carousel_self_clean_supported(self) -> bool:
+ return self._matches_non_am_variant(_CLEAN_CAROUSEL_SELF_CLEAN_DOCK_TYPES)
+
+ @property
+ def is_water_updown_drain_supported(self) -> bool:
+ return self.dock_type in _WATER_UPDOWN_DRAIN_DOCK_TYPES or self._matches_am_variant(
+ _WATER_UPDOWN_DRAIN_AM_DOCK_TYPES
+ )
+
+ @property
+ def is_hatch_door_dock_cool_fan_supported(self) -> bool:
+ return self.dock_type in _HATCH_DOOR_DOCK_COOL_FAN_DOCK_TYPES
+
+ @property
+ def is_special_support_wash_temp(self) -> bool:
+ return self.dock_type in _SPECIAL_WASH_TEMP_DOCK_TYPES
+
+
 WASH_N_FILL_DOCK_TYPES = [
- RoborockDockTypeCode.empty_wash_fill_dock,
- RoborockDockTypeCode.s8_dock,
- RoborockDockTypeCode.p10_dock,
- RoborockDockTypeCode.p10_pro_dock,
- RoborockDockTypeCode.s8_maxv_ultra_dock,
- RoborockDockTypeCode.qrevo_s_dock,
- RoborockDockTypeCode.saros_r10_dock,
- RoborockDockTypeCode.qrevo_curv_dock,
- RoborockDockTypeCode.qrevo_s5v_dock,
- RoborockDockTypeCode.saros_20_dock,
+ dock_type for dock_type in RoborockDockTypeCode if RoborockDockFeatures.from_dock_type(dock_type).is_washable
 ]
 
 
 def is_wash_n_fill_dock(dock_type: RoborockDockTypeCode) -> bool:
  """Check if the dock type is a wash and fill dock."""
- return dock_type in WASH_N_FILL_DOCK_TYPES
+ return RoborockDockFeatures.from_dock_type(dock_type).is_washable
 
 
 def is_valid_dock(dock_type: RoborockDockTypeCode) -> bool:
  """Check if device supports a dock."""
- return dock_type != RoborockDockTypeCode.no_dock
+ return RoborockDockFeatures.from_dock_type(dock_type).has_dock

roborock/devices/traits/v1/__init__.py

@@ -42,8 +42,8 @@
  - `requires_feature` - The string name of the device feature that must be supported
  for this trait to be enabled. See `DeviceFeaturesTrait` for a list of
  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.
+ - `requires_dock_features` - If set, this is a function that accepts a `RoborockDockFeatures`
+ and returns a boolean indicating whether the trait is supported for that dock.
 
 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
@@ -59,6 +59,7 @@
 
 from roborock.data.containers import HomeData, HomeDataProduct, RoborockBase
 from roborock.data.v1.v1_code_mappings import RoborockDockTypeCode
+from roborock.device_features import RoborockDockFeatures
 from roborock.devices.cache import DeviceCache
 from roborock.devices.traits import Trait
 from roborock.exceptions import RoborockException
@@ -262,9 +263,10 @@ async def discover_features(self) -> None:
  await self.device_features.refresh()
  # Dock type also acts like a device feature for some traits.
  dock_type = await self._dock_type()
+ dock_features = RoborockDockFeatures.from_dock_type(dock_type, has_am=self.status.has_am)
 
  # Initialize traits with special arguments before the generic loop
- if self.wash_towel_mode is None and self._is_supported(WashTowelModeTrait, "wash_towel_mode", dock_type):
+ if self.wash_towel_mode is None and self._is_supported(WashTowelModeTrait, "wash_towel_mode", dock_features):
  wash_towel_mode = WashTowelModeTrait(self.device_features)
  wash_towel_mode._rpc_channel = self._get_rpc_channel(wash_towel_mode) # type: ignore[assignment]
  self.wash_towel_mode = wash_towel_mode
@@ -280,19 +282,19 @@ async def discover_features(self) -> None:
 
  # Union args may not be in declared order
  item_type = union_args[0] if union_args[1] is type(None) else union_args[1]
- if not self._is_supported(item_type, item.name, dock_type):
+ if not self._is_supported(item_type, item.name, dock_features):
  _LOGGER.debug("Trait '%s' not supported, skipping", item.name)
  continue
  _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
  trait = item_type()
  setattr(self, item.name, trait)
  trait._rpc_channel = self._get_rpc_channel(trait)
 
- def _is_supported(self, trait_type: type[V1TraitMixin], name: str, dock_type: RoborockDockTypeCode) -> bool:
+ def _is_supported(self, trait_type: type[V1TraitMixin], name: str, dock_features: RoborockDockFeatures) -> bool:
  """Check if a trait is supported by the device."""
 
- if (requires_dock_type := getattr(trait_type, "requires_dock_type", None)) is not None:
- return requires_dock_type(dock_type)
+ if (requires_dock_features := getattr(trait_type, "requires_dock_features", None)) is not None:
+ return requires_dock_features(dock_features)
 
  if (feature_name := getattr(trait_type, "requires_feature", None)) is None:
  _LOGGER.debug("Optional trait missing 'requires_feature' attribute %s, skipping", name)
@@ -307,16 +309,23 @@ async def _dock_type(self) -> RoborockDockTypeCode:
  if dock_type is not None:
  _LOGGER.debug("Using cached dock type: %s", dock_type)
  try:
-return RoborockDockTypeCode(dock_type)
+dock_type = RoborockDockTypeCode(dock_type)
  except ValueError:
  _LOGGER.debug("Cached dock type %s is invalid, refreshing", dock_type)
+ else:
+ if self.status.dss is None:
+ await self.status.refresh()
+ if self.status.dock_type is not None:
+ dock_type = self.status.dock_type
+ await self._set_cached_trait_data("dock_type", dock_type)
+ return dock_type
 
  _LOGGER.debug("Starting dock type discovery")
  await self.status.refresh()
  _LOGGER.debug("Fetched dock type: %s", self.status.dock_type)
  if self.status.dock_type is None:
  # Explicitly set so we reuse cached value next type
- dock_type = RoborockDockTypeCode.no_dock
+ dock_type = RoborockDockTypeCode.o0_dock
  else:
  dock_type = self.status.dock_type
  await self._set_cached_trait_data("dock_type", dock_type)