roborock.data.code_mappings

  1import logging
  2from collections import namedtuple
  3from enum import Enum, IntEnum, StrEnum
  4from typing import Any, Self
  5
  6_LOGGER = logging.getLogger(__name__)
  7completed_warnings = set()
  8
  9
 10class RoborockEnum(IntEnum):
 11    """Roborock Enum for codes with int values"""
 12
 13    _display_name_: str | None
 14
 15    def __new__(cls, value: int, display_name: str | None = None) -> Self:
 16        member = int.__new__(cls, value)
 17        member._value_ = value
 18        member._display_name_ = display_name
 19        return member
 20
 21    @property
 22    def name(self) -> str:
 23        return super().name.lower()
 24
 25    @property
 26    def display_name(self) -> str:
 27        return self._display_name_ or self.name
 28
 29    @classmethod
 30    def _missing_(cls: type[Self], key) -> Self:
 31        if hasattr(cls, "unknown"):
 32            warning = f"Missing {cls.__name__} code: {key} - defaulting to 'unknown'"
 33            if warning not in completed_warnings:
 34                completed_warnings.add(warning)
 35                _LOGGER.warning(warning)
 36            return cls.unknown  # type: ignore
 37        default_value = next(item for item in cls)
 38        warning = f"Missing {cls.__name__} code: {key} - defaulting to {default_value}"
 39        if warning not in completed_warnings:
 40            completed_warnings.add(warning)
 41            _LOGGER.warning(warning)
 42        return default_value
 43
 44    @classmethod
 45    def as_dict(cls: type[Self]) -> dict[str, int]:
 46        result: dict[str, int] = {}
 47        for item in cls:
 48            if item.name == "missing":
 49                continue
 50            result.setdefault(item.display_name, item.value)
 51        return result
 52
 53    @classmethod
 54    def as_enum_dict(cls: type[Self]):
 55        return {i.value: i for i in cls if i.name != "missing"}
 56
 57    @classmethod
 58    def values(cls: type[Self]) -> list[int]:
 59        return list(cls.as_dict().values())
 60
 61    @classmethod
 62    def keys(cls: type[Self]) -> list[str]:
 63        return list(cls.as_dict().keys())
 64
 65    @classmethod
 66    def items(cls: type[Self]):
 67        return cls.as_dict().items()
 68
 69
 70class RoborockModeEnum(StrEnum):
 71    """A custom StrEnum that also stores an integer code for each member."""
 72
 73    code: int
 74    """The integer code associated with the enum member."""
 75
 76    def __new__(cls, value: str, code: int) -> Self:
 77        """Creates a new enum member."""
 78        member = str.__new__(cls, value)
 79        member._value_ = value
 80        member.code = code
 81        return member
 82
 83    @classmethod
 84    def from_code(cls, code: int) -> Self:
 85        for member in cls:
 86            if member.code == code:
 87                return member
 88        message = f"{code} is not a valid code for {cls.__name__}"
 89        if message not in completed_warnings:
 90            completed_warnings.add(message)
 91            _LOGGER.warning(message)
 92        raise ValueError(message)
 93
 94    @classmethod
 95    def from_code_optional(cls, code: int) -> Self | None:
 96        """Gracefully return None if the code does not exist.
 97
 98        This is the silent counterpart to :meth:`from_code`: callers use it when
 99        an unknown code is expected and tolerable (e.g. decoding a device push
100        that may include data points this library does not model yet), so it must
101        not emit the "not a valid code" warning that ``from_code`` logs.
102        """
103        for member in cls:
104            if member.code == code:
105                return member
106        return None
107
108    @classmethod
109    def from_value(cls, value: str) -> Self:
110        """Find enum member by string value (case-insensitive)."""
111        for member in cls:
112            if member.value.lower() == value.lower():
113                return member
114        raise ValueError(f"{value} is not a valid value for {cls.__name__}")
115
116    @classmethod
117    def from_name(cls, name: str) -> Self:
118        """Find enum member by name (case-insensitive)."""
119        for member in cls:
120            if member.name.lower() == name.lower():
121                return member
122        raise ValueError(f"{name} is not a valid name for {cls.__name__}")
123
124    @classmethod
125    def from_any_optional(cls, value: str | int) -> Self | None:
126        """Resolve a string or int to an enum member.
127
128        Tries to look up by enum name, string value, or integer code
129        and returns None if no match is found.
130        """
131        # Try enum name lookup (e.g. "SEEK")
132        try:
133            return cls.from_name(str(value))
134        except ValueError:
135            pass
136        # Try DP string value lookup (e.g. "dpSeek")
137        try:
138            return cls.from_value(str(value))
139        except ValueError:
140            pass
141        # Try integer code lookup (e.g. "11"). Use the silent optional variant so
142        # a value that is neither a name, a DP string, nor a known code resolves
143        # to None without logging a spurious "not a valid code" warning.
144        try:
145            int_code = int(value)
146        except (ValueError, TypeError):
147            return None
148        return cls.from_code_optional(int_code)
149
150    @classmethod
151    def keys(cls) -> list[str]:
152        """Returns a list of all member values."""
153        return [member.value for member in cls]
154
155    def __eq__(self, other: Any) -> bool:
156        if isinstance(other, str):
157            return self.value == other or self.name == other
158        if isinstance(other, int):
159            return self.code == other
160        return super().__eq__(other)
161
162    def __hash__(self) -> int:
163        """Hash a RoborockModeEnum.
164
165        It is critical that you do not mix RoborockModeEnums with raw strings or ints in hashed situations
166        (i.e. sets or keys in dictionaries)
167        """
168        return hash((self.code, self._value_))
169
170
171ProductInfo = namedtuple("ProductInfo", ["nickname", "short_models"])
172
173
174class RoborockProductNickname(Enum):
175    # Coral Series
176    CORAL = ProductInfo(nickname="Coral", short_models=("a20", "a21"))
177    CORALPRO = ProductInfo(nickname="CoralPro", short_models=("a143", "a144"))
178
179    # Pearl Series
180    PEARL = ProductInfo(nickname="Pearl", short_models=("a74", "a75"))
181    PEARLC = ProductInfo(nickname="PearlC", short_models=("a103", "a104"))
182    PEARLE = ProductInfo(nickname="PearlE", short_models=("a167", "a168"))
183    PEARLELITE = ProductInfo(nickname="PearlELite", short_models=("a169", "a170"))
184    PEARLPLUS = ProductInfo(nickname="PearlPlus", short_models=("a86", "a87"))
185    PEARLPLUSS = ProductInfo(nickname="PearlPlusS", short_models=("a116", "a117", "a136"))
186    PEARLS = ProductInfo(nickname="PearlS", short_models=("a100", "a101"))
187    PEARLSLITE = ProductInfo(nickname="PearlSLite", short_models=("a122", "a123"))
188
189    # Ruby Series
190    RUBYPLUS = ProductInfo(nickname="RubyPlus", short_models=("t4", "s4"))
191    RUBYSC = ProductInfo(nickname="RubySC", short_models=("p5", "a08"))
192    RUBYSE = ProductInfo(nickname="RubySE", short_models=("a19",))
193    RUBYSLITE = ProductInfo(nickname="RubySLite", short_models=("p6", "s5e", "a05"))
194
195    # Tanos Series
196    TANOS = ProductInfo(nickname="Tanos", short_models=("t6", "s6"))
197    TANOSE = ProductInfo(nickname="TanosE", short_models=("t7", "a11"))
198    TANOSS = ProductInfo(nickname="TanosS", short_models=("a14", "a15"))
199    TANOSSC = ProductInfo(nickname="TanosSC", short_models=("a39", "a40"))
200    TANOSSE = ProductInfo(nickname="TanosSE", short_models=("a33", "a34"))
201    TANOSSMAX = ProductInfo(nickname="TanosSMax", short_models=("a52",))
202    TANOSSLITE = ProductInfo(nickname="TanosSLite", short_models=("a37", "a38"))
203    TANOSSPLUS = ProductInfo(nickname="TanosSPlus", short_models=("a23", "a24"))
204    TANOSV = ProductInfo(nickname="TanosV", short_models=("t7p", "a09", "a10"))
205
206    # Topaz Series
207    TOPAZS = ProductInfo(nickname="TopazS", short_models=("a29", "a30", "a76"))
208    TOPAZSC = ProductInfo(nickname="TopazSC", short_models=("a64", "a65"))
209    TOPAZSPLUS = ProductInfo(nickname="TopazSPlus", short_models=("a46", "a47", "a66"))
210    TOPAZSPOWER = ProductInfo(nickname="TopazSPower", short_models=("a62",))
211    TOPAZSV = ProductInfo(nickname="TopazSV", short_models=("a26", "a27"))
212
213    # Ultron Series
214    ULTRON = ProductInfo(nickname="Ultron", short_models=("a50", "a51"))
215    ULTRONE = ProductInfo(nickname="UltronE", short_models=("a72", "a84"))
216    ULTRONLITE = ProductInfo(nickname="UltronLite", short_models=("a73", "a85"))
217    ULTRONSC = ProductInfo(nickname="UltronSC", short_models=("a94", "a95"))
218    ULTRONSE = ProductInfo(nickname="UltronSE", short_models=("a124", "a125", "a139", "a140"))
219    ULTRONSPLUS = ProductInfo(nickname="UltronSPlus", short_models=("a68", "a69", "a70"))
220    ULTRONSV = ProductInfo(nickname="UltronSV", short_models=("a96", "a97"))
221
222    # Verdelite Series
223    VERDELITE = ProductInfo(nickname="Verdelite", short_models=("a146", "a147"))
224
225    # Vivian Series
226    VIVIAN = ProductInfo(nickname="Vivian", short_models=("a134", "a135", "a155", "a156"))
227    VIVIANC = ProductInfo(nickname="VivianC", short_models=("a158", "a159"))
228
229
230SHORT_MODEL_TO_ENUM = {model: product for product in RoborockProductNickname for model in product.value.short_models}
231
232
233class RoborockCategory(Enum):
234    """Describes the category of the device."""
235
236    WET_DRY_VAC = "roborock.wetdryvac"
237    VACUUM = "robot.vacuum.cleaner"
238    WASHING_MACHINE = "roborock.wm"
239    MOWER = "roborock.mower"
240    UNKNOWN = "UNKNOWN"
241
242    @classmethod
243    def _missing_(cls, value):
244        _LOGGER.warning("Missing code %s from category", value)
245        return RoborockCategory.UNKNOWN