diff --git a/src/betterproto/enum.py b/src/betterproto/enum.py index f5f14dc..6b1b7e0 100644 --- a/src/betterproto/enum.py +++ b/src/betterproto/enum.py @@ -132,6 +132,9 @@ class Enum(IntEnum if TYPE_CHECKING else int, metaclass=EnumType): super().__setattr__(self, "value", value) return self + def __getnewargs_ex__(self) -> Tuple[Tuple[()], Dict[str, Any]]: + return (), {"name": self.name, "value": self.value} + def __str__(self) -> str: return self.name or "None" diff --git a/tests/test_pickling.py b/tests/test_pickling.py index 5650f75..32478c8 100644 --- a/tests/test_pickling.py +++ b/tests/test_pickling.py @@ -57,6 +57,11 @@ class Complex(betterproto.Message): ) +class BetterprotoEnum(betterproto.Enum): + UNSPECIFIED = 0 + ONE = 1 + + def complex_msg(): return Complex( foo_str="yep", @@ -201,3 +206,11 @@ def test_message_can_be_cached(): .string_value == "world" ) + + +def test_pickle_enum(): + enum = BetterprotoEnum.ONE + assert unpickled(enum) == enum + + enum = BetterprotoEnum.UNSPECIFIED + assert unpickled(enum) == enum