diff options
Diffstat (limited to 'Lib')
| -rw-r--r-- | Lib/enum.py | 26 | ||||
| -rw-r--r-- | Lib/test/test_enum.py | 61 |
2 files changed, 73 insertions, 14 deletions
diff --git a/Lib/enum.py b/Lib/enum.py index 794f68e..c9bd7c0 100644 --- a/Lib/enum.py +++ b/Lib/enum.py @@ -116,12 +116,14 @@ class EnumMeta(type): enum_class._value2member_map_ = {} # check for a supported pickle protocols, and if not present sabotage - # pickling, since it won't work anyway - 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) + # pickling, since it won't work anyway. + # if new class implements its own __reduce_ex__, do not sabotage + if classdict.get('__reduce_ex__') is None: + 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 @@ -167,7 +169,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__', '__reduce_ex__'): + for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'): class_method = getattr(enum_class, name) obj_method = getattr(member_type, name, None) enum_method = getattr(first_enum, name, None) @@ -192,8 +194,9 @@ class EnumMeta(type): (i.e. Color = Enum('Color', names='red green blue')). When used for the functional API: `module`, if set, will be stored in - the new class' __module__ attribute; `type`, if set, will be mixed in - as the first base class. + the new class' __module__ attribute; `qualname`, if set, will be stored + in the new class' __qualname__ attribute; `type`, if set, will be mixed + in as the first base class. Note: if `module` is not set this routine will attempt to discover the calling module by walking the frame stack; if this is unsuccessful @@ -465,14 +468,11 @@ class Enum(metaclass=EnumMeta): val = self.value return cls.__format__(val, format_spec) - def __getnewargs__(self): - return (self._value_, ) - def __hash__(self): return hash(self._name_) def __reduce_ex__(self, proto): - return self.__class__, self.__getnewargs__() + return self.__class__, (self._value_, ) # DynamicClassAttribute is used to provide access to the `name` and # `value` properties of enum members while keeping some measure of diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py index 02009af..b8ef632 100644 --- a/Lib/test/test_enum.py +++ b/Lib/test/test_enum.py @@ -956,6 +956,7 @@ class TestEnum(unittest.TestCase): test_pickle_dump_load(self.assertEqual, NI5, 5) self.assertEqual(NEI.y.value, 2) test_pickle_dump_load(self.assertIs, NEI.y) + test_pickle_dump_load(self.assertIs, NEI) def test_subclasses_with_getnewargs_ex(self): class NamedInt(int): @@ -1012,6 +1013,7 @@ class TestEnum(unittest.TestCase): 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)) + test_pickle_dump_load(self.assertIs, NEI) def test_subclasses_with_reduce(self): class NamedInt(int): @@ -1068,6 +1070,7 @@ class TestEnum(unittest.TestCase): test_pickle_dump_load(self.assertEqual, NI5, 5) self.assertEqual(NEI.y.value, 2) test_pickle_dump_load(self.assertIs, NEI.y) + test_pickle_dump_load(self.assertIs, NEI) def test_subclasses_with_reduce_ex(self): class NamedInt(int): @@ -1124,8 +1127,9 @@ class TestEnum(unittest.TestCase): test_pickle_dump_load(self.assertEqual, NI5, 5) self.assertEqual(NEI.y.value, 2) test_pickle_dump_load(self.assertIs, NEI.y) + test_pickle_dump_load(self.assertIs, NEI) - def test_subclasses_without_getnewargs(self): + def test_subclasses_without_direct_pickle_support(self): class NamedInt(int): __qualname__ = 'NamedInt' def __new__(cls, *args): @@ -1178,6 +1182,61 @@ class TestEnum(unittest.TestCase): test_pickle_exception(self.assertRaises, TypeError, NEI.x) test_pickle_exception(self.assertRaises, PicklingError, NEI) + def test_subclasses_without_direct_pickle_support_using_name(self): + class NamedInt(int): + __qualname__ = 'NamedInt' + 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 + @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' + x = ('the-x', 1) + y = ('the-y', 2) + def __reduce_ex__(self, proto): + return getattr, (self.__class__, self._name_) + + 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) + self.assertEqual(NEI.y.value, 2) + test_pickle_dump_load(self.assertIs, NEI.y) + test_pickle_dump_load(self.assertIs, NEI) + def test_tuple_subclass(self): class SomeTuple(tuple, Enum): __qualname__ = 'SomeTuple' # needed for pickle protocol 4 |