[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
This commit is contained in:
Ethan Furman 2021-06-11 01:26:32 -07:00 committed by GitHub
parent 3ce35bfbbe
commit e9726314df
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 65 additions and 3 deletions

View file

@ -242,8 +242,32 @@ class EnumMeta(type):
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
@ -572,7 +596,7 @@ class EnumMeta(type):
data_types.add(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: