diff --git a/esphome/cpp_generator.py b/esphome/cpp_generator.py index d0fc126957f..b17db965212 100644 --- a/esphome/cpp_generator.py +++ b/esphome/cpp_generator.py @@ -1270,6 +1270,9 @@ class MockObj(Expression): return MockObj(op) +_ENUM_NON_ALNUM = re.compile(r"[^A-Za-z0-9]") + + class MockObjEnum(MockObj): def __init__(self, *args, **kwargs): self._enum = kwargs.pop("enum") @@ -1289,6 +1292,17 @@ class MockObjEnum(MockObj): def __repr__(self): return f"MockObj<{str(self.base)}>" + def options(self, *names: str, prefix: str = "") -> dict[str, "MockObj"]: + """Map option names to prefix + NAME enumerators (non-alphanumerics become "_"), for cv.enum.""" + enumerators = { + name: prefix + _ENUM_NON_ALNUM.sub("_", name).upper() for name in names + } + if len(set(enumerators.values())) != len(names): + raise ValueError(f"Option names must map to distinct enumerators: {names}") + return { + name: getattr(self, enumerator) for name, enumerator in enumerators.items() + } + class MockObjClass(MockObj): def __init__(self, *args, **kwargs): diff --git a/tests/unit_tests/test_cpp_generator.py b/tests/unit_tests/test_cpp_generator.py index ec88c583baa..e78035f1f6c 100644 --- a/tests/unit_tests/test_cpp_generator.py +++ b/tests/unit_tests/test_cpp_generator.py @@ -965,3 +965,39 @@ def test_extern_progmem_array_is_an_externally_linked_global() -> None: assert [str(st) for st in CORE.global_statements if "PROGMEM" in str(st)] == [ "const uint8_t esphome::api::API_TABLE[] PROGMEM = {1, 2};" ] + + +class TestEnumOptions: + def test_class_enum_with_prefix(self) -> None: + mode = cg.MockObj("esphome::foo", "::").enum("Mode", is_class=True) + options = mode.options("fast", "slow-ish", "2.5x", prefix="MODE_") + assert list(options) == ["fast", "slow-ish", "2.5x"] + assert {k: str(v) for k, v in options.items()} == { + "fast": "esphome::foo::Mode::MODE_FAST", + "slow-ish": "esphome::foo::Mode::MODE_SLOW_ISH", + "2.5x": "esphome::foo::Mode::MODE_2_5X", + } + + def test_matches_attribute_access(self) -> None: + mode = cg.MockObj("esphome::foo", "::").enum("Mode") + options = mode.options("LOW", "HIGH") + assert str(options["LOW"]) == str(mode.LOW) == "esphome::foo::LOW" + assert str(options["HIGH"]) == str(mode.HIGH) + + def test_result_works_with_cv_enum(self) -> None: + mode = cg.MockObj("esphome::foo", "::").enum("Mode", is_class=True) + options = mode.options("on", "off", prefix="MODE_") + validator = cv.enum(options, lower=True) + assert str(validator("ON").enum_value) == "esphome::foo::Mode::MODE_ON" + assert validator(cv.SCHEMA_EXTRACT) is options + + @pytest.mark.parametrize( + "names", + [("OFF", "OFF"), ("slow-ish", "slow_ish"), ("2.5x", "2_5X")], + ) + def test_names_must_map_to_distinct_enumerators( + self, names: tuple[str, ...] + ) -> None: + mode = cg.MockObj("esphome::foo", "::").enum("Mode") + with pytest.raises(ValueError, match="distinct enumerators"): + mode.options(*names)