pythongh-105332: Fix unpickling enum.Flag with specific values

sobolevn · Jun 6, 2023 · b947792 · b947792
1 parent f04c168
commit b947792
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/Lib/enum.py b/Lib/enum.py
@@ -1308,6 +1308,8 @@ class Flag(Enum, boundary=STRICT):
 
     def __reduce_ex__(self, proto):
         cls = self.__class__
+        if self._value_ in cls._value2member_map_:
+            return cls, (self._value_,)
         unknown = self._value_ & ~cls._flag_mask_
         member_value = self._value_ & cls._flag_mask_
         if unknown and member_value:
@@ -1318,10 +1320,7 @@ def __reduce_ex__(self, proto):
                 return _or_, (cls(rest), cls._value2member_map_.get(val))
             else:
                 break
-        if self._name_ is None:
-            return cls, (self._value_,)
-        else:
-            return getattr, (cls, self._name_)
+        return cls, (self._value_,)
 
     _numeric_repr_ = repr
 

diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py
@@ -66,6 +66,7 @@ class FlagStooges(Flag):
         LARRY = 1
         CURLY = 2
         MOE = 4
+        BIG = 389
 except Exception as exc:
     FlagStooges = exc
 
@@ -74,17 +75,20 @@ class FlagStoogesWithZero(Flag):
     LARRY = 1
     CURLY = 2
     MOE = 4
+    BIG = 389
 
 class IntFlagStooges(IntFlag):
     LARRY = 1
     CURLY = 2
     MOE = 4
+    BIG = 389
 
 class IntFlagStoogesWithZero(IntFlag):
     NOFLAG = 0
     LARRY = 1
     CURLY = 2
     MOE = 4
+    BIG = 389
 
 # for pickle test and subclass tests
 class Name(StrEnum):
@@ -3252,11 +3256,17 @@ def test_pickle(self):
         test_pickle_dump_load(self.assertEqual,
                         FlagStooges.CURLY&~FlagStooges.CURLY)
         test_pickle_dump_load(self.assertIs, FlagStooges)
+        test_pickle_dump_load(self.assertEqual, FlagStooges.BIG)
+        test_pickle_dump_load(self.assertEqual,
+                        FlagStooges.CURLY|FlagStooges.BIG)
 
         test_pickle_dump_load(self.assertIs, FlagStoogesWithZero.CURLY)
         test_pickle_dump_load(self.assertEqual,
                         FlagStoogesWithZero.CURLY|FlagStoogesWithZero.MOE)
         test_pickle_dump_load(self.assertIs, FlagStoogesWithZero.NOFLAG)
+        test_pickle_dump_load(self.assertEqual, FlagStoogesWithZero.BIG)
+        test_pickle_dump_load(self.assertEqual,
+                        FlagStoogesWithZero.CURLY|FlagStoogesWithZero.BIG)
 
         test_pickle_dump_load(self.assertIs, IntFlagStooges.CURLY)
         test_pickle_dump_load(self.assertEqual,
@@ -3266,11 +3276,19 @@ def test_pickle(self):
         test_pickle_dump_load(self.assertEqual, IntFlagStooges(0))
         test_pickle_dump_load(self.assertEqual, IntFlagStooges(0x30))
         test_pickle_dump_load(self.assertIs, IntFlagStooges)
+        test_pickle_dump_load(self.assertEqual, IntFlagStooges.BIG)
+        test_pickle_dump_load(self.assertEqual, IntFlagStooges.BIG|1)
+        test_pickle_dump_load(self.assertEqual,
+                        IntFlagStooges.CURLY|IntFlagStooges.BIG)
 
         test_pickle_dump_load(self.assertIs, IntFlagStoogesWithZero.CURLY)
         test_pickle_dump_load(self.assertEqual,
                         IntFlagStoogesWithZero.CURLY|IntFlagStoogesWithZero.MOE)
         test_pickle_dump_load(self.assertIs, IntFlagStoogesWithZero.NOFLAG)
+        test_pickle_dump_load(self.assertEqual, IntFlagStoogesWithZero.BIG)
+        test_pickle_dump_load(self.assertEqual, IntFlagStoogesWithZero.BIG|1)
+        test_pickle_dump_load(self.assertEqual,
+                        IntFlagStoogesWithZero.CURLY|IntFlagStoogesWithZero.BIG)
 
     def test_contains_tf(self):
         Open = self.Open

diff --git a/Misc/NEWS.d/next/Library/2023-06-06-11-50-33.gh-issue-105332.tmpgRA.rst b/Misc/NEWS.d/next/Library/2023-06-06-11-50-33.gh-issue-105332.tmpgRA.rst
@@ -0,0 +1 @@
+Fix loading pickled data for :class:`enum.Flag` with specific values.