roborock.data.code_mappings

  1from __future__ import annotations
  2
  3import logging
  4from collections import namedtuple
  5from enum import Enum, IntEnum, StrEnum
  6from typing import Self
  7
  8_LOGGER = logging.getLogger(__name__)
  9completed_warnings = set()
 10
 11
 12class RoborockEnum(IntEnum):
 13    """Roborock Enum for codes with int values"""
 14
 15    @property
 16    def name(self) -> str:
 17        return super().name.lower()
 18
 19    @classmethod
 20    def _missing_(cls: type[RoborockEnum], key) -> RoborockEnum:
 21        if hasattr(cls, "unknown"):
 22            warning = f"Missing {cls.__name__} code: {key} - defaulting to 'unknown'"
 23            if warning not in completed_warnings:
 24                completed_warnings.add(warning)
 25                _LOGGER.warning(warning)
 26            return cls.unknown  # type: ignore
 27        default_value = next(item for item in cls)
 28        warning = f"Missing {cls.__name__} code: {key} - defaulting to {default_value}"
 29        if warning not in completed_warnings:
 30            completed_warnings.add(warning)
 31            _LOGGER.warning(warning)
 32        return default_value
 33
 34    @classmethod
 35    def as_dict(cls: type[RoborockEnum]):
 36        return {i.name: i.value for i in cls if i.name != "missing"}
 37
 38    @classmethod
 39    def as_enum_dict(cls: type[RoborockEnum]):
 40        return {i.value: i for i in cls if i.name != "missing"}
 41
 42    @classmethod
 43    def values(cls: type[RoborockEnum]) -> list[int]:
 44        return list(cls.as_dict().values())
 45
 46    @classmethod
 47    def keys(cls: type[RoborockEnum]) -> list[str]:
 48        return list(cls.as_dict().keys())
 49
 50    @classmethod
 51    def items(cls: type[RoborockEnum]):
 52        return cls.as_dict().items()
 53
 54
 55class RoborockModeEnum(StrEnum):
 56    """A custom StrEnum that also stores an integer code for each member."""
 57
 58    code: int
 59    """The integer code associated with the enum member."""
 60
 61    def __new__(cls, value: str, code: int) -> Self:
 62        """Creates a new enum member."""
 63        member = str.__new__(cls, value)
 64        member._value_ = value
 65        member.code = code
 66        return member
 67
 68    @classmethod
 69    def from_code(cls, code: int) -> Self:
 70        for member in cls:
 71            if member.code == code:
 72                return member
 73        message = f"{code} is not a valid code for {cls.__name__}"
 74        if message not in completed_warnings:
 75            completed_warnings.add(message)
 76            _LOGGER.warning(message)
 77        raise ValueError(message)
 78
 79    @classmethod
 80    def from_code_optional(cls, code: int) -> Self | None:
 81        """Gracefully return None if the code does not exist."""
 82        try:
 83            return cls.from_code(code)
 84        except ValueError:
 85            return None
 86
 87    @classmethod
 88    def from_value(cls, value: str) -> Self:
 89        """Find enum member by string value (case-insensitive)."""
 90        for member in cls:
 91            if member.value.lower() == value.lower():
 92                return member
 93        raise ValueError(f"{value} is not a valid value for {cls.__name__}")
 94
 95    @classmethod
 96    def from_name(cls, name: str) -> Self:
 97        """Find enum member by name (case-insensitive)."""
 98        for member in cls:
 99            if member.name.lower() == name.lower():
