This commit is contained in:
Charles-François Natali 2014-02-08 22:55:13 +01:00
commit 2c66918462
3 changed files with 233 additions and 21 deletions

View File

@ -369,10 +369,10 @@ The usual restrictions for pickling apply: picklable enums must be defined in
the top level of a module, since unpickling requires them to be importable
from that module.
.. warning::
.. note::
In order to support the singleton nature of enumeration members, pickle
protocol version 2 or higher must be used.
With pickle protocol version 4 it is possible to easily pickle enums
nested in other classes.
Functional API
@ -420,6 +420,14 @@ The solution is to specify the module name explicitly as follows::
>>> Animals = Enum('Animals', 'ant bee cat dog', module=__name__)
The new pickle protocol 4 also, in some circumstances, relies on
:attr:``__qualname__`` being set to the location where pickle will be able
to find the class. For example, if the class was made available in class
SomeData in the global scope::
>>> Animals = Enum('Animals', 'ant bee cat dog', qualname='SomeData.Animals')
Derived Enumerations
--------------------

View File

@ -31,9 +31,9 @@ def _is_sunder(name):
def _make_class_unpicklable(cls):
"""Make the given class un-picklable."""
def _break_on_call_reduce(self):
def _break_on_call_reduce(self, proto):
raise TypeError('%r cannot be pickled' % self)
cls.__reduce__ = _break_on_call_reduce
cls.__reduce_ex__ = _break_on_call_reduce
cls.__module__ = '<unknown>'
@ -115,12 +115,13 @@ class EnumMeta(type):
# Reverse value->name map for hashable values.
enum_class._value2member_map_ = {}
# check for a __getnewargs__, and if not present sabotage
# check for a supported pickle protocols, and if not present sabotage
# pickling, since it won't work anyway
if (member_type is not object and
member_type.__dict__.get('__getnewargs__') is None
):
_make_class_unpicklable(enum_class)
if member_type is not object:
methods = ('__getnewargs_ex__', '__getnewargs__',
'__reduce_ex__', '__reduce__')
if not any(map(member_type.__dict__.get, methods)):
_make_class_unpicklable(enum_class)
# instantiate them, checking for duplicates as we go
# we instantiate first instead of checking for duplicates first in case
@ -166,7 +167,7 @@ class EnumMeta(type):
# double check that repr and friends are not the mixin's or various
# things break (such as pickle)
for name in ('__repr__', '__str__', '__format__', '__getnewargs__'):
for name in ('__repr__', '__str__', '__format__', '__getnewargs__', '__reduce_ex__'):
class_method = getattr(enum_class, name)
obj_method = getattr(member_type, name, None)
enum_method = getattr(first_enum, name, None)
@ -183,7 +184,7 @@ class EnumMeta(type):
enum_class.__new__ = Enum.__new__
return enum_class
def __call__(cls, value, names=None, *, module=None, type=None):
def __call__(cls, value, names=None, *, module=None, qualname=None, type=None):
"""Either returns an existing member, or creates a new enum class.
This method is used both when an enum class is given a value to match
@ -202,7 +203,7 @@ class EnumMeta(type):
if names is None: # simple value lookup
return cls.__new__(cls, value)
# otherwise, functional API: we're creating a new Enum type
return cls._create_(value, names, module=module, type=type)
return cls._create_(value, names, module=module, qualname=qualname, type=type)
def __contains__(cls, member):
return isinstance(member, cls) and member.name in cls._member_map_
@ -273,7 +274,7 @@ class EnumMeta(type):
raise AttributeError('Cannot reassign members.')
super().__setattr__(name, value)
def _create_(cls, class_name, names=None, *, module=None, type=None):
def _create_(cls, class_name, names=None, *, module=None, qualname=None, type=None):
"""Convenience method to create a new Enum class.
`names` can be:
@ -315,6 +316,8 @@ class EnumMeta(type):
_make_class_unpicklable(enum_class)
else:
enum_class.__module__ = module
if qualname is not None:
enum_class.__qualname__ = qualname
return enum_class
@ -468,6 +471,9 @@ class Enum(metaclass=EnumMeta):
def __hash__(self):
return hash(self._name_)
def __reduce_ex__(self, proto):
return self.__class__, self.__getnewargs__()
# DynamicClassAttribute is used to provide access to the `name` and
# `value` properties of enum members while keeping some measure of
# protection from modification, while still allowing for an enumeration

