Add cPickle support for PROTO. Duplicated PROTO/LONG1/LONG4 code in
the hitherto unknown (to me) noload() cPickle function, which is (a)
something we don't test at all, and (b) pickle.py doesn't have.
diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py
index 2a1ca17..5ef0cf2 100644
--- a/Lib/test/pickletester.py
+++ b/Lib/test/pickletester.py
@@ -1,4 +1,6 @@
import unittest
+import pickle
+
from test.test_support import TestFailed, have_unicode, TESTFN
# Tests that try a number of pickle protocols should have a
@@ -296,6 +298,25 @@
# Tests for protocol 2
+ def test_proto(self):
+ build_none = pickle.NONE + pickle.STOP
+ for proto in protocols:
+ expected = build_none
+ if proto >= 2:
+ expected = pickle.PROTO + chr(proto) + expected
+ p = self.dumps(None, proto)
+ self.assertEqual(p, expected)
+
+ oob = protocols[-1] + 1 # a future protocol
+ badpickle = pickle.PROTO + chr(oob) + build_none
+ try:
+ self.loads(badpickle)
+ except ValueError, detail:
+ self.failUnless(str(detail).startswith(
+ "unsupported pickle protocol"))
+ else:
+ self.fail("expected bad protocol number to raise ValueError")
+
def test_long1(self):
x = 12345678910111213141516178920L
s = self.dumps(x, 2)
@@ -314,14 +335,14 @@
c = (1, 2)
d = (1, 2, 3)
e = (1, 2, 3, 4)
- for proto in 0, 1, 2:
+ for proto in protocols:
for x in a, b, c, d, e:
s = self.dumps(x, proto)
y = self.loads(s)
self.assertEqual(x, y, (proto, x, s, y))
def test_singletons(self):
- for proto in 0, 1, 2:
+ for proto in protocols:
for x in None, False, True:
s = self.dumps(x, proto)
y = self.loads(s)
diff --git a/Modules/cPickle.c b/Modules/cPickle.c
index b59f573..43a8d33 100644
--- a/Modules/cPickle.c
+++ b/Modules/cPickle.c
@@ -2213,13 +2213,22 @@
{
static char stop = STOP;
+ if (self->proto >= 2) {
+ char bytes[2];
+
+ bytes[0] = PROTO;
+ bytes[1] = CURRENT_PROTOCOL_NUMBER;
+ if (self->write_func(self, bytes, 2) < 0)
+ return -1;
+ }
+
if (save(self, args, 0) < 0)
return -1;
- if ((*self->write_func)(self, &stop, 1) < 0)
+ if (self->write_func(self, &stop, 1) < 0)
return -1;
- if ((*self->write_func)(self, NULL, 0) < 0)
+ if (self->write_func(self, NULL, 0) < 0)
return -1;
return 0;
@@ -3870,6 +3879,31 @@
return 0;
}
+/* Just raises an error if we don't know the protocol specified. PROTO
+ * is the first opcode for protocols >= 2.
+ */
+static int
+load_proto(Unpicklerobject *self)
+{
+ int i;
+ char *protobyte;
+
+ i = self->read_func(self, &protobyte, 1);
+ if (i < 0)
+ return -1;
+
+ i = calc_binint(protobyte, 1);
+ /* No point checking for < 0, since calc_binint returns an unsigned
+ * int when chewing on 1 byte.
+ */
+ assert(i >= 0);
+ if (i <= CURRENT_PROTOCOL_NUMBER)
+ return 0;
+
+ PyErr_Format(PyExc_ValueError, "unsupported pickle protocol: %d", i);
+ return -1;
+}
+
static PyObject *
load(Unpicklerobject *self)
{
@@ -4099,6 +4133,11 @@
break;
continue;
+ case PROTO:
+ if (load_proto(self) < 0)
+ break;
+ continue;
+
case '\0':
/* end of file */
PyErr_SetNone(PyExc_EOFError);
@@ -4227,6 +4266,16 @@
break;
continue;
+ case LONG1:
+ if (load_counted_long(self, 1) < 0)
+ break;
+ continue;
+
+ case LONG4:
+ if (load_counted_long(self, 4) < 0)
+ break;
+ continue;
+
case FLOAT:
if (load_float(self) < 0)
break;
@@ -4402,6 +4451,11 @@
break;
continue;
+ case PROTO:
+ if (load_proto(self) < 0)
+ break;
+ continue;
+
default:
cPickle_ErrFormat(UnpicklingError,
"invalid load key, '%s'.",