diff options
-rw-r--r-- | Lib/test/pickletester.py | 64 | ||||
-rw-r--r-- | Lib/test/test_copy_reg.py | 24 | ||||
-rw-r--r-- | Modules/cPickle.c | 7 |
3 files changed, 61 insertions, 34 deletions
diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py index c6e5494..4c039bc 100644 --- a/Lib/test/pickletester.py +++ b/Lib/test/pickletester.py @@ -1,6 +1,7 @@ import unittest import pickle import pickletools +import copy_reg from test.test_support import TestFailed, have_unicode, TESTFN @@ -17,6 +18,37 @@ def opcode_in_pickle(code, pickle): return True return False +# We can't very well test the extension registry without putting known stuff +# in it, but we have to be careful to restore its original state. Code +# should do this: +# +# e = ExtensionSaver(extension_code) +# try: +# fiddle w/ the extension registry's stuff for extension_code +# finally: +# e.restore() + +class ExtensionSaver: + # Remember current registration for code (if any), and remove it (if + # there is one). + def __init__(self, code): + self.code = code + if code in copy_reg._inverted_registry: + self.pair = copy_reg._inverted_registry[code] + copy_reg.remove_extension(self.pair[0], self.pair[1], code) + else: + self.pair = None + + # Restore previous registration for code. + def restore(self): + code = self.code + curpair = copy_reg._inverted_registry.get(code) + if curpair is not None: + copy_reg.remove_extension(curpair[0], curpair[1], code) + pair = self.pair + if pair is not None: + copy_reg.add_extension(pair[0], pair[1], code) + class C: def __cmp__(self, other): return cmp(self.__dict__, other.__dict__) @@ -586,42 +618,52 @@ class AbstractPickleTests(unittest.TestCase): # Register a type with copy_reg, with extension code extcode. Pickle # an object of that type. Check that the resulting pickle uses opcode # (EXT[124]) under proto 2, and not in proto 1. + def produce_global_ext(self, extcode, opcode): - import copy_reg - copy_reg.add_extension(__name__, "MyList", extcode) + e = ExtensionSaver(extcode) try: + copy_reg.add_extension(__name__, "MyList", extcode) x = MyList([1, 2, 3]) x.foo = 42 x.bar = "hello" # Dump using protocol 1 for comparison. s1 = self.dumps(x, 1) + self.assert_(__name__ in s1) + self.assert_("MyList" in s1) + self.assertEqual(opcode_in_pickle(opcode, s1), False) + y = self.loads(s1) self.assertEqual(list(x), list(y)) self.assertEqual(x.__dict__, y.__dict__) - self.assert_(s1.find(__name__) >= 0) - self.assert_(s1.find("MyList") >= 0) # Dump using protocol 2 for test. s2 = self.dumps(x, 2) - self.assertEqual(s2.find(__name__), -1) - self.assertEqual(s2.find("MyList"), -1) + self.assert_(__name__ not in s2) + self.assert_("MyList" not in s2) + self.assertEqual(opcode_in_pickle(opcode, s2), True) + y = self.loads(s2) self.assertEqual(list(x), list(y)) self.assertEqual(x.__dict__, y.__dict__) - self.assertEqual(opcode_in_pickle(opcode, s2), True) finally: - copy_reg.remove_extension(__name__, "MyList", extcode) + e.restore() def test_global_ext1(self): - self.produce_global_ext(0xf0, pickle.EXT1) + self.produce_global_ext(0x00000001, pickle.EXT1) # smallest EXT1 code + self.produce_global_ext(0x000000ff, pickle.EXT1) # largest EXT1 code def test_global_ext2(self): - self.produce_global_ext(0xfff0, pickle.EXT2) + self.produce_global_ext(0x00000100, pickle.EXT2) # smallest EXT2 code + self.produce_global_ext(0x0000ffff, pickle.EXT2) # largest EXT2 code + self.produce_global_ext(0x0000abcd, pickle.EXT2) # check endianness def test_global_ext4(self): - self.produce_global_ext(0xabcdef0, pickle.EXT4) + self.produce_global_ext(0x00010000, pickle.EXT4) # smallest EXT4 code + self.produce_global_ext(0x7fffffff, pickle.EXT4) # largest EXT4 code + self.produce_global_ext(0x12abcdef, pickle.EXT4) # check endianness + # XXX Temporary hack, so long as the C implementation of pickle protocol # XXX 2 isn't ready. When it is, move the methods in TempAbstractPickleTests diff --git a/Lib/test/test_copy_reg.py b/Lib/test/test_copy_reg.py index 3325708..c41946a 100644 --- a/Lib/test/test_copy_reg.py +++ b/Lib/test/test_copy_reg.py @@ -1,33 +1,13 @@ import copy_reg import unittest -from test import test_support +from test import test_support +from test.pickletester import ExtensionSaver class C: pass -class ExtensionSaver: - # Remember current registration for code (if any), and remove it (if - # there is one). - def __init__(self, code): - self.code = code - if code in copy_reg._inverted_registry: - self.pair = copy_reg._inverted_registry[code] - copy_reg.remove_extension(self.pair[0], self.pair[1], code) - else: - self.pair = None - - # Restore previous registration for code. - def restore(self): - code = self.code - curpair = copy_reg._inverted_registry.get(code) - if curpair is not None: - copy_reg.remove_extension(curpair[0], curpair[1], code) - pair = self.pair - if pair is not None: - copy_reg.add_extension(pair[0], pair[1], code) - class CopyRegTestCase(unittest.TestCase): def test_class(self): diff --git a/Modules/cPickle.c b/Modules/cPickle.c index c8b194b..e726220 100644 --- a/Modules/cPickle.c +++ b/Modules/cPickle.c @@ -1946,7 +1946,7 @@ save_global(Picklerobject *self, PyObject *args, PyObject *name) * so generate an EXT opcode. */ PyObject *py_code; /* extension code as Python object */ - long code; /* extensoin code as C value */ + long code; /* extension code as C value */ char c_str[5]; int n; @@ -5280,6 +5280,11 @@ init_stuff(PyObject *module_dict) two_tuple = PyTuple_New(2); if (two_tuple == NULL) return -1; + /* We use this temp container with no regard to refcounts, or to + * keeping containees alive. Exempt from GC, because we don't + * want anything looking at two_tuple() by magic. + */ + PyObject_GC_UnTrack(two_tuple); /* Ugh */ if (!( t=PyImport_ImportModule("__builtin__"))) return -1; |