100                return member
101        raise ValueError(f"{name} is not a valid name for {cls.__name__}")
102
103    @classmethod
104    def keys(cls) -> list[str]:
105        """Returns a list of all member values."""
106        return [member.value for member in cls]
107
108
109ProductInfo = namedtuple("ProductInfo", ["nickname", "short_models"])
110
111
112class RoborockProductNickname(Enum):
113    # Coral Series
114    CORAL = ProductInfo(nickname="Coral", short_models=("a20", "a21"))
115    CORALPRO = ProductInfo(nickname="CoralPro", short_models=("a143", "a144"))
116
117    # Pearl Series
118    PEARL = ProductInfo(nickname="Pearl", short_models=("a74", "a75"))
119    PEARLC = ProductInfo(nickname="PearlC", short_models=("a103", "a104"))
120    PEARLE = ProductInfo(nickname="PearlE", short_models=("a167", "a168"))
121    PEARLELITE = ProductInfo(nickname="PearlELite", short_models=("a169", "a170"))
122    PEARLPLUS = ProductInfo(nickname="PearlPlus", short_models=("a86", "a87"))
123    PEARLPLUSS = ProductInfo(nickname="PearlPlusS", short_models=("a116", "a117", "a136"))
124    PEARLS = ProductInfo(nickname="PearlS", short_models=("a100", "a101"))
125    PEARLSLITE = ProductInfo(nickname="PearlSLite", short_models=("a122", "a123"))
126
127    # Ruby Series
128    RUBYPLUS = ProductInfo(nickname="RubyPlus", short_models=("t4", "s4"))
129    RUBYSC = ProductInfo(nickname="RubySC", short_models=("p5", "a08"))
130    RUBYSE = ProductInfo(nickname="RubySE", short_models=("a19",))
131    RUBYSLITE = ProductInfo(nickname="RubySLite", short_models=("p6", "s5e", "a05"))
132
133    # Tanos Series
134    TANOS = ProductInfo(nickname="Tanos", short_models=("t6", "s6"))
135    TANOSE = ProductInfo(nickname="TanosE", short_models=("t7", "a11"))
136    TANOSS = ProductInfo(nickname="TanosS", short_models=("a14", "a15"))
137    TANOSSC = ProductInfo(nickname="TanosSC", short_models=("a39", "a40"))
138    TANOSSE = ProductInfo(nickname="TanosSE", short_models=("a33", "a34"))
139    TANOSSMAX = ProductInfo(nickname="TanosSMax", short_models=("a52",))
140    TANOSSLITE = ProductInfo(nickname="TanosSLite", short_models=("a37", "a38"))
141    TANOSSPLUS = ProductInfo(nickname="TanosSPlus", short_models=("a23", "a24"))
142    TANOSV = ProductInfo(nickname="TanosV", short_models=("t7p", "a09", "a10"))
143
144    # Topaz Series
145    TOPAZS = ProductInfo(nickname="TopazS", short_models=("a29", "a30", "a76"))
146    TOPAZSC = ProductInfo(nickname="TopazSC", short_models=("a64", "a65"))
147    TOPAZSPLUS = ProductInfo(nickname="TopazSPlus", short_models=("a46", "a47", "a66"))
148    TOPAZSPOWER = ProductInfo(nickname="TopazSPower", short_models=("a62",))
149    TOPAZSV = ProductInfo(nickname="TopazSV", short_models=("a26", "a27"))
150
151    # Ultron Series
152    ULTRON = ProductInfo(nickname="Ultron", short_models=("a50", "a51"))
153    ULTRONE = ProductInfo(nickname="UltronE", short_models=("a72", "a84"))
154    ULTRONLITE = ProductInfo(nickname="UltronLite", short_models=("a73", "a85"))
155    ULTRONSC = ProductInfo(nickname="UltronSC", short_models=("a94", "a95"))
156    ULTRONSE = ProductInfo(nickname="UltronSE", short_models=("a124", "a125", "a139", "a140"))
157    ULTRONSPLUS = ProductInfo(nickname="UltronSPlus", short_models=("a68", "a69", "a70"))
158    ULTRONSV = ProductInfo(nickname="UltronSV", short_models=("a96", "a97"))
159
160    # Verdelite Series
161    VERDELITE = ProductInfo(nickname="Verdelite", short_models=("a146", "a147"))
162
163    # Vivian Series
164    VIVIAN = ProductInfo(nickname="Vivian", short_models=("a134", "a135", "a155", "a156"))
165    VIVIANC = ProductInfo(nickname="VivianC", short_models=("a158", "a159"))
166
167
168SHORT_MODEL_TO_ENUM = {model: product for product in RoborockProductNickname for model in product.value.short_models}
169
170
171class RoborockCategory(Enum):
172    """Describes the category of the device."""
173
174    WET_DRY_VAC = "roborock.wetdryvac"
175    VACUUM = "robot.vacuum.cleaner"
176    WASHING_MACHINE = "roborock.wm"
177    MOWER = "roborock.mower"
178    UNKNOWN = "UNKNOWN"
179
180    @classmethod
181    def _missing_(cls, value):
182        _LOGGER.warning("Missing code %s from category", value)
183        return RoborockCategory.UNKNOWN