feat(enums): Add AliasedStrEnum helper class and tests

Author: Ishirui

src/dda/utils/enums.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: 2025-present Datadog, Inc. <dev@datadoghq.com>
+#
+# SPDX-License-Identifier: MIT
+
+from collections.abc import Iterator
+from enum import Enum, EnumMeta, StrEnum
+from typing import Self
+
+
+# Since __iter__ is defined at metaclass level for Enums, we cannot override it in AliasedStrEnum
+# So we need to create a new metaclass that overrides __iter__.
+class AliasedStrEnumMeta(EnumMeta):
+    """Metaclass for AliasedStrEnum that handles iteration correctly."""
+
+    def __iter__(cls) -> Iterator[Enum]:  # type: ignore[override]
+        """Override default StrEnum iteration - needed because of __new__ override."""
+        return iter(x for x in cls._member_map_.values())
+
+
+class AliasedStrEnum(StrEnum, metaclass=AliasedStrEnumMeta):
+    """
+    A string enumeration that supports multiple alias values for each member.
+    The first value is considered the canonical value.
+
+    When instantiating or comparing, any of the alias values will be considered equal to the canonical value.
+    Example:
+        class OS(AliasedStrEnum):
+            LINUX = "linux"
+            WINDOWS = ("windows", "nt", "win")
+
+        OS.WINDOWS == "windows"  # True
+        OS.WINDOWS == "nt"       # True
+        OS.WINDOWS == "linux"    # False
+
+        OS("win") is OS.WINDOWS  # True
+        OS("linux") is OS.LINUX  # True
+        OS("bsd")                # Raises ValueError
+
+        OS["WINDOWS"]            # OS.WINDOWS
+        OS["win"]            # Raises KeyError - use OS("win") instead
+        OS.WINDOWS.name          # "WINDOWS"
+
+        str(OS.WINDOWS)         # "windows"
+        OS.WINDOWS.value        # "windows"
+        OS.WINDOWS.values       # {"windows", "nt", "win"}
+        OS.WINDOWS.aliases      # {"nt", "win"}
+    """
+
+    # This method is called when registering new enum members, not when instantiating them.
+    def __new__(cls, value: str, *alt_values: str) -> Self:
+        # Create the StrEnum member instance
+        obj = str.__new__(cls, value)
+        obj._value_ = value
+        obj._alt_values_ = set(alt_values)  # type: ignore[attr-defined]
+
+        # Register in the reverse mapping for by-value lookup (e.g. Enum(value))
+        cls._value2member_map_[value] = obj
+        for alias in alt_values:
+            cls._value2member_map_[alias] = obj
+        return obj
+
+    @property
+    def aliases(self) -> set[str]:
+        """Returns the set of non-canonical alias values for this enum member."""
+        return self._alt_values_  # type: ignore[attr-defined]
+
+    @property
+    def values(self) -> set[str]:
+        """Returns the set of all values (canonical + aliases) for this enum member."""
+        return {self.value} | self._alt_values_  # type: ignore[attr-defined]
+
+    def __str__(self) -> str:
+        return self.value
+
+    def __repr__(self) -> str:
+        aliases_str = f" {tuple(sorted(self._alt_values_))!r}" if self._alt_values_ else ""  # type: ignore[attr-defined]
+        return f"<{self.__class__.__name__}.{self.name}: {self.value!r}{aliases_str}>"
+
+    def __eq__(self, value: object) -> bool:
+        """Override default StrEnum equality to consider equality with any of the alias values."""
+        if isinstance(value, str):
+            return value in self.values
+        return super().__eq__(value)
+
+    def __contains__(self, item: object) -> bool:
+        """Override default StrEnum containment to consider containment with any of the alias values."""
+        if isinstance(item, str):
+            return item in self.values
+        return False
+
+    def __hash__(self) -> int:
+        """Override default StrEnum hash to consider hash of any of the alias values."""
+        return hash(self.value)

tests/utils/test_enums.py

