Issue 5665: add more pickling tests.

- Add tests for the module-level load() and dump() functions.
- Add tests for cPickle's internal data structures, stressing workloads
with many gets/puts.
- Add tests for the Pickler and Unpickler classes, in particular the
memo attribute.
- test_xpickle is extended to test backwards compatibility with Python
2.4, 2.5 and 2.6 by round-tripping pickled objects through a worker
process. This is guarded with a regrtest -u xpickle resource.

5 files changed, 428 insertions, 21 deletions

diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py
index e34d55b..4ffa702 100644
--- a/Lib/test/pickletester.py
+++ b/Lib/test/pickletester.py
@@ -1,11 +1,11 @@
 import unittest
 import pickle
 import cPickle
+import cStringIO
 import pickletools
 import copy_reg
 
-from test.test_support import TestFailed, have_unicode, TESTFN, \
-                              run_with_locale
+from test.test_support import TestFailed, have_unicode, TESTFN
 
 # Tests that try a number of pickle protocols should have a
 #     for proto in protocols:
@@ -13,6 +13,42 @@ from test.test_support import TestFailed, have_unicode, TESTFN, \
 assert pickle.HIGHEST_PROTOCOL == cPickle.HIGHEST_PROTOCOL == 2
 protocols = range(pickle.HIGHEST_PROTOCOL + 1)
 
+# Copy of test.test_support.run_with_locale. This is needed to support Python
+# 2.4, which didn't include it. This is all to support test_xpickle, which
+# bounces pickled objects through older Python versions to test backwards
+# compatibility.
+def run_with_locale(catstr, *locales):
+    def decorator(func):
+        def inner(*args, **kwds):
+            try:
+                import locale
+                category = getattr(locale, catstr)
+                orig_locale = locale.setlocale(category)
+            except AttributeError:
+                # if the test author gives us an invalid category string
+                raise
+            except:
+                # cannot retrieve original locale, so do nothing
+                locale = orig_locale = None
+            else:
+                for loc in locales:
+                    try:
+                        locale.setlocale(category, loc)
+                        break
+                    except:
+                        pass
+
+            # now run the function, resetting the locale on exceptions
+            try:
+                return func(*args, **kwds)
+            finally:
+                if locale and orig_locale:
+                    locale.setlocale(category, orig_locale)
+        inner.func_name = func.func_name
+        inner.__doc__ = func.__doc__
+        return inner
+    return decorator
+
 
 # Return True if opcode code appears in the pickle, else False.
 def opcode_in_pickle(code, pickle):
@@ -409,12 +445,11 @@ class AbstractPickleTests(unittest.TestCase):
     # is a mystery.  cPickle also suppresses PUT for objects with a refcount
     # of 1.
     def dont_test_disassembly(self):
-        from cStringIO import StringIO
         from pickletools import dis
 
         for proto, expected in (0, DATA0_DIS), (1, DATA1_DIS):
             s = self.dumps(self._testdata, proto)
-            filelike = StringIO()
+            filelike = cStringIO.StringIO()
             dis(s, out=filelike)
             got = filelike.getvalue()
             self.assertEqual(expected, got)
@@ -822,7 +857,7 @@ class AbstractPickleTests(unittest.TestCase):
         self.assertEqual(x.bar, y.bar)
 
     def test_reduce_overrides_default_reduce_ex(self):
-        for proto in 0, 1, 2:
+        for proto in protocols:
             x = REX_one()
             self.assertEqual(x._reduce_called, 0)
             s = self.dumps(x, proto)
@@ -831,7 +866,7 @@ class AbstractPickleTests(unittest.TestCase):
             self.assertEqual(y._reduce_called, 0)
 
     def test_reduce_ex_called(self):
-        for proto in 0, 1, 2:
+        for proto in protocols:
             x = REX_two()
             self.assertEqual(x._proto, None)
             s = self.dumps(x, proto)
