Skip to content

Commit 5369858

Browse files
gh-93820: Pickle enum.Flag by name (GH-93891)
1 parent c834c02 commit 5369858

File tree

3 files changed

+58
-3
lines changed

3 files changed

+58
-3
lines changed

Lib/enum.py

+15-1
Original file line numberDiff line numberDiff line change
@@ -1273,7 +1273,21 @@ class Flag(Enum, boundary=STRICT):
12731273
"""
12741274

12751275
def __reduce_ex__(self, proto):
1276-
return self.__class__, (self._value_, )
1276+
cls = self.__class__
1277+
unknown = self._value_ & ~cls._flag_mask_
1278+
member_value = self._value_ & cls._flag_mask_
1279+
if unknown and member_value:
1280+
return _or_, (cls(member_value), unknown)
1281+
for val in _iter_bits_lsb(member_value):
1282+
rest = member_value & ~val
1283+
if rest:
1284+
return _or_, (cls(rest), cls._value2member_map_.get(val))
1285+
else:
1286+
break
1287+
if self._name_ is None:
1288+
return cls, (self._value_,)
1289+
else:
1290+
return getattr, (cls, self._name_)
12771291

12781292
_numeric_repr_ = repr
12791293

Lib/test/test_enum.py

+42-2
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'
@@ -2999,9 +3016,32 @@ def test_programatic_function_from_dict(self):
29993016
def test_pickle(self):
30003017
if isinstance(FlagStooges, Exception):
30013018
raise FlagStooges
3002-
test_pickle_dump_load(self.assertIs, FlagStooges.CURLY|FlagStooges.MOE)
3019+
test_pickle_dump_load(self.assertIs, FlagStooges.CURLY)
3020+
test_pickle_dump_load(self.assertEqual,
3021+
FlagStooges.CURLY|FlagStooges.MOE)
3022+
test_pickle_dump_load(self.assertEqual,
3023+
FlagStooges.CURLY&~FlagStooges.CURLY)
30033024
test_pickle_dump_load(self.assertIs, FlagStooges)
30043025

3026+
test_pickle_dump_load(self.assertIs, FlagStoogesWithZero.CURLY)
3027+
test_pickle_dump_load(self.assertEqual,
3028+
FlagStoogesWithZero.CURLY|FlagStoogesWithZero.MOE)
3029+
test_pickle_dump_load(self.assertIs, FlagStoogesWithZero.NOFLAG)
3030+
3031+
test_pickle_dump_load(self.assertIs, IntFlagStooges.CURLY)
3032+
test_pickle_dump_load(self.assertEqual,
3033+
IntFlagStooges.CURLY|IntFlagStooges.MOE)
3034+
test_pickle_dump_load(self.assertEqual,
3035+
IntFlagStooges.CURLY|IntFlagStooges.MOE|0x30)
3036+
test_pickle_dump_load(self.assertEqual, IntFlagStooges(0))
3037+
test_pickle_dump_load(self.assertEqual, IntFlagStooges(0x30))
3038+
test_pickle_dump_load(self.assertIs, IntFlagStooges)
3039+
3040+
test_pickle_dump_load(self.assertIs, IntFlagStoogesWithZero.CURLY)
3041+
test_pickle_dump_load(self.assertEqual,
3042+
IntFlagStoogesWithZero.CURLY|IntFlagStoogesWithZero.MOE)
3043+
test_pickle_dump_load(self.assertIs, IntFlagStoogesWithZero.NOFLAG)
3044+
30053045
def test_contains_tf(self):
30063046
Open = self.Open
30073047
Color = self.Color
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Pickle :class:`enum.Flag` by name.

0 commit comments

Comments
 (0)