gh-105332: [Enum] Fix unpickling flags in edge-cases (GH-105348)

* revert enum pickling from by-name to by-value

Co-authored-by: Ethan Furman <ethan@stoneleaf.us>
This commit is contained in:
Nikita Sobolev 2023-06-08 21:40:15 +03:00 committed by GitHub
parent e822a676f1
commit 4ff5690e59
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 47 additions and 23 deletions

View file

@ -517,7 +517,16 @@ from that module.
nested in other classes.
It is possible to modify how enum members are pickled/unpickled by defining
:meth:`__reduce_ex__` in the enumeration class.
:meth:`__reduce_ex__` in the enumeration class. The default method is by-value,
but enums with complicated values may want to use by-name::
>>> class MyEnum(Enum):
... __reduce_ex__ = enum.pickle_by_enum_name
.. note::
Using by-name for flags is not recommended, as unnamed aliases will
not unpickle.
Functional API

View file

@ -12,6 +12,7 @@ __all__ = [
'FlagBoundary', 'STRICT', 'CONFORM', 'EJECT', 'KEEP',
'global_flag_repr', 'global_enum_repr', 'global_str', 'global_enum',
'EnumCheck', 'CONTINUOUS', 'NAMED_FLAGS', 'UNIQUE',
'pickle_by_global_name', 'pickle_by_enum_name',
]
@ -913,7 +914,6 @@ class EnumType(type):
body['__module__'] = module
tmp_cls = type(name, (object, ), body)
cls = _simple_enum(etype=cls, boundary=boundary or KEEP)(tmp_cls)
cls.__reduce_ex__ = _reduce_ex_by_global_name
if as_global:
global_enum(cls)
else:
@ -1216,7 +1216,7 @@ class Enum(metaclass=EnumType):
return hash(self._name_)
def __reduce_ex__(self, proto):
return getattr, (self.__class__, self._name_)
return self.__class__, (self._value_, )
# enum.property is used to provide access to the `name` and
# `value` attributes of enum members while keeping some measure of
@ -1283,8 +1283,14 @@ class StrEnum(str, ReprEnum):
return name.lower()
def _reduce_ex_by_global_name(self, proto):
def pickle_by_global_name(self, proto):
# should not be used with Flag-type enums
return self.name
_reduce_ex_by_global_name = pickle_by_global_name
def pickle_by_enum_name(self, proto):
# should not be used with Flag-type enums
return getattr, (self.__class__, self._name_)
class FlagBoundary(StrEnum):
"""
@ -1306,23 +1312,6 @@ class Flag(Enum, boundary=STRICT):
Support for flags
"""
def __reduce_ex__(self, proto):
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
@staticmethod
@ -2049,7 +2038,6 @@ def _old_convert_(etype, name, module, filter, source=None, *, boundary=None):
# unless some values aren't comparable, in which case sort by name
members.sort(key=lambda t: t[0])
cls = etype(name, members, module=module, boundary=boundary or KEEP)
cls.__reduce_ex__ = _reduce_ex_by_global_name
return cls
_stdlib_enums = IntEnum, StrEnum, IntFlag

View file

@ -31,6 +31,11 @@ def load_tests(loader, tests, ignore):
'../../Doc/library/enum.rst',
optionflags=doctest.ELLIPSIS|doctest.NORMALIZE_WHITESPACE,
))
if os.path.exists('Doc/howto/enum.rst'):
tests.addTests(doctest.DocFileSuite(
'../../Doc/howto/enum.rst',
optionflags=doctest.ELLIPSIS|doctest.NORMALIZE_WHITESPACE,
))
return tests
MODULE = __name__
@ -66,6 +71,7 @@ try:
LARRY = 1
CURLY = 2
MOE = 4
BIG = 389
except Exception as exc:
FlagStooges = exc
@ -74,17 +80,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):
@ -1942,7 +1951,6 @@ class TestSpecial(unittest.TestCase):
__qualname__ = 'NEI'
x = ('the-x', 1)
y = ('the-y', 2)
self.assertIs(NEI.__new__, Enum.__new__)
self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)")
globals()['NamedInt'] = NamedInt
@ -1950,6 +1958,10 @@ class TestSpecial(unittest.TestCase):
NI5 = NamedInt('test', 5)
self.assertEqual(NI5, 5)
self.assertEqual(NEI.y.value, 2)
with self.assertRaisesRegex(TypeError, "name and value must be specified"):
test_pickle_dump_load(self.assertIs, NEI.y)
# fix pickle support and try again
NEI.__reduce_ex__ = enum.pickle_by_enum_name
test_pickle_dump_load(self.assertIs, NEI.y)
test_pickle_dump_load(self.assertIs, NEI)
@ -3252,11 +3264,17 @@ class OldTestFlag(unittest.TestCase):
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 +3284,19 @@ class OldTestFlag(unittest.TestCase):
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

View file

@ -0,0 +1 @@
Revert pickling method from by-name back to by-value.