Greatly simplify enums implementation (#72)

* Greatly simplify enums implementation

* Remove try_value

* Fix documentation

* Update version

Author: AdrienVannson

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "betterproto2"
-version = "0.2.3"
+version = "0.3.0"
 description = "A better Protobuf / gRPC generator & library"
 authors = [
     {name = "Adrien Vannson", email = "[email protected]"},

src/betterproto2/__init__.py

@@ -823,7 +823,7 @@ def _postprocess_single(self, wire_type: int, meta: FieldMetadata, field_name: s
                 value = value > 0
             elif meta.proto_type == TYPE_ENUM:
                 # Convert enum ints to python enum instances
-                value = self._betterproto.cls_by_field[field_name].try_value(value)
+                value = self._betterproto.cls_by_field[field_name](value)
         elif wire_type in (WIRE_FIXED_32, WIRE_FIXED_64):
             fmt = _pack_fmt(meta.proto_type)
             value = struct.unpack(fmt, value)[0]

src/betterproto2/enum.py

@@ -1,177 +1,45 @@
-from __future__ import annotations
+from enum import IntEnum
 
-from enum import (
-    EnumMeta,
-    IntEnum,
-)
-from types import MappingProxyType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-)
+from typing_extensions import Self
 
-if TYPE_CHECKING:
-    from collections.abc import (
-        Generator,
-        Mapping,
-    )
-
-    from typing_extensions import (
-        Never,
-        Self,
-    )
-
-
-def _is_descriptor(obj: object) -> bool:
-    return hasattr(obj, "__get__") or hasattr(obj, "__set__") or hasattr(obj, "__delete__")
-
-
-class EnumType(EnumMeta if TYPE_CHECKING else type):
-    _value_map_: Mapping[int, Enum]
-    _member_map_: Mapping[str, Enum]
-
-    def __new__(mcs, name: str, bases: tuple[type, ...], namespace: dict[str, Any]) -> Self:
-        value_map = {}
-        member_map = {}
-
-        new_mcs = type(
-            f"{name}Type",
-            tuple(
-                dict.fromkeys([base.__class__ for base in bases if base.__class__ is not type] + [EnumType, type])
-            ),  # reorder the bases so EnumType and type are last to avoid conflicts
-            {"_value_map_": value_map, "_member_map_": member_map},
-        )
-
-        members = {
-            name: value for name, value in namespace.items() if not _is_descriptor(value) and not name.startswith("__")
-        }
-
-        cls = type.__new__(
-            new_mcs,
-            name,
-            bases,
-            {key: value for key, value in namespace.items() if key not in members},
-        )
-        # this allows us to disallow member access from other members as
-        # members become proper class variables
-
-        for name, value in members.items():
-            member = value_map.get(value)
-            if member is None:
-                member = cls.__new__(cls, name=name, value=value)  # type: ignore
-                value_map[value] = member
-            member_map[name] = member
-            type.__setattr__(new_mcs, name, member)
-
-        return cls
-
-    if not TYPE_CHECKING:
-
-        def __call__(cls, value: int) -> Enum:
-            try:
-                return cls._value_map_[value]
-            except (KeyError, TypeError):
-                raise ValueError(f"{value!r} is not a valid {cls.__name__}") from None
-
-        def __iter__(cls) -> Generator[Enum, None, None]:
-            yield from cls._member_map_.values()
-
-        def __reversed__(cls) -> Generator[Enum, None, None]:
-            yield from reversed(cls._member_map_.values())
-
-        def __getitem__(cls, key: str) -> Enum:
-            return cls._member_map_[key]
-
-        @property
-        def __members__(cls) -> MappingProxyType[str, Enum]:
-            return MappingProxyType(cls._member_map_)
-
-    def __repr__(cls) -> str:
-        return f"<enum {cls.__name__!r}>"
-
-    def __len__(cls) -> int:
-        return len(cls._member_map_)
-
-    def __setattr__(cls, name: str, value: Any) -> Never:
-        raise AttributeError(f"{cls.__name__}: cannot reassign Enum members.")
-
-    def __delattr__(cls, name: str) -> Never:
-        raise AttributeError(f"{cls.__name__}: cannot delete Enum members.")
-
-    def __contains__(cls, member: object) -> bool:
-        return isinstance(member, cls) and member.name in cls._member_map_
-
-
-class Enum(IntEnum if TYPE_CHECKING else int, metaclass=EnumType):
-    """
-    The base class for protobuf enumerations, all generated enumerations will
-    inherit from this. Emulates `enum.IntEnum`.
-    """
-
-    name: str | None
-    value: int
-
-    if not TYPE_CHECKING:
-
-        def __new__(cls, *, name: str | None, value: int) -> Self:
-            self = super().__new__(cls, value)
-            super().__setattr__(self, "name", name)
-            super().__setattr__(self, "value", value)
-            return self
-
-    def __str__(self) -> str:
-        return self.name or "None"
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}.{self.name}"
-
-    def __setattr__(self, key: str, value: Any) -> Never:
-        raise AttributeError(f"{self.__class__.__name__} Cannot reassign a member's attributes.")
-
-    def __delattr__(self, item: Any) -> Never:
-        raise AttributeError(f"{self.__class__.__name__} Cannot delete a member's attributes.")
-
-    def __copy__(self) -> Self:
-        return self
-
-    def __deepcopy__(self, memo: Any) -> Self:
-        return self
 
+class Enum(IntEnum):
     @classmethod
-    def try_value(cls, value: int = 0) -> Self:
-        """Return the value which corresponds to the value.
-
-        Parameters
-        -----------
-        value: :class:`int`
-            The value of the enum member to get.
+    def _missing_(cls, value):
+        # If the given value is not an integer, let the standard enum implementation raise an error
+        if not isinstance(value, int):
+            return None
+
+        # Create a new "unknown" instance with the given value.
+        obj = int.__new__(cls, value)
+        obj._value_ = value
+        obj._name_ = ""
+        return obj
+
+    def __str__(self):
+        if not self.name:
+            return f"{self.__class__.__name__}.~UNKNOWN({self.value})"
+        return f"{self.__class__.__name__}.{self.name}"
 
-        Returns
-        -------
-        :class:`Enum`
-            The corresponding member or a new instance of the enum if
-            ``value`` isn't actually a member.
-        """
-        try:
-            return cls._value_map_[value]
-        except (KeyError, TypeError):
-            return cls.__new__(cls, name=None, value=value)
+    def __repr__(self):
+        if not self.name:
+            return f"<{self.__class__.__name__}.~UNKNOWN: {self.value}>"
+        return super().__repr__()
 
     @classmethod
     def from_string(cls, name: str) -> Self:
         """Return the value which corresponds to the string name.
 
-        Parameters
-        -----------
-        name: :class:`str`
-            The name of the enum member to get.
+        Parameters:
+            name: The name of the enum member to get.
+
+        Raises:
+            ValueError: The member was not found in the Enum.
 
-        Raises
-        -------
-        :exc:`ValueError`
-            The member was not found in the Enum.
+        Returns:
+            The corresponding value
         """
         try:
-            return cls._member_map_[name]
+            return cls[name]
         except KeyError as e:
             raise ValueError(f"Unknown value {name} for enum {cls.__name__}") from e

src/betterproto2/internal_lib/google/protobuf/__init__.py

@@ -9,15 +9,16 @@ class Colour(betterproto2.Enum):
     BLUE = 3
 
 
-PURPLE = Colour.__new__(Colour, name=None, value=4)
+PURPLE = Colour(4)
 
 
 @pytest.mark.parametrize(
     "member, str_value",
     [
-        (Colour.RED, "RED"),
-        (Colour.GREEN, "GREEN"),
-        (Colour.BLUE, "BLUE"),
+        (Colour.RED, "Colour.RED"),
+        (Colour.GREEN, "Colour.GREEN"),
+        (Colour.BLUE, "Colour.BLUE"),
+        (PURPLE, "Colour.~UNKNOWN(4)"),
     ],
 )
 def test_str(member: Colour, str_value: str) -> None:
@@ -27,9 +28,10 @@ def test_str(member: Colour, str_value: str) -> None:
 @pytest.mark.parametrize(
     "member, repr_value",
     [
-        (Colour.RED, "Colour.RED"),
-        (Colour.GREEN, "Colour.GREEN"),
-        (Colour.BLUE, "Colour.BLUE"),
+        (Colour.RED, "<Colour.RED: 1>"),
+        (Colour.GREEN, "<Colour.GREEN: 2>"),
+        (Colour.BLUE, "<Colour.BLUE: 3>"),
+        (PURPLE, "<Colour.~UNKNOWN: 4>"),
     ],
 )
 def test_repr(member: Colour, repr_value: str) -> None:
@@ -42,7 +44,7 @@ def test_repr(member: Colour, repr_value: str) -> None:
         (Colour.RED, ("RED", 1)),
         (Colour.GREEN, ("GREEN", 2)),
         (Colour.BLUE, ("BLUE", 3)),
-        (PURPLE, (None, 4)),
+        (PURPLE, ("", 4)),
     ],
 )
 def test_name_values(member: Colour, values: tuple[str | None, int]) -> None:
@@ -70,5 +72,5 @@ def test_from_string(member: Colour, input_str: str) -> None:
         (PURPLE, 4),
     ],
 )
-def test_try_value(member: Colour, input_int: int) -> None:
-    assert Colour.try_value(input_int) == member
+def test_construction(member: Colour, input_int: int) -> None:
+    assert Colour(input_int) == member