add getnewargs_ex to Enum, fix pickling (#663)

Author: benedikt-bartscher

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"
 

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