@@ -840,7 +875,7 @@ class AbstractPickleTests(unittest.TestCase):
             self.assertEqual(y._proto, None)
 
     def test_reduce_ex_overrides_reduce(self):
-        for proto in 0, 1, 2:
+        for proto in protocols:
             x = REX_three()
             self.assertEqual(x._proto, None)
             s = self.dumps(x, proto)
@@ -849,7 +884,7 @@ class AbstractPickleTests(unittest.TestCase):
             self.assertEqual(y._proto, None)
 
     def test_reduce_ex_calls_base(self):
-        for proto in 0, 1, 2:
+        for proto in protocols:
             x = REX_four()
             self.assertEqual(x._proto, None)
             s = self.dumps(x, proto)
@@ -858,7 +893,7 @@ class AbstractPickleTests(unittest.TestCase):
             self.assertEqual(y._proto, proto)
 
     def test_reduce_calls_base(self):
-        for proto in 0, 1, 2:
+        for proto in protocols:
             x = REX_five()
             self.assertEqual(x._reduce_called, 0)
             s = self.dumps(x, proto)
@@ -879,7 +914,7 @@ class AbstractPickleTests(unittest.TestCase):
                 return dict, (), None, None, []
 
         # Protocol 0 is less strict and also accept iterables.
-        for proto in 0, 1, 2:
+        for proto in protocols:
             try:
                 self.dumps(C(), proto)
             except (AttributeError, pickle.PickleError, cPickle.PickleError):
@@ -889,6 +924,21 @@ class AbstractPickleTests(unittest.TestCase):
             except (AttributeError, pickle.PickleError, cPickle.PickleError):
                 pass
 
+    def test_many_puts_and_gets(self):
+        # Test that internal data structures correctly deal with lots of
+        # puts/gets.
+        keys = ("aaa" + str(i) for i in xrange(100))
+        large_dict = dict((k, [4, 5, 6]) for k in keys)
+        obj = [dict(large_dict), dict(large_dict), dict(large_dict)]
+
+        for proto in protocols:
+            dumped = self.dumps(obj, proto)
+            loaded = self.loads(dumped)
+            self.assertEqual(loaded, obj,
+                             "Failed protocol %d: %r != %r"
+                             % (proto, obj, loaded))
+
+
 # Test classes for reduce_ex
 
 class REX_one(object):
@@ -990,13 +1040,20 @@ class AbstractPickleModuleTests(unittest.TestCase):
         finally:
             os.remove(TESTFN)
 
+    def test_load_from_and_dump_to_file(self):
+        stream = cStringIO.StringIO()
+        data = [123, {}, 124]
+        self.module.dump(data, stream)
+        stream.seek(0)
+        unpickled = self.module.load(stream)
+        self.assertEqual(unpickled, data)
+
     def test_highest_protocol(self):
         # Of course this needs to be changed when HIGHEST_PROTOCOL changes.
         self.assertEqual(self.module.HIGHEST_PROTOCOL, 2)
 
     def test_callapi(self):
-        from cStringIO import StringIO
-        f = StringIO()
+        f = cStringIO.StringIO()
         # With and without keyword arguments
         self.module.dump(123, f, -1)
         self.module.dump(123, file=f, protocol=-1)
@@ -1039,3 +1096,116 @@ class AbstractPersistentPicklerTests(unittest.TestCase):
         self.assertEqual(self.loads(self.dumps(L, 1)), L)
         self.assertEqual(self.id_count, 5)
         self.assertEqual(self.load_count, 5)
