feat: add better cleaning mode support (#817)

* feat: add better cleaning mode support

* chore: address comments from copilot

* chore: address comments from copilot

* chore: address comments from copilot

* chore: address comments

* chore: drop int typing and improve comment

* chore: add some comments that better explain 'get_mop_only_vacuum_mode'

Author: Lash-L

roborock/data/v1/v1_clean_modes.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
 import typing
+from enum import StrEnum
+from typing import TypeVar
 
+from ...exceptions import RoborockUnsupportedFeature
 from ..code_mappings import RoborockModeEnum
 
 if typing.TYPE_CHECKING:
@@ -68,6 +71,22 @@ class WashTowelModes(RoborockModeEnum):
     SUPER_DEEP = ("super_deep", 8)
 
 
+class CleaningMode(StrEnum):
+    """High-level cleaning intent derived from the lower-level motor settings.
+
+    Prefer this abstraction when you want to present or switch between the
+    user-facing cleaning behaviors exposed by the app. The lower-level
+    `VacuumModes`, `WaterModes`, and `CleanRoutes` enums are still useful for
+    fine-grained selection.
+    """
+
+    VACUUM = "vacuum"
+    VAC_AND_MOP = "vac_and_mop"
+    MOP = "mop"
+    CUSTOM = "custom"
+    SMART_MODE = "smart_mode"
+
+
 WATER_SLIDE_MODE_MAPPING: dict[int, WaterModes] = {
     200: WaterModes.OFF,
     221: WaterModes.PURE_WATER_FLOW_START,
@@ -78,6 +97,8 @@ class WashTowelModes(RoborockModeEnum):
     250: WaterModes.PURE_WATER_FLOW_END,
 }
 
+ModeEnumT = TypeVar("ModeEnumT", bound=RoborockModeEnum)
+
 
 def get_wash_towel_modes(features: DeviceFeatures) -> list[WashTowelModes]:
     """Get the valid wash towel modes for the device"""
@@ -174,7 +195,151 @@ def get_water_mode_mapping(features: DeviceFeatures) -> dict[int, str]:
     return {mode.code: mode.value for mode in get_water_modes(features)}
 
 
-def is_mode_customized(clean_mode: VacuumModes, water_mode: WaterModes, mop_mode: CleanRoutes) -> bool:
+def get_cleaning_mode_options(features: DeviceFeatures) -> list[CleaningMode]:
+    """Return the supported high-level cleaning modes for the device.
+
+    These options are the preferred user-facing choices because they bundle the
+    correct fan, water, and mop-route settings together for the device. Callers
+    should generally present these instead of mixing lower-level mode enums
+    unless they explicitly need fine-grained control.
+    """
+    if not features.is_support_water_mode:
+        return []
+
+    supported_water_modes = get_water_modes(features)
+    options = [CleaningMode.VACUUM, CleaningMode.VAC_AND_MOP]
+    if features.is_pure_clean_mop_supported:
+        options.append(CleaningMode.MOP)
+    if features.is_customized_clean_supported and WaterModes.CUSTOMIZED in supported_water_modes:
+        options.append(CleaningMode.CUSTOM)
+    if features.is_smart_clean_mode_set_supported and WaterModes.SMART_MODE in supported_water_modes:
+        options.append(CleaningMode.SMART_MODE)
+    return options
+
+
+def get_mop_only_vacuum_mode(features: DeviceFeatures) -> VacuumModes:
+    """Determine the vacuum mode to use when you just want to mop.
+
+    There are three cases that must be handled:
+    1. The device does not support only mopping.
+    2. The device supports raising the vacuum brush while mopping
+    3. All other cases.
+    """
+    if not features.is_pure_clean_mop_supported:
+        raise RoborockUnsupportedFeature("Mop-only cleaning is not supported")
+    if features.is_support_main_brush_up_down_supported:
+        return VacuumModes.OFF_RAISE_MAIN_BRUSH
+    return VacuumModes.OFF
+
+
+def _get_default_mopping_water_mode(features: DeviceFeatures) -> WaterModes:
+    """Pick a sensible default water mode when mopping for the device."""
+    # Water-slide devices use a disjoint set of water codes; pick a mid-flow
+    # slide code instead of the standard 202, which they don't accept.
+    if features.is_water_slide_mode_supported:
+        return WaterModes.PURE_WATER_FLOW_MIDDLE
+    return WaterModes.STANDARD
+
+
+def _get_clean_motor_mode_params(
+    mode: CleaningMode,
+    features: DeviceFeatures,
+) -> tuple[VacuumModes, WaterModes, CleanRoutes]:
+    """Return (fan_power, water_box_mode, mop_mode) enums for the high-level mode."""
+    if mode == CleaningMode.VACUUM:
+        return (VacuumModes.BALANCED, WaterModes.OFF, CleanRoutes.STANDARD)
+    if mode == CleaningMode.VAC_AND_MOP:
+        return (VacuumModes.BALANCED, _get_default_mopping_water_mode(features), CleanRoutes.STANDARD)
+    if mode == CleaningMode.MOP:
+        return (
+            get_mop_only_vacuum_mode(features),
+            _get_default_mopping_water_mode(features),
+            CleanRoutes.STANDARD,
+        )
+    if mode == CleaningMode.CUSTOM:
+        return (VacuumModes.CUSTOMIZED, WaterModes.CUSTOMIZED, CleanRoutes.CUSTOMIZED)
+    if mode == CleaningMode.SMART_MODE:
+        return (VacuumModes.SMART_MODE, WaterModes.SMART_MODE, CleanRoutes.SMART_MODE)
+    raise RoborockUnsupportedFeature(f"Cleaning mode {mode.value!r} is not supported")
+
+
+def resolve_cleaning_mode(cleaning_mode: str | CleaningMode) -> CleaningMode:
+    """Resolve a string or enum into a CleaningMode value."""
+    if isinstance(cleaning_mode, CleaningMode):
+        return cleaning_mode
+    try:
+        return CleaningMode(cleaning_mode)
+    except ValueError as err:
+        raise RoborockUnsupportedFeature(f"Cleaning mode {cleaning_mode!r} is not supported") from err
+
+
+def get_cleaning_mode_parameters(cleaning_mode: CleaningMode, features: DeviceFeatures) -> list[dict[str, int]]:
+    """Get the RPC payload for switching the high-level cleaning mode."""
+    if cleaning_mode not in get_cleaning_mode_options(features):
+        raise RoborockUnsupportedFeature(f"Cleaning mode {cleaning_mode.value!r} is not supported")
+
+    fan_power, water_box_mode, mop_mode = _get_clean_motor_mode_params(cleaning_mode, features)
+    params: dict[str, int] = {"fan_power": fan_power.code, "water_box_mode": water_box_mode.code}
+    if features.is_clean_route_setting_supported:
+        params["mop_mode"] = mop_mode.code
+    return [params]
+
+
+def _resolve_mode_code(value: int | ModeEnumT | None, mode_cls: type[ModeEnumT]) -> ModeEnumT | None:
+    """Resolve a raw code or enum into a RoborockModeEnum."""
+    if value is None:
+        return None
+    if isinstance(value, mode_cls):
+        return value
+    return mode_cls.from_code_optional(int(value))
+
+
+def _resolve_clean_mode(value: int | VacuumModes | None, features: DeviceFeatures) -> VacuumModes | None:
+    """Resolve a vacuum mode code, accounting for feature-specific code aliases."""
+    if value is None or isinstance(value, VacuumModes):
+        return value
+    if value == VacuumModes.OFF.code:
+        if features.is_pure_clean_mop_supported:
+            return get_mop_only_vacuum_mode(features)
+        return VacuumModes.GENTLE
+    return VacuumModes.from_code_optional(value)
+
+
+def get_current_cleaning_mode(
+    clean_mode: int | VacuumModes | None,
+    water_mode: int | WaterModes | None,
+    mop_mode: int | CleanRoutes | None,
+    features: DeviceFeatures,
+) -> CleaningMode | None:
+    """Classify the current high-level cleaning mode from individual mode codes."""
+    if not features.is_support_water_mode:
+        return None
+    clean_mode_enum = _resolve_clean_mode(clean_mode, features)
+    water_mode_enum = _resolve_mode_code(water_mode, WaterModes)
+    mop_mode_enum = _resolve_mode_code(mop_mode, CleanRoutes)
+    if clean_mode_enum is None or water_mode_enum is None:
+        return None
+
+    if is_smart_mode_set(water_mode_enum, clean_mode_enum, mop_mode_enum):
+        return CleaningMode.SMART_MODE
+    if is_mode_customized(clean_mode_enum, water_mode_enum, mop_mode_enum):
+        return CleaningMode.CUSTOM
+    if water_mode_enum != WaterModes.OFF:
+        try:
+            if clean_mode_enum == get_mop_only_vacuum_mode(features):
+                return CleaningMode.MOP
+        except RoborockUnsupportedFeature:
+            pass
+    if water_mode_enum == WaterModes.OFF:
+        return CleaningMode.VACUUM
+    return CleaningMode.VAC_AND_MOP
+
+
+def is_mode_customized(
+    clean_mode: VacuumModes | None,
+    water_mode: WaterModes | None,
+    mop_mode: CleanRoutes | None,
+) -> bool:
     """Check if any of the cleaning modes are set to a custom value."""
     return (
         clean_mode == VacuumModes.CUSTOMIZED
@@ -183,7 +348,11 @@ def is_mode_customized(clean_mode: VacuumModes, water_mode: WaterModes, mop_mode
     )
 
 
-def is_smart_mode_set(water_mode: WaterModes, clean_mode: VacuumModes, mop_mode: CleanRoutes) -> bool:
+def is_smart_mode_set(
+    water_mode: WaterModes | None,
+    clean_mode: VacuumModes | None,
+    mop_mode: CleanRoutes | None,
+) -> bool:
     """Check if the smart mode is set for the given water mode and clean mode"""
     return (
         water_mode == WaterModes.SMART_MODE

roborock/devices/traits/v1/status.py

@@ -3,14 +3,19 @@
 from typing import Any
 
 from roborock import (
+    CleaningMode,
     CleanRoutes,
     StatusV2,
     VacuumModes,
     WaterModes,
     get_clean_modes,
     get_clean_routes,
+    get_cleaning_mode_options,
+    get_cleaning_mode_parameters,
+    get_current_cleaning_mode,
     get_water_mode_mapping,
     get_water_modes,
+    resolve_cleaning_mode,
 )
 from roborock.devices.traits.common import DpsDataConverter, TraitUpdateListener
 from roborock.roborock_message import RoborockDataProtocol
@@ -42,10 +47,11 @@ class StatusTrait(StatusV2, common.V1TraitMixin, TraitUpdateListener):
     - Water Mode
     - Mop Route
 
-    You should call the _options() version of the attribute to know which are supported for your device
-    (i.e. fan_speed_options())
-    Then you can call the _mapping to convert an int value to the actual Enum. (i.e. fan_speed_mapping())
-    You can call the _name property to get the str value of the enum. (i.e. fan_speed_name)
+    You should use the _options version of the attribute to know which are
+    supported for your device (i.e. fan_speed_options)
+    Then you can use the _mapping to convert an int value to the actual Enum.
+    (i.e. fan_speed_mapping)
+    You can use the _name property to get the str value of the enum. (i.e. fan_speed_name)
 
     """
 
@@ -83,6 +89,10 @@ def mop_route_options(self) -> list[CleanRoutes]:
     def mop_route_mapping(self) -> dict[int, str]:
         return {route.code: route.value for route in self.mop_route_options}
 
+    @cached_property
+    def cleaning_mode_options(self) -> list[CleaningMode]:
+        return get_cleaning_mode_options(self._device_features_trait)
+
     @property
     def fan_speed_name(self) -> str | None:
         if self.fan_power is None:
@@ -101,6 +111,28 @@ def mop_route_name(self) -> str | None:
             return None
         return self.mop_route_mapping.get(self.mop_mode)
 
+    @property
+    def current_cleaning_mode(self) -> CleaningMode | None:
+        return get_current_cleaning_mode(
+            clean_mode=self.fan_power,
+            water_mode=self.water_box_mode,
+            mop_mode=self.mop_mode,
+            features=self._device_features_trait,
+        )
+
+    @property
+    def current_cleaning_mode_name(self) -> str | None:
+        if (cleaning_mode := self.current_cleaning_mode) is None:
+            return None
+        return cleaning_mode.value
+
+    async def set_cleaning_mode(self, cleaning_mode: str | CleaningMode) -> None:
+        """Set the preferred high-level cleaning mode for the device."""
+        await self.rpc_channel.send_command(
+            RoborockCommand.SET_CLEAN_MOTOR_MODE,
+            params=get_cleaning_mode_parameters(resolve_cleaning_mode(cleaning_mode), self._device_features_trait),
+        )
+
     def update_from_dps(self, decoded_dps: dict[RoborockDataProtocol, Any]) -> None:
         """Update the trait from data protocol push message data.
 

tests/devices/__snapshots__/test_v1_device.ambr

@@ -870,7 +870,7 @@
   })
 # ---
 # name: test_device_trait_command_parsing[status]
-  StatusTrait(adbumper_status=None, auto_dust_collection=None, avoid_count=None, back_type=None, battery=100, camera_status=None, charge_status=None, clean_area=91287500, clean_fluid_status=None, clean_percent=None, clean_time=5405, clear_water_box_status=None, collision_avoid_status=None, command=<RoborockCommand.GET_STATUS: 'get_status'>, common_status=None, converter=DefaultConverter, corner_clean_mode=None, current_map=0, debug_mode=None, dirty_water_box_status=None, distance_off=0, dnd_enabled=1, dock_cool_fan_status=None, dock_error_status=None, dock_state=<RoborockDockState.full: 'full'>, dock_type=None, dry_status=None, dss=None, dust_bag_status=None, dust_collection_status=None, error_code=<RoborockErrorCode.none: 0>, error_code_name='none', fan_power=106, fan_speed_mapping={101: 'quiet', 102: 'balanced', 103: 'turbo', 104: 'max', 108: 'max_plus', 105: 'off', 106: 'custom'}, fan_speed_name='custom', fan_speed_options=[<VacuumModes.QUIET: 'quiet'>, <VacuumModes.BALANCED: 'balanced'>, <VacuumModes.TURBO: 'turbo'>, <VacuumModes.MAX: 'max'>, <VacuumModes.MAX_PLUS: 'max_plus'>, <VacuumModes.OFF: 'off'>, <VacuumModes.CUSTOMIZED: 'custom'>], hatch_door_status=None, home_sec_enable_password=None, home_sec_status=None, in_cleaning=<RoborockInCleaning.complete: 0>, in_fresh_state=1, in_returning=0, in_warmup=None, is_exploring=None, is_locating=0, kct=None, lab_status=1, last_clean_t=None, lock_status=0, map_present=1, map_status=3, mop_forbidden_enable=0, mop_mode=None, mop_route_mapping={300: 'standard', 301: 'deep', 302: 'custom'}, mop_route_name=None, mop_route_options=[<CleanRoutes.STANDARD: 'standard'>, <CleanRoutes.DEEP: 'deep'>, <CleanRoutes.CUSTOMIZED: 'custom'>], msg_seq=515, msg_ver=2, rdt=None, repeat=None, replenish_mode=None, rss=None, square_meter_clean_area=91.3, state=<RoborockStateCode.charging: 8>, state_name='charging', subdivision_sets=None, switch_map_mode=None, unsave_map_flag=0, unsave_map_reason=4, wash_phase=None, wash_ready=None, wash_status=None, water_box_carriage_status=0, water_box_filter_status=None, water_box_mode=204, water_box_status=0, water_mode_mapping={200: 'off', 201: 'mild', 202: 'standard', 203: 'intense', 207: 'custom_water_flow', 204: 'custom'}, water_mode_name='custom', water_mode_options=[<WaterModes.OFF: 'off'>, <WaterModes.MILD: 'mild'>, <WaterModes.STANDARD: 'standard'>, <WaterModes.INTENSE: 'intense'>, <WaterModes.CUSTOM: 'custom_water_flow'>, <WaterModes.CUSTOMIZED: 'custom'>], water_shortage_status=None)
+  StatusTrait(adbumper_status=None, auto_dust_collection=None, avoid_count=None, back_type=None, battery=100, camera_status=None, charge_status=None, clean_area=91287500, clean_fluid_status=None, clean_percent=None, clean_time=5405, cleaning_mode_options=[<CleaningMode.VACUUM: 'vacuum'>, <CleaningMode.VAC_AND_MOP: 'vac_and_mop'>, <CleaningMode.MOP: 'mop'>, <CleaningMode.CUSTOM: 'custom'>], clear_water_box_status=None, collision_avoid_status=None, command=<RoborockCommand.GET_STATUS: 'get_status'>, common_status=None, converter=DefaultConverter, corner_clean_mode=None, current_cleaning_mode=<CleaningMode.CUSTOM: 'custom'>, current_cleaning_mode_name='custom', current_map=0, debug_mode=None, dirty_water_box_status=None, distance_off=0, dnd_enabled=1, dock_cool_fan_status=None, dock_error_status=None, dock_state=<RoborockDockState.full: 'full'>, dock_type=None, dry_status=None, dss=None, dust_bag_status=None, dust_collection_status=None, error_code=<RoborockErrorCode.none: 0>, error_code_name='none', fan_power=106, fan_speed_mapping={101: 'quiet', 102: 'balanced', 103: 'turbo', 104: 'max', 108: 'max_plus', 105: 'off', 106: 'custom'}, fan_speed_name='custom', fan_speed_options=[<VacuumModes.QUIET: 'quiet'>, <VacuumModes.BALANCED: 'balanced'>, <VacuumModes.TURBO: 'turbo'>, <VacuumModes.MAX: 'max'>, <VacuumModes.MAX_PLUS: 'max_plus'>, <VacuumModes.OFF: 'off'>, <VacuumModes.CUSTOMIZED: 'custom'>], hatch_door_status=None, home_sec_enable_password=None, home_sec_status=None, in_cleaning=<RoborockInCleaning.complete: 0>, in_fresh_state=1, in_returning=0, in_warmup=None, is_exploring=None, is_locating=0, kct=None, lab_status=1, last_clean_t=None, lock_status=0, map_present=1, map_status=3, mop_forbidden_enable=0, mop_mode=None, mop_route_mapping={300: 'standard', 301: 'deep', 302: 'custom'}, mop_route_name=None, mop_route_options=[<CleanRoutes.STANDARD: 'standard'>, <CleanRoutes.DEEP: 'deep'>, <CleanRoutes.CUSTOMIZED: 'custom'>], msg_seq=515, msg_ver=2, rdt=None, repeat=None, replenish_mode=None, rss=None, square_meter_clean_area=91.3, state=<RoborockStateCode.charging: 8>, state_name='charging', subdivision_sets=None, switch_map_mode=None, unsave_map_flag=0, unsave_map_reason=4, wash_phase=None, wash_ready=None, wash_status=None, water_box_carriage_status=0, water_box_filter_status=None, water_box_mode=204, water_box_status=0, water_mode_mapping={200: 'off', 201: 'mild', 202: 'standard', 203: 'intense', 207: 'custom_water_flow', 204: 'custom'}, water_mode_name='custom', water_mode_options=[<WaterModes.OFF: 'off'>, <WaterModes.MILD: 'mild'>, <WaterModes.STANDARD: 'standard'>, <WaterModes.INTENSE: 'intense'>, <WaterModes.CUSTOM: 'custom_water_flow'>, <WaterModes.CUSTOMIZED: 'custom'>], water_shortage_status=None)
 # ---
 # name: test_device_trait_command_parsing[status].1
   dict({