From db17f0c521d7ff16f608962ca328d65507787c82 Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Sat, 29 Mar 2025 02:23:15 +0100 Subject: [PATCH 1/2] add getnewargs_ex to Enum, fix pickling --- src/betterproto/enum.py | 3 +++ 1 file changed, 3 insertions(+) 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" From ce1e042187b0e6155d03e6ae58364f2be344296e Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Sat, 29 Mar 2025 11:41:59 +0100 Subject: [PATCH 2/2] add minimal test for enum pickling --- tests/test_pickling.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) 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