8000 [3.9] bpo-44342: [Enum] improve pickle support by ethanfurman · Pull Request #26666 · python/cpython · GitHub
[go: up one dir, main page]

Skip to content

[3.9] bpo-44342: [Enum] improve pickle support #26666

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 11, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 27 additions & 3 deletions Lib/enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -242,8 +242,32 @@ def __new__(metacls, cls, bases, classdict, **kwds):
methods = ('__getnewargs_ex__', '__getnewargs__',
'__reduce_ex__', '__reduce__')
if not any(m in member_type.__dict__ for m in methods):
_make_class_unpicklable(enum_class)

if '__new__' in classdict:
# too late, sabotage
_make_class_unpicklable(enum_class)
else:
# final attempt to verify that pickling would work:
# travel mro until __new__ is found, checking for
# __reduce__ and friends along the way -- if any of them
# are found before/when __new__ is found, pickling should
# work
sabotage = None
for chain in bases:
for base in chain.__mro__:
if base is object:
continue
elif any(m in base.__dict__ for m in methods):
# found one, we're good
sabotage = False
break
elif '__new__' in base.__dict__:
# not good
sabotage = True
break
if sabotage is not None:
break
if sabotage:
_make_class_unpicklable(enum_class)
# instantiate them, checking for duplicates as we go
# we instantiate first instead of checking for duplicates first in case
# a custom __new__ is doing something funky with the values -- such as
Expand Down Expand Up @@ -572,7 +596,7 @@ def _find_data_type(bases):
data_types.append(candidate or base)
break
else:
candidate = base
candidate = candidate or base
if len(data_types) > 1:
raise TypeError('%r: too many data types: %r' % (class_name, data_types))
elif data_types:
Expand Down
36 changes: 36 additions & 0 deletions Lib/test/test_enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -594,13 +594,49 @@ class Test2Enum(MyStrEnum, MyMethodEnum):

def test_inherited_data_type(self):
class HexInt(int):
__qualname__ = 'HexInt'
def __repr__(self):
return hex(self)
class MyEnum(HexInt, enum.Enum):
__qualname__ = 'MyEnum'
A = 1
B = 2
C = 3
self.assertEqual(repr(MyEnum.A), '<MyEnum.A: 0x1>')
globals()['HexInt'] = HexInt
globals()['MyEnum'] = MyEnum
test_pickle_dump_load(self.assertIs, MyEnum.A)
test_pickle_dump_load(self.assertIs, MyEnum)
#
class SillyInt(HexInt):
__qualname__ = 'SillyInt'
pass
class MyOtherEnum(SillyInt, enum.Enum):
__qualname__ = 'MyOtherEnum'
D = 4
E = 5
F = 6
self.assertIs(MyOtherEnum._member_type_, SillyInt)
globals()['SillyInt'] = SillyInt
globals()['MyOtherEnum'] = MyOtherEnum
test_pickle_dump_load(self.assertIs, MyOtherEnum.E)
test_pickle_dump_load(self.assertIs, MyOtherEnum)
#
class BrokenInt(int):
__qualname__ = 'BrokenInt'
def __new__(cls, value):
return int.__new__(cls, value)
class MyBrokenEnum(BrokenInt, Enum):
__qualname__ = 'MyBrokenEnum'
G = 7
H = 8
I = 9
self.assertIs(MyBrokenEnum._member_type_, BrokenInt)
self.assertIs(MyBrokenEnum(7), MyBrokenEnum.G)
globals()['BrokenInt'] = BrokenInt
globals()['MyBrokenEnum'] = MyBrokenEnum
test_pickle_exception(self.assertRaises, TypeError, MyBrokenEnum.G)
test_pickle_exception(self.assertRaises, PicklingError, MyBrokenEnum)

def test_too_many_data_types(self):
with self.assertRaisesRegex(TypeError, 'too many data types'):
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
[Enum] Be more robust in searching for pickle support before making an enum
class unpicklable.
0