Pickle enums by value instead of name (restores pre-3.11 behavior) to support dataclasses.asdict

construct_typed/tenum.py

@@ -74,6 +74,13 @@ class EnumBase(enum.IntEnum):
             return pseudo_member
         return None  # will raise the ValueError in Enum.__new__
 
+    def __reduce_ex__(self, proto: t.Any):
+        """
+        Pickle enums by value instead of name (restores pre-3.11 behavior).
+        See https://github.com/python/cpython/pull/26658 for why this exists.
+        """
+        return self.__class__, (self._value_,)
+
 
 EnumType = t.TypeVar("EnumType", bound=EnumBase)
 
@@ -171,6 +178,13 @@ class FlagsEnumBase(enum.IntFlag):
         new_member.__doc__ = "missing value"
         return new_member
 
+    def __reduce_ex__(self, proto: t.Any):
+        """
+        Pickle enums by value instead of name (restores pre-3.11 behavior).
+        See https://github.com/python/cpython/pull/26658 for why this exists.
+        """
+        return self.__class__, (self._value_,)
+
 
 FlagsEnumType = t.TypeVar("FlagsEnumType", bound=FlagsEnumBase)
 

tests/test_typed.py

@@ -384,6 +384,32 @@ def test_tenum_no_enumbase() -> None:
     assert raises(lambda: cst.TEnum(cs.Byte, cls)) == TypeError
 
 
+def test_tenum_asdict():
+    # see: https://github.com/timrid/construct-typing/issues/21
+    import construct_typed as cst
+    import dataclasses
+
+    class TestEnum(cst.EnumBase):
+        one = 1
+        two = 2
+        four = 4
+        eight = 8
+
+    @dataclasses.dataclass
+    class SomeDataclass:
+        a: TestEnum
+
+    dc = SomeDataclass(TestEnum.one)
+    dc_dict = dataclasses.asdict(dc)
+    assert dc_dict["a"] == dc.a
+    assert dc_dict["a"] is dc.a
+
+    dc = SomeDataclass(TestEnum(5))
+    dc_dict = dataclasses.asdict(dc)
+    assert dc_dict["a"] == dc.a
+    assert dc_dict["a"] is dc.a
+
+
 def test_tenum_docstring() -> None:
     class TestEnum(cst.EnumBase):
         """
@@ -472,6 +498,31 @@ def test_tenum_flags() -> None:
     assert raises(d.build, 2) == TypeError
 
 
+def test_tenum_flags_asdict():
+    import construct_typed as cst
+    import dataclasses
+
+    class TestEnum(cst.FlagsEnumBase):
+        one = 1
+        two = 2
+        four = 4
+        eight = 8
+
+    @dataclasses.dataclass
+    class SomeDataclass:
+        a: TestEnum
+
+    dc = SomeDataclass(TestEnum.one)
+    dc_dict = dataclasses.asdict(dc)
+    assert dc_dict["a"] == dc.a
+    assert dc_dict["a"] is dc.a
+
+    dc = SomeDataclass(TestEnum(5))
+    dc_dict = dataclasses.asdict(dc)
+    assert dc_dict["a"] == dc.a
+    assert dc_dict["a"] is dc.a
+
+
 def test_tenum_flags_docstring() -> None:
     class TestEnum(cst.FlagsEnumBase):
         """