feat: make status dynamic (#611)

* chore: do dynamic status

* feat: add more data and region

* chore: change str

* chore: add testing

* chore: add hash

* chore: add dss to status

* chore: update e2e

* chore: docstring

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* chore: address comments

* chore: add some docs and a basic test

* chore: add warning about hash for RoborockModeEnum

* chore: add more info about the dynamic attributes

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: Lash-L

roborock/data/code_mappings.py

@@ -3,7 +3,7 @@
 import logging
 from collections import namedtuple
 from enum import Enum, IntEnum, StrEnum
-from typing import Self
+from typing import Any, Self
 
 _LOGGER = logging.getLogger(__name__)
 completed_warnings = set()
@@ -105,6 +105,21 @@ def keys(cls) -> list[str]:
         """Returns a list of all member values."""
         return [member.value for member in cls]
 
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, str):
+            return self.value == other or self.name == other
+        if isinstance(other, int):
+            return self.code == other
+        return super().__eq__(other)
+
+    def __hash__(self) -> int:
+        """Hash a RoborockModeEnum.
+
+        It is critical that you do not mix RoborockModeEnums with raw strings or ints in hashed situations
+        (i.e. sets or keys in dictionaries)
+        """
+        return hash((self.code, self._value_))
+
 
 ProductInfo = namedtuple("ProductInfo", ["nickname", "short_models"])
 

roborock/data/v1/v1_clean_modes.py

@@ -16,6 +16,8 @@ class VacuumModes(RoborockModeEnum):
     TURBO = ("turbo", 103)
     MAX = ("max", 104)
     MAX_PLUS = ("max_plus", 108)
+    CARPET = ("carpet", 107)
+    OFF_RAISE_MAIN_BRUSH = ("off_raise_main_brush", 109)
     CUSTOMIZED = ("custom", 106)
     SMART_MODE = ("smart_mode", 110)
 
@@ -45,6 +47,8 @@ class WaterModes(RoborockModeEnum):
     STANDARD = ("standard", 202)
     HIGH = ("high", 203)
     INTENSE = ("intense", 203)
+    MIN = ("min", 205)
+    MAX = ("max", 206)
     CUSTOMIZED = ("custom", 204)
     CUSTOM = ("custom_water_flow", 207)
     EXTREME = ("extreme", 208)
@@ -81,9 +85,14 @@ def get_clean_modes(features: DeviceFeatures) -> list[VacuumModes]:
     if features.is_max_plus_mode_supported or features.is_none_pure_clean_mop_with_max_plus:
         # If the vacuum has max plus mode supported
         modes.append(VacuumModes.MAX_PLUS)
+    if features.is_carpet_deep_clean_supported:
+        modes.append(VacuumModes.CARPET)
     if features.is_pure_clean_mop_supported:
         # If the vacuum is capable of 'pure mop clean' aka no vacuum
-        modes.append(VacuumModes.OFF)
+        if features.is_support_main_brush_up_down_supported:
+            modes.append(VacuumModes.OFF_RAISE_MAIN_BRUSH)
+        else:
+            modes.append(VacuumModes.OFF)
     else:
         # If not, we can add gentle
         modes.append(VacuumModes.GENTLE)

roborock/data/v1/v1_containers.py

@@ -121,6 +121,8 @@ def _requires_schema_code(requires_schema_code: str, default=None) -> Any:
 
 @dataclass
 class Status(RoborockBase):
+    """This status will be deprecated in favor of StatusV2."""
+
     msg_ver: int | None = None
     msg_seq: int | None = None
     state: RoborockStateCode | None = _requires_schema_code("state", default=None)
@@ -282,6 +284,140 @@ def __repr__(self) -> str:
         return _attr_repr(self)
 
 