@@ -0,0 +1,144 @@
+# SPDX-FileCopyrightText: 2024-present Datadog, Inc. <dev@datadoghq.com>
+#
+# SPDX-License-Identifier: MIT
+from __future__ import annotations
+
+import pytest
+
+from dda.utils.enums import AliasedStrEnum
+
+
+class ExampleEnum(AliasedStrEnum):
+    """Test enum for testing AliasedStrEnum functionality."""
+
+    LINUX = "linux"
+    WINDOWS = ("windows", "nt", "win")
+    MACOS = ("macos", "darwin", "osx")
+
+
+class TestAliasedStrEnum:
+    def test_canonical_value_access(self):
+        """Test that canonical values work correctly."""
+        assert ExampleEnum.LINUX == "linux"
+        assert ExampleEnum.WINDOWS == "windows"
+        assert ExampleEnum.MACOS == "macos"
+
+    def test_alias_equality(self):
+        """Test that alias values are equal to the enum member."""
+        assert ExampleEnum.WINDOWS == "windows"
+        assert ExampleEnum.WINDOWS == "nt"
+        assert ExampleEnum.WINDOWS == "win"
+
+        assert ExampleEnum.MACOS == "macos"
+        assert ExampleEnum.MACOS == "darwin"
+        assert ExampleEnum.MACOS == "osx"
+
+    def test_alias_inequality(self):
+        """Test that non-alias values are not equal."""
+        assert ExampleEnum.WINDOWS != "linux"
+        assert ExampleEnum.LINUX != "windows"
+        assert ExampleEnum.MACOS != "linux"
+
+    def test_instantiation_by_canonical_value(self):
+        """Test instantiation using canonical values."""
+        assert ExampleEnum("linux") is ExampleEnum.LINUX
+        assert ExampleEnum("windows") is ExampleEnum.WINDOWS
+        assert ExampleEnum("macos") is ExampleEnum.MACOS
+
+    def test_instantiation_by_alias(self):
+        """Test instantiation using alias values."""
+        assert ExampleEnum("nt") is ExampleEnum.WINDOWS
+        assert ExampleEnum("win") is ExampleEnum.WINDOWS
+        assert ExampleEnum("darwin") is ExampleEnum.MACOS
+        assert ExampleEnum("osx") is ExampleEnum.MACOS
+
+    def test_instantiation_invalid_value(self):
+        """Test that invalid values raise ValueError."""
+        with pytest.raises(match="'bsd' is not a valid TestOS"):
+            ExampleEnum("bsd")
+        with pytest.raises(match="'invalid' is not a valid TestOS"):
+            ExampleEnum("invalid")
+
+    def test_name_access(self):
+        """Test that name property returns the enum member name."""
+        assert ExampleEnum.LINUX.name == "LINUX"
+        assert ExampleEnum.WINDOWS.name == "WINDOWS"
+        assert ExampleEnum.MACOS.name == "MACOS"
+
+    def test_value_property(self):
+        """Test that value property returns the canonical value."""
+        assert ExampleEnum.LINUX.value == "linux"
+        assert ExampleEnum.WINDOWS.value == "windows"
+        assert ExampleEnum.MACOS.value == "macos"
+
+    def test_aliases_property(self):
+        """Test that aliases property returns the set of non-canonical values."""
+        assert ExampleEnum.LINUX.aliases == set()
+        assert ExampleEnum.WINDOWS.aliases == {"nt", "win"}
+        assert ExampleEnum.MACOS.aliases == {"darwin", "osx"}
+
+    def test_values_property(self):
+        """Test that values property returns all values (canonical + aliases)."""
+        assert ExampleEnum.LINUX.values == {"linux"}
+        assert ExampleEnum.WINDOWS.values == {"windows", "nt", "win"}
+        assert ExampleEnum.MACOS.values == {"macos", "darwin", "osx"}
+
+    def test_str_representation(self):
+        """Test string representation returns canonical value."""
+        assert str(ExampleEnum.LINUX) == "linux"
+        assert str(ExampleEnum.WINDOWS) == "windows"
+        assert str(ExampleEnum.MACOS) == "macos"
+
+    def test_repr_representation(self):
+        """Test repr representation includes aliases."""
+        assert repr(ExampleEnum.LINUX) == "<TestOS.LINUX: 'linux'>"
+
+        # For members with aliases, check that aliases are included
+        assert repr(ExampleEnum.WINDOWS) == "<TestOS.WINDOWS: 'windows' ('nt', 'win')>"
+        assert repr(ExampleEnum.MACOS) == "<TestOS.MACOS: 'macos' ('darwin', 'osx')>"
+
+    def test_containment(self):
+        """Test containment operator with alias values."""
+        assert "linux" in ExampleEnum.LINUX
+        assert "windows" in ExampleEnum.WINDOWS
+        assert "nt" in ExampleEnum.WINDOWS
+        assert "win" in ExampleEnum.WINDOWS
+        assert "darwin" in ExampleEnum.MACOS
+        assert "osx" in ExampleEnum.MACOS
+
+        assert "bsd" not in ExampleEnum.LINUX
+        assert "linux" not in ExampleEnum.WINDOWS
+
+    def test_hash_consistency(self):
+        """Test that hash is consistent for the same enum member."""
+        assert hash(ExampleEnum.LINUX) == hash(ExampleEnum.LINUX)
+        assert hash(ExampleEnum.WINDOWS) == hash(ExampleEnum.WINDOWS)
+
+        # Hash should be based on canonical value
+        assert hash(ExampleEnum.WINDOWS) == hash("windows")
+
+    def test_member_access_by_name(self):
+        """Test accessing members by name using bracket notation."""
+        assert ExampleEnum["LINUX"] is ExampleEnum.LINUX
+        assert ExampleEnum["WINDOWS"] is ExampleEnum.WINDOWS
+        assert ExampleEnum["MACOS"] is ExampleEnum.MACOS
+
+    def test_member_access_by_alias_fails(self):
+        """Test that accessing by alias value using bracket notation fails."""
+        with pytest.raises(KeyError):
+            ExampleEnum["nt"]
+        with pytest.raises(KeyError):
+            ExampleEnum["darwin"]
+
+    def test_iteration(self):
+        """Test that enum can be iterated over."""
+        members = list(ExampleEnum)
+        assert len(members) == 3
+        assert ExampleEnum.LINUX in members
+        assert ExampleEnum.WINDOWS in members
+        assert ExampleEnum.MACOS in members
+
+        members2 = []
+        for member in ExampleEnum:
+            members2.append(member)  # noqa: PERF402
+        assert members == members2