8000 [Enum] improve pickle support (#26666) · python/cpython@e972631 · GitHub
[go: up one dir, main page]

Skip to content

Commit e972631

Browse files
authored
[Enum] improve pickle support (#26666)
search all bases for a __reduce__ style method; if a __new__ method is found first the enum will be made unpicklable
1 parent 3ce35bf commit e972631

File tree

3 files changed

+65
-3
lines changed

3 files changed

+65
-3
lines changed

Lib/enum.py

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -242,8 +242,32 @@ def __new__(metacls, cls, bases, classdict, **kwds):
242242
methods = ('__getnewargs_ex__', '__getnewargs__',
243243
'__reduce_ex__', '__reduce__')
244244
if not any(m in member_type.__dict__ for m in methods):
245-
_make_class_unpicklable(enum_class)
246-
245+
if '__new__' in classdict:
246+
# too late, sabotage
247+
_make_class_unpicklable(enum_class)
248+
else:
249+
# final attempt to verify that pickling would work:
250+
# travel mro until __new__ is found, checking for
251+
# __reduce__ and friends along the way -- if any of them
252+
# are found before/when __new__ is found, pickling should
253+
# work
254+
sabotage = None
255+
for chain in bases:
256+
for base in chain.__mro__:
257+
if base is object:
258+
continue
259+
elif any(m in base.__dict__ for m in methods):
260+
# found one, we're good
261+
sabotage = False
262+
break
263+
elif '__new__' in base.__dict__:
264+
# not good
265+
sabotage = True
266+
break
267+
if sabotage is not None:
268+
break
269+
if sabotage:
270+
_make_class_unpicklable(enum_class)
247271
# instantiate them, checking for duplicates as we go
248272
# we instantiate first instead of checking for duplicates first in case
249273
# a custom __new__ is doing something funky with the values -- such as
@@ -572,7 +596,7 @@ def _find_data_type(bases):
572596
data_types.add(candidate or base)
573597
break
574598
else:
575-
candidate = base
599+
candidate = candidate or base
576600
if len(data_types) > 1:
577601
raise TypeError('%r: too many data types: %r' % (class_name, data_types))
578602
elif data_types:

Lib/test/test_enum.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -594,13 +594,49 @@ class Test2Enum(MyStrEnum, MyMethodEnum):
594594

595595
def test_inherited_data_type(self):
596596
class HexInt(int):
597+
__qualname__ = 'HexInt'
597598
def __repr__(self):
598599
return hex(self)
599600
class MyEnum(HexInt, enum.Enum):
601+
__qualname__ = 'MyEnum'
600602
A = 1
601603
B = 2
602604
C = 3
603605
self.assertEqual(repr(MyEnum.A), '<MyEnum.A: 0x1>')
606+
globals()['HexInt'] = HexInt
607+
globals()['MyEnum'] = MyEnum
608+
test_pickle_dump_load(self.assertIs, MyEnum.A)
609+
test_pickle_dump_load(self.assertIs, MyEnum)
610+
#
611+
class SillyInt(HexInt):
612+
__qualname__ = 'SillyInt'
613+
pass
614+
class MyOtherEnum(SillyInt, enum.Enum):
615+
__qualname__ = 'MyOtherEnum'
616+
D = 4
617+
E = 5
618+
F = 6
619+
self.assertIs(MyOtherEnum._member_type_, SillyInt)
620+
globals()['SillyInt'] = SillyInt
621+
globals()['MyOtherEnum'] = MyOtherEnum
622+
test_pickle_dump_load(self.assertIs, MyOtherEnum.E)
623+
test_pickle_dump_load(self.assertIs, MyOtherEnum)
624+
#
625+
class BrokenInt(int):
626+
__qualname__ = 'BrokenInt'
627+
def __new__(cls, value):
628+
return int.__new__(cls, value)
629+
class MyBrokenEnum(BrokenInt, Enum):
630+
__qualname__ = 'MyBrokenEnum'
631+
G = 7
632+
H = 8
633+
I = 9
634+
self.assertIs(MyBrokenEnum._member_type_, BrokenInt)
635+
self.assertIs(MyBrokenEnum(7), MyBrokenEnum.G)
636+
globals()['BrokenInt'] = BrokenInt
637+
globals()['MyBrokenEnum'] = MyBrokenEnum
638+
test_pickle_exception(self.assertRaises, TypeError, MyBrokenEnum.G)
639+
test_pickle_exception(self.assertRaises, PicklingError, MyBrokenEnum)
604640

605641
def test_too_many_data_types(self):
606642
with self.assertRaisesRegex(TypeError, 'too many data types'):
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
[Enum] Be more robust in searching for pickle support before making an enum
2+
class unpicklable.

0 commit comments

Comments
 (0)
0