From 4190fb84293b4c54127e1b7be2d9bdf06edf2910 Mon Sep 17 00:00:00 2001
From: Tim Peters <tim.peters@gmail.com>
Date: Sun, 2 Feb 2003 16:09:05 +0000
Subject: 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.

---
 Lib/test/pickletester.py | 25 +++++++++++++++++++--
 Modules/cPickle.c        | 58 ++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 79 insertions(+), 4 deletions(-)

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 @@ class AbstractPickleTests(unittest.TestCase):
 
     # 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 @@ class AbstractPickleTests(unittest.TestCase):
         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 @@ dump(Picklerobject *self, PyObject *args)
 {
 	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 @@ load_reduce(Unpicklerobject *self)
 	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 @@ load(Unpicklerobject *self)
 				break;
 			continue;
 
+		case PROTO:
+			if (load_proto(self) < 0)
+				break;
+			continue;
+
 		case '\0':
 			/* end of file */
 			PyErr_SetNone(PyExc_EOFError);
@@ -4227,6 +4266,16 @@ noload(Unpicklerobject *self)
 				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 @@ noload(Unpicklerobject *self)
 				break;
 			continue;
 
+		case PROTO:
+			if (load_proto(self) < 0)
+				break;
+			continue;
+
 		default:
 			cPickle_ErrFormat(UnpicklingError,
 					  "invalid load key, '%s'.",
-- 
cgit v0.12