feat: add better cleaning mode support

roborock/data/v1/v1_clean_modes.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import typing
+from enum import StrEnum
 
+from ...exceptions import RoborockUnsupportedFeature
 from ..code_mappings import RoborockModeEnum
 
 if typing.TYPE_CHECKING:
@@ -68,6 +70,14 @@ class WashTowelModes(RoborockModeEnum):
     SUPER_DEEP = ("super_deep", 8)
 
 
+class CleaningModes(StrEnum):
+    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,
@@ -174,7 +184,105 @@ 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[CleaningModes]:
+    """Get the supported high-level cleaning modes for the device."""
+    if not features.is_support_water_mode:
+        return []
+
+    options = [CleaningModes.VACUUM, CleaningModes.VAC_AND_MOP]
+    if features.is_pure_clean_mop_supported:
+        options.append(CleaningModes.MOP)
+    if features.is_customized_clean_supported:
+        options.append(CleaningModes.CUSTOM)
+    if features.is_smart_clean_mode_set_supported:
+        options.append(CleaningModes.SMART_MODE)
+    return options
+
+
+def get_mop_only_vacuum_mode(features: DeviceFeatures) -> VacuumModes:
+    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_code(features: DeviceFeatures) -> int:
+    """Pick a sensible default water code 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.code
+    return WaterModes.STANDARD.code
+
+
+def _get_clean_motor_mode_params(mode: CleaningModes, features: DeviceFeatures) -> tuple[int, int, int]:
+    """Return (fan_power, water_box_mode, mop_mode) codes for the high-level mode."""
+    if mode == CleaningModes.VACUUM:
+        return (VacuumModes.BALANCED.code, WaterModes.OFF.code, CleanRoutes.STANDARD.code)
+    if mode == CleaningModes.VAC_AND_MOP:
+        return (VacuumModes.BALANCED.code, _get_default_mopping_water_code(features), CleanRoutes.STANDARD.code)
+    if mode == CleaningModes.MOP:
+        return (
+            get_mop_only_vacuum_mode(features).code,
+            _get_default_mopping_water_code(features),
+            CleanRoutes.STANDARD.code,
+        )
+    if mode == CleaningModes.CUSTOM:
+        return (VacuumModes.CUSTOMIZED.code, WaterModes.CUSTOMIZED.code, CleanRoutes.CUSTOMIZED.code)
+    if mode == CleaningModes.SMART_MODE:
+        return (VacuumModes.SMART_MODE.code, WaterModes.SMART_MODE.code, CleanRoutes.SMART_MODE.code)
+    raise RoborockUnsupportedFeature(f"Cleaning mode {mode.value!r} is not supported")
+
+
+def get_cleaning_mode_parameters(cleaning_mode: str | CleaningModes, features: DeviceFeatures) -> list[dict[str, int]]:
+    """Get the RPC payload for switching the high-level cleaning mode."""
+    try:
+        mode = CleaningModes(cleaning_mode)
+    except ValueError as err:
+        raise RoborockUnsupportedFeature(f"Cleaning mode {cleaning_mode!r} is not supported") from err
+    if mode not in get_cleaning_mode_options(features):
+        raise RoborockUnsupportedFeature(f"Cleaning mode {mode.value!r} is not supported")
+
+    fan_power, water_box_mode, mop_mode = _get_clean_motor_mode_params(mode, features)
+    params: dict[str, int] = {"fan_power": fan_power, "water_box_mode": water_box_mode}
+    if features.is_clean_route_setting_supported:
+        params["mop_mode"] = mop_mode
+    return [params]
+
+
+def get_current_cleaning_mode(
+    clean_mode: int | None,
+    water_mode: int | None,
+    mop_mode: int | None,
+    features: DeviceFeatures,
+) -> CleaningModes | None:
+    """Classify the current high-level cleaning mode from individual mode codes."""
+    if not features.is_support_water_mode:
+        return None
+    if clean_mode is None or water_mode is None:
+        return None
+
+    if is_smart_mode_set(water_mode, clean_mode, mop_mode):
+        return CleaningModes.SMART_MODE
+    if is_mode_customized(clean_mode, water_mode, mop_mode):
+        return CleaningModes.CUSTOM
+    if water_mode != WaterModes.OFF.code:
+        try:
+            if clean_mode == get_mop_only_vacuum_mode(features).code:
+                return CleaningModes.MOP
+        except RoborockUnsupportedFeature:
+            pass
+    if water_mode == WaterModes.OFF.code:
+        return CleaningModes.VACUUM
+    return CleaningModes.VAC_AND_MOP
+
+
+def is_mode_customized(
+    clean_mode: int | VacuumModes | None,
+    water_mode: int | WaterModes | None,
+    mop_mode: int | CleanRoutes | None,
+) -> bool:
     """Check if any of the cleaning modes are set to a custom value."""
     return (
         clean_mode == VacuumModes.CUSTOMIZED
@@ -183,7 +291,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: int | WaterModes | None,
+    clean_mode: int | VacuumModes | None,
+    mop_mode: int | 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

@@ -1,12 +1,16 @@
 from functools import cached_property
 
 from roborock import (
+    CleaningModes,
     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,
 )
@@ -34,10 +38,11 @@ class StatusTrait(StatusV2, common.V1TraitMixin):
     - 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)
 
     """
 
@@ -74,6 +79,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[CleaningModes]:
+        return get_cleaning_mode_options(self._device_features_trait)
+
     @property
     def fan_speed_name(self) -> str | None:
         if self.fan_power is None:
@@ -91,3 +100,29 @@ def mop_route_name(self) -> str | None:
         if self.mop_mode is None:
             return None
         return self.mop_route_mapping.get(self.mop_mode)
+
+    @property
+    def current_cleaning_mode(self) -> CleaningModes | 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 cleaning_mode_name(self) -> str | None:
+        if (cleaning_mode := self.current_cleaning_mode) is None:
+            return None
+        return cleaning_mode.value
+
+    def get_cleaning_mode_parameters(self, cleaning_mode: str | CleaningModes) -> list[dict[str, int]]:
+        """Get the RPC payload for the selected high-level cleaning mode."""
+        return get_cleaning_mode_parameters(cleaning_mode, self._device_features_trait)
+
+    async def set_cleaning_mode(self, cleaning_mode: str | CleaningModes) -> None:
+        """Set the high-level cleaning mode."""
+        await self.rpc_channel.send_command(
+            RoborockCommand.SET_CLEAN_MOTOR_MODE,
+            params=self.get_cleaning_mode_parameters(cleaning_mode),
+        )

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_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_name='custom', cleaning_mode_options=[<CleaningModes.VACUUM: 'vacuum'>, <CleaningModes.VAC_AND_MOP: 'vac_and_mop'>, <CleaningModes.MOP: 'mop'>, <CleaningModes.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=<CleaningModes.CUSTOM: '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_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({