+
+class AbstractPicklerUnpicklerObjectTests(unittest.TestCase):
+
+    pickler_class = None
+    unpickler_class = None
+
+    def setUp(self):
+        assert self.pickler_class
+        assert self.unpickler_class
+
+    def test_clear_pickler_memo(self):
+        # To test whether clear_memo() has any effect, we pickle an object,
+        # then pickle it again without clearing the memo; the two serialized
+        # forms should be different. If we clear_memo() and then pickle the
+        # object again, the third serialized form should be identical to the
+        # first one we obtained.
+        data = ["abcdefg", "abcdefg", 44]
+        f = cStringIO.StringIO()
+        pickler = self.pickler_class(f)
+
+        pickler.dump(data)
+        first_pickled = f.getvalue()
+
+        # Reset StringIO object.
+        f.seek(0)
+        f.truncate()
+
+        pickler.dump(data)
+        second_pickled = f.getvalue()
+
+        # Reset the Pickler and StringIO objects.
+        pickler.clear_memo()
+        f.seek(0)
+        f.truncate()
+
+        pickler.dump(data)
+        third_pickled = f.getvalue()
+
+        self.assertNotEqual(first_pickled, second_pickled)
+        self.assertEqual(first_pickled, third_pickled)
+
+    def test_priming_pickler_memo(self):
+        # Verify that we can set the Pickler's memo attribute.
+        data = ["abcdefg", "abcdefg", 44]
+        f = cStringIO.StringIO()
+        pickler = self.pickler_class(f)
+
+        pickler.dump(data)
+        first_pickled = f.getvalue()
+
+        f = cStringIO.StringIO()
+        primed = self.pickler_class(f)
+        primed.memo = pickler.memo
+
+        primed.dump(data)
+        primed_pickled = f.getvalue()
+
+        self.assertNotEqual(first_pickled, primed_pickled)
+
+    def test_priming_unpickler_memo(self):
+        # Verify that we can set the Unpickler's memo attribute.
+        data = ["abcdefg", "abcdefg", 44]
+        f = cStringIO.StringIO()
+        pickler = self.pickler_class(f)
+
+        pickler.dump(data)
+        first_pickled = f.getvalue()
+
+        f = cStringIO.StringIO()
+        primed = self.pickler_class(f)
+        primed.memo = pickler.memo
+
+        primed.dump(data)
+        primed_pickled = f.getvalue()
+
+        unpickler = self.unpickler_class(cStringIO.StringIO(first_pickled))
+        unpickled_data1 = unpickler.load()
+
+        self.assertEqual(unpickled_data1, data)
+
+        primed = self.unpickler_class(cStringIO.StringIO(primed_pickled))
+        primed.memo = unpickler.memo
+        unpickled_data2 = primed.load()
+
+        primed.memo.clear()
+
+        self.assertEqual(unpickled_data2, data)
+        self.assertTrue(unpickled_data2 is unpickled_data1)
+
+    def test_reusing_unpickler_objects(self):
+        data1 = ["abcdefg", "abcdefg", 44]
+        f = cStringIO.StringIO()
+        pickler = self.pickler_class(f)
+        pickler.dump(data1)
+        pickled1 = f.getvalue()
+
+        data2 = ["abcdefg", 44, 44]
+        f = cStringIO.StringIO()
+        pickler = self.pickler_class(f)
+        pickler.dump(data2)
+        pickled2 = f.getvalue()
+
+        f = cStringIO.StringIO()
+        f.write(pickled1)
+        f.seek(0)
+        unpickler = self.unpickler_class(f)
+        self.assertEqual(unpickler.load(), data1)
+
+        f.seek(0)
+        f.truncate()
+        f.write(pickled2)
+        f.seek(0)
+        self.assertEqual(unpickler.load(), data2)
diff --git a/Lib/test/regrtest.py b/Lib/test/regrtest.py
index 29f84fa..fe7f347 100755
--- a/Lib/test/regrtest.py
+++ b/Lib/test/regrtest.py
@@ -122,6 +122,10 @@ resources to test.  Currently only the following are defined:
 
     gui -       Run tests that require a running GUI.
 
+    xpickle -   Test pickle and cPickle against Python 2.4, 2.5 and 2.6 to
+                test backwards compatibility. These tests take a long time
+                to run.
+
 To enable all resources except one, use '-uall,-<resource>'.  For
 example, to run all the tests except for the bsddb tests, give the
 option '-uall,-bsddb'.
