Skip to content

Commit

Permalink
[3.11] gh-93820: Pickle enum.Flag by name (GH-93891). (GH-94288)
Browse files Browse the repository at this point in the history
(cherry picked from commit 5369858)

Co-authored-by: Serhiy Storchaka <storchaka@gmail.com>
  • Loading branch information
serhiy-storchaka committed Jun 26, 2022
1 parent 5ce819f commit 1b27ec5
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 3 deletions.
16 changes: 15 additions & 1 deletion Lib/enum.py
Expand Up @@ -1286,7 +1286,21 @@ class Flag(Enum, boundary=STRICT):
"""

def __reduce_ex__(self, proto):
return self.__class__, (self._value_, )
cls = self.__class__
unknown = self._value_ & ~cls._flag_mask_
member_value = self._value_ & cls._flag_mask_
if unknown and member_value:
return _or_, (cls(member_value), unknown)
for val in _iter_bits_lsb(member_value):
rest = member_value & ~val
if rest:
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_)

_numeric_repr_ = repr

Expand Down
44 changes: 42 additions & 2 deletions Lib/test/test_enum.py
Expand Up @@ -66,10 +66,27 @@ class FloatStooges(float, Enum):
class FlagStooges(Flag):
LARRY = 1
CURLY = 2
MOE = 3
MOE = 4
except Exception as exc:
FlagStooges = exc

class FlagStoogesWithZero(Flag):
NOFLAG = 0
LARRY = 1
CURLY = 2
MOE = 4

class IntFlagStooges(IntFlag):
LARRY = 1
CURLY = 2
MOE = 4

class IntFlagStoogesWithZero(IntFlag):
NOFLAG = 0
LARRY = 1
CURLY = 2
MOE = 4

# for pickle test and subclass tests
class Name(StrEnum):
BDFL = 'Guido van Rossum'
Expand Down Expand Up @@ -2964,9 +2981,32 @@ def test_programatic_function_from_dict(self):
def test_pickle(self):
if isinstance(FlagStooges, Exception):
raise FlagStooges
test_pickle_dump_load(self.assertIs, FlagStooges.CURLY|FlagStooges.MOE)
test_pickle_dump_load(self.assertIs, FlagStooges.CURLY)
test_pickle_dump_load(self.assertEqual,
FlagStooges.CURLY|FlagStooges.MOE)
test_pickle_dump_load(self.assertEqual,
FlagStooges.CURLY&~FlagStooges.CURLY)
test_pickle_dump_load(self.assertIs, FlagStooges)

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.assertIs, IntFlagStooges.CURLY)
test_pickle_dump_load(self.assertEqual,
IntFlagStooges.CURLY|IntFlagStooges.MOE)
test_pickle_dump_load(self.assertEqual,
IntFlagStooges.CURLY|IntFlagStooges.MOE|0x30)
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.assertIs, IntFlagStoogesWithZero.CURLY)
test_pickle_dump_load(self.assertEqual,
IntFlagStoogesWithZero.CURLY|IntFlagStoogesWithZero.MOE)
test_pickle_dump_load(self.assertIs, IntFlagStoogesWithZero.NOFLAG)

@unittest.skipIf(
python_version >= (3, 12),
'__contains__ now returns True/False for all inputs',
Expand Down
@@ -0,0 +1 @@
Pickle :class:`enum.Flag` by name.

0 comments on commit 1b27ec5

Please sign in to comment.