cPickle now generates proto 2 EXT[124] when appropriate.

Moved such EXT tests as currently exist from TempAbstractPickleTests to
AbstractPickleTests, so that test_cpickle runs them too.

diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py
index 87acec6aa2fe73ca3e20f52be3c9fa30a5956e78..c6e549435d3507d73a7eed77a4aa0e12ca3f3961 100644 (file)
--- a/Lib/test/pickletester.py
+++ b/Lib/test/pickletester.py
@@ -583,23 +583,6 @@ class AbstractPickleTests(unittest.TestCase):
                 self.assertEqual(B(x), B(y), detail)
                 self.assertEqual(x.__dict__, y.__dict__, detail)
 
-# 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
-# XXX into AbstractPickleTests above, and get rid of TempAbstractPickleTests
-# XXX along with the references to it in test_pickle.py.
-class TempAbstractPickleTests(unittest.TestCase):
-
-    def test_newobj_list_slots(self):
-        x = SlotList([1, 2, 3])
-        x.foo = 42
-        x.bar = "hello"
-        s = self.dumps(x, 2)
-        y = self.loads(s)
-        self.assertEqual(list(x), list(y))
-        self.assertEqual(x.__dict__, y.__dict__)
-        self.assertEqual(x.foo, y.foo)
-        self.assertEqual(x.bar, y.bar)
-
     # 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.
@@ -638,7 +621,24 @@ class TempAbstractPickleTests(unittest.TestCase):
         self.produce_global_ext(0xfff0, pickle.EXT2)
 
     def test_global_ext4(self):
-        self.produce_global_ext(0xffffff0, pickle.EXT4)
+        self.produce_global_ext(0xabcdef0, pickle.EXT4)
+
+# 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
+# XXX into AbstractPickleTests above, and get rid of TempAbstractPickleTests
+# XXX along with the references to it in test_pickle.py.
+class TempAbstractPickleTests(unittest.TestCase):
+
+    def test_newobj_list_slots(self):
+        x = SlotList([1, 2, 3])
+        x.foo = 42
+        x.bar = "hello"
+        s = self.dumps(x, 2)
+        y = self.loads(s)
+        self.assertEqual(list(x), list(y))
+        self.assertEqual(x.__dict__, y.__dict__)
+        self.assertEqual(x.foo, y.foo)
+        self.assertEqual(x.bar, y.bar)
 
 
 class MyInt(int):

diff --git a/Modules/cPickle.c b/Modules/cPickle.c
index 6f09c1797b72550f65331c8d692774d5787461f3..c8b194ba40271156e42d29c2f72a715fcc62edd7 100644 (file)
--- a/Modules/cPickle.c
+++ b/Modules/cPickle.c
@@ -109,6 +109,9 @@ static PyObject *inverted_registry;
 /* copy_reg._extension_cache, {code: object} */
 static PyObject *extension_cache;
 
+/* For looking up name pairs in copy_reg._extension_registry. */
+static PyObject *two_tuple;
+
 static PyObject *__class___str, *__getinitargs___str, *__dict___str,
   *__getstate___str, *__setstate___str, *__name___str, *__reduce___str,
   *write_str, *append_str,
@@ -1931,12 +1934,70 @@ save_global(Picklerobject *self, PyObject *args, PyObject *name)
        if (klass != args) {
                Py_DECREF(klass);
                cPickle_ErrFormat(PicklingError,
-                                 "Can't pickle %s: it's not the same object as %s.%s",
+                                 "Can't pickle %s: it's not the same object "
+                                       "as %s.%s",
                                  "OSS", args, module, global_name);
                goto finally;
        }
        Py_DECREF(klass);
 
+       if (self->proto >= 2) {
+               /* See whether this is in the extension registry, and if
+                * so generate an EXT opcode.
+                */
+               PyObject *py_code;      /* extension code as Python object */
+               long code;              /* extensoin code as C value */
+               char c_str[5];
+               int n;
+
+               PyTuple_SET_ITEM(two_tuple, 0, module);
+               PyTuple_SET_ITEM(two_tuple, 1, global_name);
+               py_code = PyDict_GetItem(extension_registry, two_tuple);
+               if (py_code == NULL)
+                       goto gen_global;        /* not registered */
+
+               /* Verify py_code has the right type and value. */
+               if (!PyInt_Check(py_code)) {
+                       cPickle_ErrFormat(PicklingError, "Can't pickle %s: "
+                               "extension code %s isn't n integer",
+                               "OO", args, py_code);
+                       goto finally;
+               }
+               code = PyInt_AS_LONG(py_code);
+               if (code <= 0 ||  code > 0x7fffffffL) {
+                       cPickle_ErrFormat(PicklingError, "Can't pickle %s: "
+                               "extension code %ld is out of range",
+                               "Ol", args, code);
+                       goto finally;
+               }
+
+               /* Generate an EXT opcode. */
+               if (code <= 0xff) {
+                       c_str[0] = EXT1;
+                       c_str[1] = (char)code;
+                       n = 2;
+               }
+               else if (code <= 0xffff) {
+                       c_str[0] = EXT2;
+                       c_str[1] = (char)(code & 0xff);
+                       c_str[2] = (char)((code >> 8) & 0xff);
+                       n = 3;
+               }
+               else {
+                       c_str[0] = EXT4;
+                       c_str[1] = (char)(code & 0xff);
+                       c_str[2] = (char)((code >> 8) & 0xff);
+                       c_str[3] = (char)((code >> 16) & 0xff);
+                       c_str[4] = (char)((code >> 24) & 0xff);
+                       n = 5;
+               }
+
+               if (self->write_func(self, c_str, n) >= 0)
+                       res = 0;
+               goto finally;   /* and don't memoize */
+       }
+
+  gen_global:
        if (self->write_func(self, &global, 1) < 0)
                goto finally;
 
@@ -5213,7 +5274,11 @@ init_stuff(PyObject *module_dict)
 
        Py_DECREF(copy_reg);
 
-       if (!( empty_tuple = PyTuple_New(0)))
+       if (!(empty_tuple = PyTuple_New(0)))
+               return -1;
+
+       two_tuple = PyTuple_New(2);
+       if (two_tuple == NULL)
                return -1;
 
        /* Ugh */