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