@@ -175,7 +179,8 @@ if sys.platform == 'darwin':
 from test import test_support
 
 RESOURCE_NAMES = ('audio', 'curses', 'largefile', 'network', 'bsddb',
-                  'decimal', 'compiler', 'subprocess', 'urlfetch', 'gui')
+                  'decimal', 'compiler', 'subprocess', 'urlfetch', 'gui',
+                  'xpickle')
 
 
 def usage(code, msg=''):
diff --git a/Lib/test/test_cpickle.py b/Lib/test/test_cpickle.py
index 88057c7..0541873 100644
--- a/Lib/test/test_cpickle.py
+++ b/Lib/test/test_cpickle.py
@@ -1,6 +1,7 @@
 import cPickle, unittest
 from cStringIO import StringIO
 from test.pickletester import AbstractPickleTests, AbstractPickleModuleTests
+from test.pickletester import AbstractPicklerUnpicklerObjectTests
 from test import test_support
 
 class cPickleTests(AbstractPickleTests, AbstractPickleModuleTests):
@@ -90,6 +91,12 @@ class cPickleFastPicklerTests(AbstractPickleTests):
         b = self.loads(self.dumps(a))
         self.assertEqual(a, b)
 
+class cPicklePicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):
+
+    pickler_class = cPickle.Pickler
+    unpickler_class = cPickle.Unpickler
+
+
 class Node(object):
     pass
 
@@ -120,6 +127,7 @@ def test_main():
         cPickleListPicklerTests,
         cPickleFastPicklerTests,
         cPickleDeepRecursive,
+        cPicklePicklerUnpicklerObjectTests,
     )
 
 if __name__ == "__main__":
diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py
index fb10ffe..6c83811 100644
--- a/Lib/test/test_pickle.py
+++ b/Lib/test/test_pickle.py
@@ -6,6 +6,7 @@ from test import test_support
 from test.pickletester import AbstractPickleTests
 from test.pickletester import AbstractPickleModuleTests
 from test.pickletester import AbstractPersistentPicklerTests
+from test.pickletester import AbstractPicklerUnpicklerObjectTests
 
 class PickleTests(AbstractPickleTests, AbstractPickleModuleTests):
 
@@ -60,11 +61,18 @@ class PersPicklerTests(AbstractPersistentPicklerTests):
         u = PersUnpickler(f)
         return u.load()
 
+class PicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):
+
+    pickler_class = pickle.Pickler
+    unpickler_class = pickle.Unpickler
+
+
 def test_main():
     test_support.run_unittest(
         PickleTests,
         PicklerTests,
-        PersPicklerTests
+        PersPicklerTests,
+        PicklerUnpicklerObjectTests,
     )
     test_support.run_doctest(pickle)
 
diff --git a/Lib/test/test_xpickle.py b/Lib/test/test_xpickle.py
index 3ffb744..3ed7ef7 100644
--- a/Lib/test/test_xpickle.py
+++ b/Lib/test/test_xpickle.py
@@ -1,19 +1,42 @@
 # test_pickle dumps and loads pickles via pickle.py.
 # test_cpickle does the same, but via the cPickle module.
 # This test covers the other two cases, making pickles with one module and
-# loading them via the other.
+# loading them via the other. It also tests backwards compatibility with
+# previous version of Python by bouncing pickled objects through Python 2.4
+# and Python 2.5 running this file.
 
-import pickle
 import cPickle
+import os
+import os.path
+import pickle
+import subprocess
+import sys
+import types
+import unittest
 
 from test import test_support