View File

@ -52,6 +52,11 @@ try:
except Exception as exc:
Answer = exc
try:
Theory = Enum('Theory', 'rule law supposition', qualname='spanish_inquisition')
except Exception as exc:
Theory = exc
# for doctests
try:
class Fruit(Enum):
@ -61,14 +66,18 @@ try:
except Exception:
pass
def test_pickle_dump_load(assertion, source, target=None):
def test_pickle_dump_load(assertion, source, target=None,
*, protocol=(0, HIGHEST_PROTOCOL)):
start, stop = protocol
if target is None:
target = source
for protocol in range(2, HIGHEST_PROTOCOL+1):
for protocol in range(start, stop+1):
assertion(loads(dumps(source, protocol=protocol)), target)
def test_pickle_exception(assertion, exception, obj):
for protocol in range(2, HIGHEST_PROTOCOL+1):
def test_pickle_exception(assertion, exception, obj,
*, protocol=(0, HIGHEST_PROTOCOL)):
start, stop = protocol
for protocol in range(start, stop+1):
with assertion(exception):
dumps(obj, protocol=protocol)
@ -101,6 +110,7 @@ class TestHelpers(unittest.TestCase):
class TestEnum(unittest.TestCase):
def setUp(self):
class Season(Enum):
SPRING = 1
@ -540,11 +550,31 @@ class TestEnum(unittest.TestCase):
test_pickle_dump_load(self.assertIs, Question.who)
test_pickle_dump_load(self.assertIs, Question)
def test_enum_function_with_qualname(self):
if isinstance(Theory, Exception):
raise Theory
self.assertEqual(Theory.__qualname__, 'spanish_inquisition')
def test_class_nested_enum_and_pickle_protocol_four(self):
# would normally just have this directly in the class namespace
class NestedEnum(Enum):
twigs = 'common'
shiny = 'rare'
self.__class__.NestedEnum = NestedEnum
self.NestedEnum.__qualname__ = '%s.NestedEnum' % self.__class__.__name__
test_pickle_exception(
self.assertRaises, PicklingError, self.NestedEnum.twigs,
protocol=(0, 3))
test_pickle_dump_load(self.assertIs, self.NestedEnum.twigs,
protocol=(4, HIGHEST_PROTOCOL))
def test_exploding_pickle(self):
BadPickle = Enum('BadPickle', 'dill sweet bread-n-butter')
BadPickle.__qualname__ = 'BadPickle' # needed for pickle protocol 4
BadPickle = Enum(
'BadPickle', 'dill sweet bread-n-butter', module=__name__)
globals()['BadPickle'] = BadPickle
enum._make_class_unpicklable(BadPickle) # will overwrite __qualname__
# now break BadPickle to test exception raising
enum._make_class_unpicklable(BadPickle)
test_pickle_exception(self.assertRaises, TypeError, BadPickle.dill)
test_pickle_exception(self.assertRaises, PicklingError, BadPickle)
@ -927,6 +957,174 @@ class TestEnum(unittest.TestCase):
self.assertEqual(NEI.y.value, 2)
test_pickle_dump_load(self.assertIs, NEI.y)
def test_subclasses_with_getnewargs_ex(self):
class NamedInt(int):
__qualname__ = 'NamedInt' # needed for pickle protocol 4
def __new__(cls, *args):
_args = args
name, *args = args
if len(args) == 0:
raise TypeError("name and value must be specified")
self = int.__new__(cls, *args)
self._intname = name
self._args = _args
return self
def __getnewargs_ex__(self):
return self._args, {}
@property
def __name__(self):
return self._intname
def __repr__(self):
# repr() is updated to include the name and type info
return "{}({!r}, {})".format(type(self).__name__,
self.__name__,
int.__repr__(self))
def __str__(self):
# str() is unchanged, even if it relies on the repr() fallback
base = int
base_str = base.__str__
if base_str.__objclass__ is object:
return base.__repr__(self)
return base_str(self)
# for simplicity, we only define one operator that
# propagates expressions
def __add__(self, other):
temp = int(self) + int( other)
if isinstance(self, NamedInt) and isinstance(other, NamedInt):
return NamedInt(
'({0} + {1})'.format(self.__name__, other.__name__),
temp )
else:
return temp
class NEI(NamedInt, Enum):
__qualname__ = 'NEI' # needed for pickle protocol 4
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
globals()['NEI'] = NEI
NI5 = NamedInt('test', 5)
self.assertEqual(NI5, 5)
test_pickle_dump_load(self.assertEqual, NI5, 5, protocol=(4, 4))
self.assertEqual(NEI.y.value, 2)
test_pickle_dump_load(self.assertIs, NEI.y, protocol=(4, 4))
def test_subclasses_with_reduce(self):
class NamedInt(int):
__qualname__ = 'NamedInt' # needed for pickle protocol 4
def __new__(cls, *args):
_args = args
name, *args = args
if len(args) == 0:
raise TypeError("name and value must be specified")
self = int.__new__(cls, *args)
self._intname = name
self._args = _args
return self
def __reduce__(self):
return self.__class__, self._args
@property
def __name__(self):
return self._intname
def __repr__(self):
# repr() is updated to include the name and type info
return "{}({!r}, {})".format(type(self).__name__,
self.__name__,
int.__repr__(self))
def __str__(self):
# str() is unchanged, even if it relies on the repr() fallback
base = int
base_str = base.__str__
if base_str.__objclass__ is object:
return base.__repr__(self)
return base_str(self)
# for simplicity, we only define one operator that
# propagates expressions
def __add__(self, other):
temp = int(self) + int( other)
if isinstance(self, NamedInt) and isinstance(other, NamedInt):
return NamedInt(
'({0} + {1})'.format(self.__name__, other.__name__),
temp )
else:
return temp
class NEI(NamedInt, Enum):
__qualname__ = 'NEI' # needed for pickle protocol 4
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
globals()['NEI'] = NEI
NI5 = NamedInt('test', 5)
self.assertEqual(NI5, 5)
test_pickle_dump_load(self.assertEqual, NI5, 5)
self.assertEqual(NEI.y.value, 2)
test_pickle_dump_load(self.assertIs, NEI.y)
def test_subclasses_with_reduce_ex(self):
class NamedInt(int):
__qualname__ = 'NamedInt' # needed for pickle protocol 4
def __new__(cls, *args):
_args = args
name, *args = args
if len(args) == 0:
raise TypeError("name and value must be specified")
self = int.__new__(cls, *args)
self._intname = name
self._args = _args
return self
def __reduce_ex__(self, proto):
return self.__class__, self._args
@property
def __name__(self):
return self._intname
def __repr__(self):
# repr() is updated to include the name and type info
return "{}({!r}, {})".format(type(self).__name__,
self.__name__,
int.__repr__(self))
def __str__(self):
# str() is unchanged, even if it relies on the repr() fallback
base = int
base_str = base.__str__
if base_str.__objclass__ is object:
return base.__repr__(self)
return base_str(self)
# for simplicity, we only define one operator that
# propagates expressions
def __add__(self, other):
temp = int(self) + int( other)
if isinstance(self, NamedInt) and isinstance(other, NamedInt):
return NamedInt(
'({0} + {1})'.format(self.__name__, other.__name__),
temp )
else:
return temp
class NEI(NamedInt, Enum):
__qualname__ = 'NEI' # needed for pickle protocol 4
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
globals()['NEI'] = NEI
NI5 = NamedInt('test', 5)
self.assertEqual(NI5, 5)
test_pickle_dump_load(self.assertEqual, NI5, 5)
self.assertEqual(NEI.y.value, 2)
test_pickle_dump_load(self.assertIs, NEI.y)
def test_subclasses_without_getnewargs(self):
class NamedInt(int):
__qualname__ = 'NamedInt'