Skip to content

Commit 1b27ec5

Browse files
[3.11] gh-93820: Pickle enum.Flag by name (GH-93891). (GH-94288)
(cherry picked from commit 5369858) Co-authored-by: Serhiy Storchaka <[email protected]>
1 parent 5ce819f commit 1b27ec5

File tree

3 files changed

+58
-3
lines changed

3 files changed

+58
-3
lines changed

Lib/enum.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1286,7 +1286,21 @@ class Flag(Enum, boundary=STRICT):
12861286
"""
12871287

12881288
def __reduce_ex__(self, proto):
1289-
return self.__class__, (self._value_, )
1289+
cls = self.__class__
1290+
unknown = self._value_ & ~cls._flag_mask_
1291+
member_value = self._value_ & cls._flag_mask_
1292+
if unknown and member_value:
1293+
return _or_, (cls(member_value), unknown)
1294+
for val in _iter_bits_lsb(member_value):
1295+
rest = member_value & ~val
1296+
if rest:
1297+
return _or_, (cls(rest), cls._value2member_map_.get(val))
1298+
else:
1299+
break
1300+
if self._name_ is None:
1301+
return cls, (self._value_,)
1302+
else:
1303+
return getattr, (cls, self._name_)
12901304

12911305
_numeric_repr_ = repr
12921306

Lib/test/test_enum.py

Lines changed: 42 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,10 +66,27 @@ class FloatStooges(float, Enum):
6666
class FlagStooges(Flag):
6767
LARRY = 1
6868
CURLY = 2
69-
MOE = 3
69+
MOE = 4
7070
except Exception as exc:
7171
FlagStooges = exc
7272

73+
class FlagStoogesWithZero(Flag):
74+
NOFLAG = 0
75+
LARRY = 1
76+
CURLY = 2
77+
MOE = 4
78+
79+
class IntFlagStooges(IntFlag):
80+
LARRY = 1
81+
CURLY = 2
82+
MOE = 4
83+
84+
class IntFlagStoogesWithZero(IntFlag):
85+
NOFLAG = 0
86+
LARRY = 1
87+
CURLY = 2
88+
MOE = 4
89+
7390
# for pickle test and subclass tests
7491
class Name(StrEnum):
7592
BDFL = 'Guido van Rossum'
@@ -2964,9 +2981,32 @@ def test_programatic_function_from_dict(self):
29642981
def test_pickle(self):
29652982
if isinstance(FlagStooges, Exception):
29662983
raise FlagStooges
2967-
test_pickle_dump_load(self.assertIs, FlagStooges.CURLY|FlagStooges.MOE)
2984+
test_pickle_dump_load(self.assertIs, FlagStooges.CURLY)
2985+
test_pickle_dump_load(self.assertEqual,
2986+
FlagStooges.CURLY|FlagStooges.MOE)
2987+
test_pickle_dump_load(self.assertEqual,
2988+
FlagStooges.CURLY&~FlagStooges.CURLY)
29682989
test_pickle_dump_load(self.assertIs, FlagStooges)
29692990

2991+
test_pickle_dump_load(self.assertIs, FlagStoogesWithZero.CURLY)
2992+
test_pickle_dump_load(self.assertEqual,
2993+
FlagStoogesWithZero.CURLY|FlagStoogesWithZero.MOE)
2994+
test_pickle_dump_load(self.assertIs, FlagStoogesWithZero.NOFLAG)
2995+
2996+
test_pickle_dump_load(self.assertIs, IntFlagStooges.CURLY)
2997+
test_pickle_dump_load(self.assertEqual,
2998+
IntFlagStooges.CURLY|IntFlagStooges.MOE)
2999+
test_pickle_dump_load(self.assertEqual,
3000+
IntFlagStooges.CURLY|IntFlagStooges.MOE|0x30)
3001+
test_pickle_dump_load(self.assertEqual, IntFlagStooges(0))
3002+
test_pickle_dump_load(self.assertEqual, IntFlagStooges(0x30))
3003+
test_pickle_dump_load(self.assertIs, IntFlagStooges)
3004+
3005+
test_pickle_dump_load(self.assertIs, IntFlagStoogesWithZero.CURLY)
3006+
test_pickle_dump_load(self.assertEqual,
3007+
IntFlagStoogesWithZero.CURLY|IntFlagStoogesWithZero.MOE)
3008+
test_pickle_dump_load(self.assertIs, IntFlagStoogesWithZero.NOFLAG)
3009+
29703010
@unittest.skipIf(
29713011
python_version >= (3, 12),
29723012
'__contains__ now returns True/False for all inputs',
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Pickle :class:`enum.Flag` by name.

0 commit comments

Comments
 (0)