-from test.pickletester import AbstractPickleTests
+
+# Most distro-supplied Pythons don't include the tests
+# or test support files, and some don't include a way to get these back even if
+# you're will to install extra packages (like Ubuntu). Doing things like this
+# "provides" a pickletester module for older versions of Python that may be
+# installed without it. Note that one other design for this involves messing
+# with sys.path, which is less precise.
+mod_path = os.path.abspath(os.path.join(os.path.dirname(__file__),
+                                        "pickletester.py"))
+pickletester = types.ModuleType("test.pickletester")
+execfile(mod_path, pickletester.__dict__, pickletester.__dict__)
+AbstractPickleTests = pickletester.AbstractPickleTests
+if pickletester.__name__ in sys.modules:
+    raise RuntimeError("Did not expect to find test.pickletester loaded")
+sys.modules[pickletester.__name__] = pickletester
+
 
 class DumpCPickle_LoadPickle(AbstractPickleTests):
 
     error = KeyError
 
-    def dumps(self, arg, proto=0, fast=0):
+    def dumps(self, arg, proto=0, fast=False):
         # Ignore fast
         return cPickle.dumps(arg, proto)
 
@@ -25,7 +48,7 @@ class DumpPickle_LoadCPickle(AbstractPickleTests):
 
     error = cPickle.BadPickleGet
 
-    def dumps(self, arg, proto=0, fast=0):
+    def dumps(self, arg, proto=0, fast=False):
         # Ignore fast
         return pickle.dumps(arg, proto)
 
@@ -33,11 +56,204 @@ class DumpPickle_LoadCPickle(AbstractPickleTests):
         # Ignore fast
         return cPickle.loads(buf)
 
