gh-93820: Pickle enum.Flag by name (GH-93891)

This commit is contained in:
Serhiy Storchaka 2022-06-26 10:54:00 +03:00 committed by GitHub
parent c834c02569
commit 536985814a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 58 additions and 3 deletions

View File

@ -1273,7 +1273,21 @@ class Flag(Enum, boundary=STRICT):
""" """
def __reduce_ex__(self, proto): 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 _numeric_repr_ = repr

View File

@ -66,10 +66,27 @@ try:
class FlagStooges(Flag): class FlagStooges(Flag):
LARRY = 1 LARRY = 1
CURLY = 2 CURLY = 2
MOE = 3 MOE = 4
except Exception as exc: except Exception as exc:
FlagStooges = 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 # for pickle test and subclass tests
class Name(StrEnum): class Name(StrEnum):
BDFL = 'Guido van Rossum' BDFL = 'Guido van Rossum'
@ -2999,9 +3016,32 @@ class OldTestFlag(unittest.TestCase):
def test_pickle(self): def test_pickle(self):
if isinstance(FlagStooges, Exception): if isinstance(FlagStooges, Exception):
raise FlagStooges 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, 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)
def test_contains_tf(self): def test_contains_tf(self):
Open = self.Open Open = self.Open
Color = self.Color Color = self.Color

View File

@ -0,0 +1 @@
Pickle :class:`enum.Flag` by name.