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