+def have_python_version(name):
+    """Check whether the given name is a valid Python binary.
+
+    This respects your PATH.
+
+    Args:
+        name: short string name of a Python binary such as "python2.4".
+
+    Returns:
+        True if the name is valid, False otherwise.
+    """
+    return os.system(name + " -c 'import sys; sys.exit()'") == 0
+
+
+class AbstractCompatTests(AbstractPickleTests):
+
+    module = None
+    python = None
+    error = None
+
+    def setUp(self):
+        self.assertTrue(self.python)
+        self.assertTrue(self.module)
+        self.assertTrue(self.error)
+
+    def send_to_worker(self, python, obj, proto):
+        """Bounce a pickled object through another version of Python.
+
+        This will pickle the object, send it to a child process where it will be
+        unpickled, then repickled and sent back to the parent process.
+
+        Args:
+            python: the name of the Python binary to start.
+            obj: object to pickle.
+            proto: pickle protocol number to use.
+
+        Returns:
+            The pickled data received from the child process.
+        """
+        # Prevent the subprocess from picking up invalid .pyc files.
+        target = __file__
+        if target[-1] in ("c", "o"):
+            target = target[:-1]
+
+        data = self.module.dumps((proto, obj), proto)
+        worker = subprocess.Popen([python, target, "worker"],
+                                  stdin=subprocess.PIPE,
+                                  stdout=subprocess.PIPE,
+                                  stderr=subprocess.PIPE)
+        stdout, stderr = worker.communicate(data)
+        if worker.returncode != 0:
+            raise RuntimeError(stderr)
+        return stdout
+
+    def dumps(self, arg, proto=0, fast=False):
+        return self.send_to_worker(self.python, arg, proto)
+
+    def loads(self, input):
+        return self.module.loads(input)
+
+    # These tests are disabled because they require some special setup
+    # on the worker that's hard to keep in sync.
+    def test_global_ext1(self):
+        pass
+
+    def test_global_ext2(self):
+        pass
+
+    def test_global_ext4(self):
+        pass
+
+    # This is a cut-down version of pickletester's test_float. Backwards
+    # compatibility for the values in for_bin_protos was explicitly broken in
+    # r68903 to fix a bug.
+    def test_float(self):
+        for_bin_protos = [4.94e-324, 1e-310]
+        neg_for_bin_protos = [-x for x in for_bin_protos]
+        test_values = [0.0, 7e-308, 6.626e-34, 0.1, 0.5,
+                       3.14, 263.44582062374053, 6.022e23, 1e30]
+        test_proto0_values = test_values + [-x for x in test_values]
+        test_values = test_proto0_values + for_bin_protos + neg_for_bin_protos
+
+        for value in test_proto0_values:
+            pickle = self.dumps(value, 0)
+            got = self.loads(pickle)
+            self.assertEqual(value, got)
+
+        for proto in pickletester.protocols[1:]:
+            for value in test_values:
+                pickle = self.dumps(value, proto)
+                got = self.loads(pickle)
+                self.assertEqual(value, got)
+
+    # Backwards compatibility was explicitly broken in r67934 to fix a bug.
+    def test_unicode_high_plane(self):
+        pass
+
+    if test_support.have_unicode:
+        # This is a cut-down version of pickletester's test_unicode. Backwards
+        # compatibility was explicitly broken in r67934 to fix a bug.
+        def test_unicode(self):
+            endcases = [u'', u'<\\u>', u'<\\\u1234>', u'<\n>', u'<\\>']
+            for proto in pickletester.protocols:
+                for u in endcases:
+                    p = self.dumps(u, proto)
+                    u2 = self.loads(p)
+                    self.assertEqual(u2, u)
+
+
+def run_compat_test(python_name):
+    return (test_support.is_resource_enabled("xpickle") and
+            have_python_version(python_name))
+
+
+# Test backwards compatibility with Python 2.4.
+if not run_compat_test("python2.4"):
+    class CPicklePython24Compat(unittest.TestCase):
+        pass
+else:
+    class CPicklePython24Compat(AbstractCompatTests):
+
+        module = cPickle
+        python = "python2.4"
+        error = cPickle.BadPickleGet
+
+        # Disable these tests for Python 2.4. Making them pass would require
+        # nontrivially monkeypatching the pickletester module in the worker.
+        def test_reduce_calls_base(self):
+            pass
+
+        def test_reduce_ex_calls_base(self):
+            pass
+
+class PicklePython24Compat(CPicklePython24Compat):
+
+    module = pickle
+    error = KeyError
+
+
+# Test backwards compatibility with Python 2.5.
+if not run_compat_test("python2.5"):
+    class CPicklePython25Compat(unittest.TestCase):
+        pass
+else:
+    class CPicklePython25Compat(AbstractCompatTests):
+
+        module = cPickle
+        python = "python2.5"
+        error = cPickle.BadPickleGet
+
+class PicklePython25Compat(CPicklePython25Compat):
+
+    module = pickle
+    error = KeyError
+
+
+# Test backwards compatibility with Python 2.6.
+if not run_compat_test("python2.6"):
+    class CPicklePython26Compat(unittest.TestCase):
+        pass
+else:
+    class CPicklePython26Compat(AbstractCompatTests):
+
+        module = cPickle
+        python = "python2.6"
+        error = cPickle.BadPickleGet
+
+class PicklePython26Compat(CPicklePython26Compat):
+
+    module = pickle
+    error = KeyError
+
+
+def worker_main(in_stream, out_stream):
+    message = cPickle.load(in_stream)
+    protocol, obj = message
+    cPickle.dump(obj, out_stream, protocol)
+
+
 def test_main():
+    if not test_support.is_resource_enabled("xpickle"):
+        print >>sys.stderr, "test_xpickle -- skipping backwards compat tests."
+        print >>sys.stderr, "Use 'regrtest.py -u xpickle' to run them."
+        sys.stderr.flush()
+
     test_support.run_unittest(
         DumpCPickle_LoadPickle,
-        DumpPickle_LoadCPickle
+        DumpPickle_LoadCPickle,
+        CPicklePython24Compat,
+        CPicklePython25Compat,
+        CPicklePython26Compat,
+        PicklePython24Compat,
+        PicklePython25Compat,
+        PicklePython26Compat,
     )
 
 if __name__ == "__main__":
-    test_main()
+    if "worker" in sys.argv:
+        worker_main(sys.stdin, sys.stdout)
+    else:
+        test_main()