feat: add device_features to automatically determine what is supported (#365)

* feat: add device_features to automatically determine what is supported

* chore: update to the version from the other PR

* chore: move device_features to seperate file and add some tests and rework device_features

Author: Lash-L

roborock/code_mappings.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from collections import namedtuple
 from enum import Enum, IntEnum
 
 _LOGGER = logging.getLogger(__name__)
@@ -50,6 +51,68 @@ def items(cls: type[RoborockEnum]):
         return cls.as_dict().items()
 
 
+ProductInfo = namedtuple("ProductInfo", ["nickname", "short_models"])
+
+
+class RoborockProductNickname(Enum):
+    # Coral Series
+    CORAL = ProductInfo(nickname="Coral", short_models=("a20", "a21"))
+    CORALPRO = ProductInfo(nickname="CoralPro", short_models=("a143", "a144"))
+
+    # Pearl Series
+    PEARL = ProductInfo(nickname="Pearl", short_models=("a74", "a75"))
+    PEARLC = ProductInfo(nickname="PearlC", short_models=("a103", "a104"))
+    PEARLE = ProductInfo(nickname="PearlE", short_models=("a167", "a168"))
+    PEARLELITE = ProductInfo(nickname="PearlELite", short_models=("a169", "a170"))
+    PEARLPLUS = ProductInfo(nickname="PearlPlus", short_models=("a86", "a87"))
+    PEARLPLUSS = ProductInfo(nickname="PearlPlusS", short_models=("a116", "a117", "a136"))
+    PEARLS = ProductInfo(nickname="PearlS", short_models=("a100", "a101"))
+    PEARLSLITE = ProductInfo(nickname="PearlSLite", short_models=("a122", "a123"))
+
+    # Ruby Series
+    RUBYPLUS = ProductInfo(nickname="RubyPlus", short_models=("t4", "s4"))
+    RUBYSC = ProductInfo(nickname="RubySC", short_models=("p5", "a08"))
+    RUBYSE = ProductInfo(nickname="RubySE", short_models=("a19",))
+    RUBYSLITE = ProductInfo(nickname="RubySLite", short_models=("p6", "s5e", "a05"))
+
+    # Tanos Series
+    TANOS = ProductInfo(nickname="Tanos", short_models=("t6", "s6"))
+    TANOSE = ProductInfo(nickname="TanosE", short_models=("t7", "a11"))
+    TANOSS = ProductInfo(nickname="TanosS", short_models=("a14", "a15"))
+    TANOSSC = ProductInfo(nickname="TanosSC", short_models=("a39", "a40"))
+    TANOSSE = ProductInfo(nickname="TanosSE", short_models=("a33", "a34"))
+    TANOSSMAX = ProductInfo(nickname="TanosSMax", short_models=("a52",))
+    TANOSSLITE = ProductInfo(nickname="TanosSLite", short_models=("a37", "a38"))
+    TANOSSPLUS = ProductInfo(nickname="TanosSPlus", short_models=("a23", "a24"))
+    TANOSV = ProductInfo(nickname="TanosV", short_models=("t7p", "a09", "a10"))
+
+    # Topaz Series
+    TOPAZS = ProductInfo(nickname="TopazS", short_models=("a29", "a30", "a76"))
+    TOPAZSC = ProductInfo(nickname="TopazSC", short_models=("a64", "a65"))
+    TOPAZSPLUS = ProductInfo(nickname="TopazSPlus", short_models=("a46", "a47", "a66"))
+    TOPAZSPOWER = ProductInfo(nickname="TopazSPower", short_models=("a62",))
+    TOPAZSV = ProductInfo(nickname="TopazSV", short_models=("a26", "a27"))
+
+    # Ultron Series
+    ULTRON = ProductInfo(nickname="Ultron", short_models=("a50", "a51"))
+    ULTRONE = ProductInfo(nickname="UltronE", short_models=("a72", "a84"))
+    ULTRONLITE = ProductInfo(nickname="UltronLite", short_models=("a73", "a85"))
+    ULTRONSC = ProductInfo(nickname="UltronSC", short_models=("a94", "a95"))
+    ULTRONSE = ProductInfo(nickname="UltronSE", short_models=("a124", "a125", "a139", "a140"))
+    ULTRONSPLUS = ProductInfo(nickname="UltronSPlus", short_models=("a68", "a69", "a70"))
+    ULTRONSV = ProductInfo(nickname="UltronSV", short_models=("a96", "a97"))
+
+    # Verdelite Series
+    VERDELITE = ProductInfo(nickname="Verdelite", short_models=("a146", "a147"))
+
+    # Vivian Series
+    VIVIAN = ProductInfo(nickname="Vivian", short_models=("a134", "a135", "a155", "a156"))
+    VIVIANC = ProductInfo(nickname="VivianC", short_models=("a158", "a159"))
+
+
+SHORT_MODEL_TO_ENUM = {model: product for product in RoborockProductNickname for model in product.value.short_models}
+
+
 class RoborockStateCode(RoborockEnum):
     unknown = 0
     starting = 1

roborock/containers.py

@@ -11,6 +11,7 @@
 from typing import Any, NamedTuple, get_args, get_origin
 
 from .code_mappings import (
+    SHORT_MODEL_TO_ENUM,
     RoborockCategory,
     RoborockCleanType,
     RoborockDockDustCollectionModeCode,
@@ -53,6 +54,7 @@
     RoborockMopModeS8ProUltra,
     RoborockMopModeSaros10,
     RoborockMopModeSaros10R,
+    RoborockProductNickname,
     RoborockStartType,
     RoborockStateCode,
 )
@@ -87,6 +89,7 @@
     STRAINER_REPLACE_TIME,
     ROBOROCK_G20S_Ultra,
 )
