Close issue20534: test_enum now tests all supported pickle protocols (2 - HIGHEST_PROTOCOL, inclusive).

This commit is contained in:
Ethan Furman 2014-02-06 17:28:50 -08:00
parent 7e6a19d9a9
commit 2ddb39a695
1 changed files with 36 additions and 24 deletions

View File

@ -5,7 +5,7 @@ import unittest
from collections import OrderedDict from collections import OrderedDict
from enum import Enum, IntEnum, EnumMeta, unique from enum import Enum, IntEnum, EnumMeta, unique
from io import StringIO from io import StringIO
from pickle import dumps, loads, PicklingError from pickle import dumps, loads, PicklingError, HIGHEST_PROTOCOL
# for pickle tests # for pickle tests
try: try:
@ -61,6 +61,16 @@ try:
except Exception: except Exception:
pass pass
def test_pickle_dump_load(assertion, source, target=None):
if target is None:
target = source
for protocol in range(2, HIGHEST_PROTOCOL+1):
assertion(loads(dumps(source, protocol=protocol)), target)
def test_pickle_exception(assertion, exception, obj):
for protocol in range(2, HIGHEST_PROTOCOL+1):
with assertion(exception):
dumps(obj, protocol=protocol)
class TestHelpers(unittest.TestCase): class TestHelpers(unittest.TestCase):
# _is_descriptor, _is_sunder, _is_dunder # _is_descriptor, _is_sunder, _is_dunder
@ -503,41 +513,40 @@ class TestEnum(unittest.TestCase):
def test_pickle_enum(self): def test_pickle_enum(self):
if isinstance(Stooges, Exception): if isinstance(Stooges, Exception):
raise Stooges raise Stooges
self.assertIs(Stooges.CURLY, loads(dumps(Stooges.CURLY))) test_pickle_dump_load(self.assertIs, Stooges.CURLY)
self.assertIs(Stooges, loads(dumps(Stooges))) test_pickle_dump_load(self.assertIs, Stooges)
def test_pickle_int(self): def test_pickle_int(self):
if isinstance(IntStooges, Exception): if isinstance(IntStooges, Exception):
raise IntStooges raise IntStooges
self.assertIs(IntStooges.CURLY, loads(dumps(IntStooges.CURLY))) test_pickle_dump_load(self.assertIs, IntStooges.CURLY)
self.assertIs(IntStooges, loads(dumps(IntStooges))) test_pickle_dump_load(self.assertIs, IntStooges)
def test_pickle_float(self): def test_pickle_float(self):
if isinstance(FloatStooges, Exception): if isinstance(FloatStooges, Exception):
raise FloatStooges raise FloatStooges
self.assertIs(FloatStooges.CURLY, loads(dumps(FloatStooges.CURLY))) test_pickle_dump_load(self.assertIs, FloatStooges.CURLY)
self.assertIs(FloatStooges, loads(dumps(FloatStooges))) test_pickle_dump_load(self.assertIs, FloatStooges)
def test_pickle_enum_function(self): def test_pickle_enum_function(self):
if isinstance(Answer, Exception): if isinstance(Answer, Exception):
raise Answer raise Answer
self.assertIs(Answer.him, loads(dumps(Answer.him))) test_pickle_dump_load(self.assertIs, Answer.him)
self.assertIs(Answer, loads(dumps(Answer))) test_pickle_dump_load(self.assertIs, Answer)
def test_pickle_enum_function_with_module(self): def test_pickle_enum_function_with_module(self):
if isinstance(Question, Exception): if isinstance(Question, Exception):
raise Question raise Question
self.assertIs(Question.who, loads(dumps(Question.who))) test_pickle_dump_load(self.assertIs, Question.who)
self.assertIs(Question, loads(dumps(Question))) test_pickle_dump_load(self.assertIs, Question)
def test_exploding_pickle(self): def test_exploding_pickle(self):
BadPickle = Enum('BadPickle', 'dill sweet bread-n-butter') BadPickle = Enum('BadPickle', 'dill sweet bread-n-butter')
enum._make_class_unpicklable(BadPickle) BadPickle.__qualname__ = 'BadPickle' # needed for pickle protocol 4
globals()['BadPickle'] = BadPickle globals()['BadPickle'] = BadPickle
with self.assertRaises(TypeError): enum._make_class_unpicklable(BadPickle) # will overwrite __qualname__
dumps(BadPickle.dill) test_pickle_exception(self.assertRaises, TypeError, BadPickle.dill)
with self.assertRaises(PicklingError): test_pickle_exception(self.assertRaises, PicklingError, BadPickle)
dumps(BadPickle)
def test_string_enum(self): def test_string_enum(self):
class SkillLevel(str, Enum): class SkillLevel(str, Enum):
@ -690,7 +699,7 @@ class TestEnum(unittest.TestCase):
self.assertEqual(Name.BDFL, 'Guido van Rossum') self.assertEqual(Name.BDFL, 'Guido van Rossum')
self.assertTrue(Name.BDFL, Name('Guido van Rossum')) self.assertTrue(Name.BDFL, Name('Guido van Rossum'))
self.assertIs(Name.BDFL, getattr(Name, 'BDFL')) self.assertIs(Name.BDFL, getattr(Name, 'BDFL'))
self.assertIs(Name.BDFL, loads(dumps(Name.BDFL))) test_pickle_dump_load(self.assertIs, Name.BDFL)
def test_extending(self): def test_extending(self):
class Color(Enum): class Color(Enum):
@ -864,6 +873,7 @@ class TestEnum(unittest.TestCase):
def test_subclasses_with_getnewargs(self): def test_subclasses_with_getnewargs(self):
class NamedInt(int): class NamedInt(int):
__qualname__ = 'NamedInt' # needed for pickle protocol 4
def __new__(cls, *args): def __new__(cls, *args):
_args = args _args = args
name, *args = args name, *args = args
@ -902,6 +912,7 @@ class TestEnum(unittest.TestCase):
return temp return temp
class NEI(NamedInt, Enum): class NEI(NamedInt, Enum):
__qualname__ = 'NEI' # needed for pickle protocol 4
x = ('the-x', 1) x = ('the-x', 1)
y = ('the-y', 2) y = ('the-y', 2)
@ -912,12 +923,13 @@ class TestEnum(unittest.TestCase):
globals()['NEI'] = NEI globals()['NEI'] = NEI
NI5 = NamedInt('test', 5) NI5 = NamedInt('test', 5)
self.assertEqual(NI5, 5) self.assertEqual(NI5, 5)
self.assertEqual(loads(dumps(NI5)), 5) test_pickle_dump_load(self.assertEqual, NI5, 5)
self.assertEqual(NEI.y.value, 2) self.assertEqual(NEI.y.value, 2)
self.assertIs(loads(dumps(NEI.y)), NEI.y) test_pickle_dump_load(self.assertIs, NEI.y)
def test_subclasses_without_getnewargs(self): def test_subclasses_without_getnewargs(self):
class NamedInt(int): class NamedInt(int):
__qualname__ = 'NamedInt'
def __new__(cls, *args): def __new__(cls, *args):
_args = args _args = args
name, *args = args name, *args = args
@ -954,6 +966,7 @@ class TestEnum(unittest.TestCase):
return temp return temp
class NEI(NamedInt, Enum): class NEI(NamedInt, Enum):
__qualname__ = 'NEI'
x = ('the-x', 1) x = ('the-x', 1)
y = ('the-y', 2) y = ('the-y', 2)
@ -964,13 +977,12 @@ class TestEnum(unittest.TestCase):
NI5 = NamedInt('test', 5) NI5 = NamedInt('test', 5)
self.assertEqual(NI5, 5) self.assertEqual(NI5, 5)
self.assertEqual(NEI.y.value, 2) self.assertEqual(NEI.y.value, 2)
with self.assertRaises(TypeError): test_pickle_exception(self.assertRaises, TypeError, NEI.x)
dumps(NEI.x) test_pickle_exception(self.assertRaises, PicklingError, NEI)
with self.assertRaises(PicklingError):
dumps(NEI)
def test_tuple_subclass(self): def test_tuple_subclass(self):
class SomeTuple(tuple, Enum): class SomeTuple(tuple, Enum):
__qualname__ = 'SomeTuple' # needed for pickle protocol 4
first = (1, 'for the money') first = (1, 'for the money')
second = (2, 'for the show') second = (2, 'for the show')
third = (3, 'for the music') third = (3, 'for the music')
@ -978,7 +990,7 @@ class TestEnum(unittest.TestCase):
self.assertIsInstance(SomeTuple.second, tuple) self.assertIsInstance(SomeTuple.second, tuple)
self.assertEqual(SomeTuple.third, (3, 'for the music')) self.assertEqual(SomeTuple.third, (3, 'for the music'))
globals()['SomeTuple'] = SomeTuple globals()['SomeTuple'] = SomeTuple
self.assertIs(loads(dumps(SomeTuple.first)), SomeTuple.first) test_pickle_dump_load(self.assertIs, SomeTuple.first)
def test_duplicate_values_give_unique_enum_items(self): def test_duplicate_values_give_unique_enum_items(self):
class AutoNumber(Enum): class AutoNumber(Enum):