+@dataclass
+class StatusV2(RoborockBase):
+    """
+    This is a new version of the Status object.
+    This is the result of GET_STATUS from the api.
+    """
+
+    msg_ver: int | None = None
+    msg_seq: int | None = None
+    state: RoborockStateCode | None = None
+    battery: int | None = None
+    clean_time: int | None = None
+    clean_area: int | None = None
+    error_code: RoborockErrorCode | None = None
+    map_present: int | None = None
+    in_cleaning: RoborockInCleaning | None = None
+    in_returning: int | None = None
+    in_fresh_state: int | None = None
+    lab_status: int | None = None
+    water_box_status: int | None = None
+    back_type: int | None = None
+    wash_phase: int | None = None
+    wash_ready: int | None = None
+    fan_power: int | None = None
+    dnd_enabled: int | None = None
+    map_status: int | None = None
+    is_locating: int | None = None
+    lock_status: int | None = None
+    water_box_mode: int | None = None
+    water_box_carriage_status: int | None = None
+    mop_forbidden_enable: int | None = None
+    camera_status: int | None = None
+    is_exploring: int | None = None
+    home_sec_status: int | None = None
+    home_sec_enable_password: int | None = None
+    adbumper_status: list[int] | None = None
+    water_shortage_status: int | None = None
+    dock_type: RoborockDockTypeCode | None = None
+    dust_collection_status: int | None = None
+    auto_dust_collection: int | None = None
+    avoid_count: int | None = None
+    mop_mode: int | None = None
+    debug_mode: int | None = None
+    collision_avoid_status: int | None = None
+    switch_map_mode: int | None = None
+    dock_error_status: RoborockDockErrorCode | None = None
+    charge_status: int | None = 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
+    rdt: int | None = None
+    clean_percent: int | None = None
+    rss: int | None = None
+    dss: int | None = None
+    common_status: int | None = None
+    corner_clean_mode: int | None = None
+    last_clean_t: int | None = None
+    replenish_mode: int | None = None
+    repeat: int | None = None
+    kct: int | None = None
+    subdivision_sets: int | None = None
+
+    @property
+    def square_meter_clean_area(self) -> float | None:
+        return round(self.clean_area / 1000000, 1) if self.clean_area is not None else None
+
+    @property
+    def error_code_name(self) -> str | None:
+        return self.error_code.name if self.error_code is not None else None
+
+    @property
+    def state_name(self) -> str | None:
+        return self.state.name if self.state is not None else None
+
+    @property
+    def current_map(self) -> int | None:
+        """Returns the current map ID if the map is present."""
+        if self.map_status is not None:
+            map_flag = self.map_status >> 2
+            if map_flag != NO_MAP:
+                return map_flag
+        return None
+
+    @property
+    def clear_water_box_status(self) -> ClearWaterBoxStatus | None:
+        if self.dss:
+            return ClearWaterBoxStatus((self.dss >> 2) & 3)
+        return None
+
+    @property
+    def dirty_water_box_status(self) -> DirtyWaterBoxStatus | None:
+        if self.dss:
+            return DirtyWaterBoxStatus((self.dss >> 4) & 3)
+        return None
+
+    @property
+    def dust_bag_status(self) -> DustBagStatus | None:
+        if self.dss:
+            return DustBagStatus((self.dss >> 6) & 3)
+        return None
+
+    @property
+    def water_box_filter_status(self) -> int | None:
+        if self.dss:
+            return (self.dss >> 8) & 3
+        return None
+
+    @property
+    def clean_fluid_status(self) -> CleanFluidStatus | None:
+        if self.dss:
+            value = (self.dss >> 10) & 3
+            if value == 0:
+                return None  # Feature not supported by this device
+        return None
+
+    @property
+    def hatch_door_status(self) -> int | None:
+        if self.dss:
+            return (self.dss >> 12) & 7
+        return None
+
+    @property
+    def dock_cool_fan_status(self) -> int | None:
+        if self.dss:
+            return (self.dss >> 15) & 3
+        return None
+
+    def __repr__(self) -> str:
+        return _attr_repr(self)
+
+
 @dataclass
 class S4MaxStatus(Status):
     fan_power: RoborockFanSpeedS6Pure | None = None

roborock/devices/device_manager.py

@@ -239,6 +239,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                     web_api,
                     device_cache=device_cache,
                     map_parser_config=map_parser_config,
+                    region=user_data.region,
                 )
             case DeviceVersion.A01:
                 channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)

roborock/devices/traits/v1/__init__.py

@@ -54,7 +54,6 @@
 
 import logging
 from dataclasses import dataclass, field, fields
-from functools import cache
 from typing import Any, get_args
 
 from roborock.data.containers import HomeData, HomeDataProduct, RoborockBase
@@ -180,6 +179,7 @@ def __init__(
         web_api: UserWebApiClient,
         device_cache: DeviceCache,
         map_parser_config: MapParserConfig | None = None,
+        region: str | None = None,
     ) -> None:
         """Initialize the V1TraitProps."""
         self._device_uid = device_uid
@@ -188,14 +188,15 @@ def __init__(
         self._map_rpc_channel = map_rpc_channel
         self._web_api = web_api
         self._device_cache = device_cache
+        self._region = region
 
-        self.status = StatusTrait(product)
+        self.device_features = DeviceFeaturesTrait(product, self._device_cache)
+        self.status = StatusTrait(self.device_features, region=self._region)
         self.consumables = ConsumableTrait()
         self.rooms = RoomsTrait(home_data)
         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, self._device_cache)
         self.network_info = NetworkInfoTrait(device_uid, self._device_cache)
         self.routines = RoutinesTrait(device_uid, web_api)
 
