From 08e8a688938cbac15f54d2cc8a450cc30772fc6b Mon Sep 17 00:00:00 2001
From: benedikt-bartscher
 <31854409+benedikt-bartscher@users.noreply.github.com>
Date: Sun, 30 Mar 2025 23:48:58 +0200
Subject: [PATCH] add getnewargs_ex to Enum, fix pickling (#663)

---
 src/betterproto/enum.py |  3 +++
 tests/test_pickling.py  | 13 +++++++++++++
 2 files changed, 16 insertions(+)

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