diff --git a/src/betterproto/enum.py b/src/betterproto/enum.py index f5f14dcc..6b1b7e0a 100644 --- a/src/betterproto/enum.py +++ b/src/betterproto/enum.py @@ -132,6 +132,9 @@ def __new__(cls, *, name: Optional[str], value: int) -> Self: 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 5650f750..32478c84 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 use_cache(): .string_value == "world" ) + + +def test_pickle_enum(): + enum = BetterprotoEnum.ONE + assert unpickled(enum) == enum + + enum = BetterprotoEnum.UNSPECIFIED + assert unpickled(enum) == enum