@@ -206,6 +207,8 @@ def __init__(
                 if (union_args := get_args(item.type)) is None or len(union_args) > 0:
                     continue
                 _LOGGER.debug("Trait '%s' is supported, initializing", item.name)
+                if not callable(item.type):
+                    continue
                 trait = item.type()
                 setattr(self, item.name, trait)
             # This is a hack to allow setting the rpc_channel on all traits. This is
@@ -324,6 +327,7 @@ def create(
     web_api: UserWebApiClient,
     device_cache: DeviceCache,
     map_parser_config: MapParserConfig | None = None,
+    region: str | None = None,
 ) -> PropertiesApi:
     """Create traits for V1 devices."""
     return PropertiesApi(
@@ -336,4 +340,5 @@ def create(
         web_api,
         device_cache,
         map_parser_config,
+        region=region,
     )

roborock/devices/traits/v1/status.py

@@ -1,24 +1,92 @@
+from functools import cached_property
 from typing import Self
 
-from roborock.data import HomeDataProduct, ModelStatus, S7MaxVStatus, Status
-from roborock.devices.traits.v1 import common
+from roborock import CleanRoutes, StatusV2, VacuumModes, WaterModes, get_clean_modes, get_clean_routes, get_water_modes
 from roborock.roborock_typing import RoborockCommand
 
+from . import common
+from .device_features import DeviceFeaturesTrait
 
-class StatusTrait(Status, common.V1TraitMixin):
-    """Trait for managing the status of Roborock devices."""
+
+class StatusTrait(StatusV2, common.V1TraitMixin):
+    """Trait for managing the status of Roborock devices.
+
+    The StatusTrait gives you the access to the state of a Roborock vacuum.
+    The various attribute options on state change per each device.
+    Values like fan speed, mop mode, etc. have different options for every device
+    and are dynamically determined.
+
+    Usage:
+        Before accessing status properties, you should call `refresh()` to fetch
+        the latest data from the device. You must pass in the device feature trait
+        to this trait so that the dynamic attributes can be pre-determined.
+
+    The current dynamic attributes are:
+    - Fan Speed
+    - 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)
+
+    """
 
     command = RoborockCommand.GET_STATUS
 
-    def __init__(self, product_info: HomeDataProduct) -> None:
+    def __init__(self, device_feature_trait: DeviceFeaturesTrait, region: str | None = None) -> None:
         """Initialize the StatusTrait."""
-        self._product_info = product_info
+        super().__init__()
+        self._device_features_trait = device_feature_trait
+        self._region = region
+
+    @cached_property
+    def fan_speed_options(self) -> list[VacuumModes]:
+        return get_clean_modes(self._device_features_trait)
+
+    @cached_property
+    def fan_speed_mapping(self) -> dict[int, str]:
+        return {fan.code: fan.value for fan in self.fan_speed_options}
+
+    @cached_property
+    def water_mode_options(self) -> list[WaterModes]:
+        return get_water_modes(self._device_features_trait)
+
+    @cached_property
+    def water_mode_mapping(self) -> dict[int, str]:
+        return {mop.code: mop.value for mop in self.water_mode_options}
+
+    @cached_property
+    def mop_route_options(self) -> list[CleanRoutes]:
+        return get_clean_routes(self._device_features_trait, self._region or "us")
+
+    @cached_property
+    def mop_route_mapping(self) -> dict[int, str]:
+        return {route.code: route.value for route in self.mop_route_options}
+
+    @property
+    def fan_speed_name(self) -> str | None:
+        if self.fan_power is None:
+            return None
+        return self.fan_speed_mapping.get(self.fan_power)
+
+    @property
+    def water_mode_name(self) -> str | None:
+        if self.water_box_mode is None:
+            return None
+        return self.water_mode_mapping.get(self.water_box_mode)
+
+    @property
+    def mop_route_name(self) -> str | None:
+        if self.mop_mode is None:
+            return None
+        return self.mop_route_mapping.get(self.mop_mode)
 
     def _parse_response(self, response: common.V1ResponseData) -> Self:
-        """Parse the response from the device into a CleanSummary."""
-        status_type: type[Status] = ModelStatus.get(self._product_info.model, S7MaxVStatus)
+        """Parse the response from the device into a StatusV2-based status object."""
         if isinstance(response, list):
             response = response[0]
         if isinstance(response, dict):
-            return status_type.from_dict(response)
+            return StatusV2.from_dict(response)
         raise ValueError(f"Unexpected status format: {response!r}")