+from .device_features import DeviceFeatures
 from .exceptions import RoborockException
 
 _LOGGER = logging.getLogger(__name__)
@@ -306,144 +309,6 @@ class HomeDataDevice(RoborockBase):
     silent_ota_switch: bool | None = None
     setting: Any | None = None
     f: bool | None = None
-    device_features: DeviceFeatures | None = None
-
-    # seemingly not just str like I thought - example: '0000000000002000' and '0000000000002F63'
-
-    # def __post_init__(self):
-    #     if self.feature_set is not None and self.new_feature_set is not None and self.new_feature_set != "":
-    #         self.device_features = build_device_features(self.feature_set, self.new_feature_set)
-
-
-@dataclass
-class DeviceFeatures(RoborockBase):
-    map_carpet_add_supported: bool
-    show_clean_finish_reason_supported: bool
-    resegment_supported: bool
-    video_monitor_supported: bool
-    any_state_transit_goto_supported: bool
-    fw_filter_obstacle_supported: bool
-    video_setting_supported: bool
-    ignore_unknown_map_object_supported: bool
-    set_child_supported: bool
-    carpet_supported: bool
-    mop_path_supported: bool
-    multi_map_segment_timer_supported: bool
-    custom_water_box_distance_supported: bool
-    wash_then_charge_cmd_supported: bool
-    room_name_supported: bool
-    current_map_restore_enabled: bool
-    photo_upload_supported: bool
-    shake_mop_set_supported: bool
-    map_beautify_internal_debug_supported: bool
-    new_data_for_clean_history: bool
-    new_data_for_clean_history_detail: bool
-    flow_led_setting_supported: bool
-    dust_collection_setting_supported: bool
-    rpc_retry_supported: bool
-    avoid_collision_supported: bool
-    support_set_switch_map_mode: bool
-    support_smart_scene: bool
-    support_floor_edit: bool
-    support_furniture: bool
-    support_room_tag: bool
-    support_quick_map_builder: bool
-    support_smart_global_clean_with_custom_mode: bool
-    record_allowed: bool
-    careful_slow_map_supported: bool
-    egg_mode_supported: bool
-    unsave_map_reason_supported: bool
-    carpet_show_on_map: bool
-    supported_valley_electricity: bool
-    drying_supported: bool
-    download_test_voice_supported: bool
-    support_backup_map: bool
-    support_custom_mode_in_cleaning: bool
-    support_remote_control_in_call: bool
-    support_set_volume_in_call: bool
-    support_clean_estimate: bool
-    support_custom_dnd: bool
-    carpet_deep_clean_supported: bool
-    stuck_zone_supported: bool
-    custom_door_sill_supported: bool
-    clean_route_fast_mode_supported: bool
-    cliff_zone_supported: bool
-    smart_door_sill_supported: bool
-    support_floor_direction: bool
-    wifi_manage_supported: bool
-    back_charge_auto_wash_supported: bool
-    support_incremental_map: bool
-    offline_map_supported: bool
-
-
-def build_device_features(feature_set: str, new_feature_set: str) -> DeviceFeatures:
-    new_feature_set_int = int(new_feature_set)
-    feature_set_int = int(feature_set)
-    new_feature_set_divided = int(new_feature_set_int / (2**32))
-    # Convert last 8 digits of new feature set into hexadecimal number
-    converted_new_feature_set = int("0x" + new_feature_set[-8:], 16)
-    new_feature_set_mod_8: bool = len(new_feature_set) % 8 == 0
-    return DeviceFeatures(
-        map_carpet_add_supported=bool(1073741824 & new_feature_set_int),
-        show_clean_finish_reason_supported=bool(1 & new_feature_set_int),
-        resegment_supported=bool(4 & new_feature_set_int),
-        video_monitor_supported=bool(8 & new_feature_set_int),
-        any_state_transit_goto_supported=bool(16 & new_feature_set_int),
-        fw_filter_obstacle_supported=bool(32 & new_feature_set_int),
-        video_setting_supported=bool(64 & new_feature_set_int),
-        ignore_unknown_map_object_supported=bool(128 & new_feature_set_int),
-        set_child_supported=bool(256 & new_feature_set_int),
-        carpet_supported=bool(512 & new_feature_set_int),
-        mop_path_supported=bool(2048 & new_feature_set_int),
-        multi_map_segment_timer_supported=bool(feature_set_int and 4096 & new_feature_set_int),
-        custom_water_box_distance_supported=bool(new_feature_set_int and 2147483648 & new_feature_set_int),
-        wash_then_charge_cmd_supported=bool((new_feature_set_divided >> 5) & 1),
-        room_name_supported=bool(16384 & new_feature_set_int),
-        current_map_restore_enabled=bool(8192 & new_feature_set_int),
-        photo_upload_supported=bool(65536 & new_feature_set_int),
-        shake_mop_set_supported=bool(262144 & new_feature_set_int),
-        map_beautify_internal_debug_supported=bool(2097152 & new_feature_set_int),
-        new_data_for_clean_history=bool(4194304 & new_feature_set_int),
-        new_data_for_clean_history_detail=bool(8388608 & new_feature_set_int),
-        flow_led_setting_supported=bool(16777216 & new_feature_set_int),
-        dust_collection_setting_supported=bool(33554432 & new_feature_set_int),
-        rpc_retry_supported=bool(67108864 & new_feature_set_int),
-        avoid_collision_supported=bool(134217728 & new_feature_set_int),
-        support_set_switch_map_mode=bool(268435456 & new_feature_set_int),
-        support_smart_scene=bool(new_feature_set_divided & 2),
-        support_floor_edit=bool(new_feature_set_divided & 8),
-        support_furniture=bool((new_feature_set_divided >> 4) & 1),
-        support_room_tag=bool((new_feature_set_divided >> 6) & 1),
-        support_quick_map_builder=bool((new_feature_set_divided >> 7) & 1),
-        support_smart_global_clean_with_custom_mode=bool((new_feature_set_divided >> 8) & 1),
-        record_allowed=bool(1024 & new_feature_set_int),
-        careful_slow_map_supported=bool((new_feature_set_divided >> 9) & 1),
-        egg_mode_supported=bool((new_feature_set_divided >> 10) & 1),
-        unsave_map_reason_supported=bool((new_feature_set_divided >> 14) & 1),
-        carpet_show_on_map=bool((new_feature_set_divided >> 12) & 1),
-        supported_valley_electricity=bool((new_feature_set_divided >> 13) & 1),
-        # This one could actually be incorrect
-        # ((t.robotNewFeatures / 2 ** 32) >> 15) & 1 && (module422.DMM.isTopazSV_CE || 'cn' == t.deviceLocation));
-        drying_supported=bool((new_feature_set_divided >> 15) & 1),
-        download_test_voice_supported=bool((new_feature_set_divided >> 16) & 1),
-        support_backup_map=bool((new_feature_set_divided >> 17) & 1),
-        support_custom_mode_in_cleaning=bool((new_feature_set_divided >> 18) & 1),
-        support_remote_control_in_call=bool((new_feature_set_divided >> 19) & 1),
-        support_set_volume_in_call=new_feature_set_mod_8 and bool(1 & converted_new_feature_set),
-        support_clean_estimate=new_feature_set_mod_8 and bool(2 & converted_new_feature_set),
-        support_custom_dnd=new_feature_set_mod_8 and bool(4 & converted_new_feature_set),
-        carpet_deep_clean_supported=bool(8 & converted_new_feature_set),
-        stuck_zone_supported=new_feature_set_mod_8 and bool(16 & converted_new_feature_set),
-        custom_door_sill_supported=new_feature_set_mod_8 and bool(32 & converted_new_feature_set),
-        clean_route_fast_mode_supported=bool(256 & converted_new_feature_set),
-        cliff_zone_supported=new_feature_set_mod_8 and bool(512 & converted_new_feature_set),
-        smart_door_sill_supported=new_feature_set_mod_8 and bool(1024 & converted_new_feature_set),
-        support_floor_direction=new_feature_set_mod_8 and bool(2048 & converted_new_feature_set),
-        wifi_manage_supported=bool(128 & converted_new_feature_set),
-        back_charge_auto_wash_supported=bool(4096 & converted_new_feature_set),
-        support_incremental_map=bool(8192 & converted_new_feature_set),
-        offline_map_supported=bool(16384 & converted_new_feature_set),
-    )
 
 
 @dataclass
@@ -890,6 +755,11 @@ class DeviceData(RoborockBase):
     device: HomeDataDevice
     model: str
     host: str | None = None
+    product_nickname: RoborockProductNickname | None = None
+    device_features: DeviceFeatures | None = None
+
+    def __post_init__(self):
+        self.product_nickname = SHORT_MODEL_TO_ENUM.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS)
 
 
 @dataclass