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)