cPickle.c, load_build(): Taught cPickle how to pick apart

the optional proto 2 slot state.

pickle.py, load_build():  CAUTION:  Noted that cPickle's
load_build and pickle's load_build really don't do the same
things with the state, and didn't before this patch either.
cPickle never tries to do .update(), and has no backoff if
instance.__dict__ can't be retrieved.  There are no tests
that can tell the difference, and part of what cPickle's
load_build() did looked accidental to me, so I don't know
what the true intent is here.

pickletester.py, test_pickle.py:  Got rid of the hack for
exempting cPickle from running some of the proto 2 tests.

dictobject.c, PyDict_Next():  documented intended use.

5 files changed, 88 insertions, 30 deletions

diff --git a/Lib/pickle.py b/Lib/pickle.py
index 74748f8..c62bddc 100644
--- a/Lib/pickle.py
+++ b/Lib/pickle.py
@@ -1249,6 +1249,10 @@ class Unpickler:
                 # the instance variables.  This is a semantic
                 # difference when unpickling in restricted
                 # vs. unrestricted modes.
+                # Note, however, that cPickle has never tried to do the
+                # .update() business, and always uses
+                #     PyObject_SetItem(inst.__dict__, key, value) in a
+                # loop over state.items().
                 for k, v in state.items():
                     setattr(inst, k, v)
         if slotstate:
diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py
index 57e051c..d541194 100644
--- a/Lib/test/pickletester.py
+++ b/Lib/test/pickletester.py
@@ -728,12 +728,6 @@ class AbstractPickleTests(unittest.TestCase):
             self.assertEqual(y.abc, 666)
             self.assertEqual(x.__dict__, y.__dict__)
 
-# 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
@@ -745,6 +739,7 @@ class TempAbstractPickleTests(unittest.TestCase):
         self.assertEqual(x.foo, y.foo)
         self.assertEqual(x.bar, y.bar)
 
+
 class MyInt(int):
     sample = 1
 
diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py
index ac2a596..9cfb9b7 100644
--- a/Lib/test/test_pickle.py
+++ b/Lib/test/test_pickle.py
@@ -5,11 +5,10 @@ from cStringIO import StringIO
 from test import test_support
 
 from test.pickletester import AbstractPickleTests
-from test.pickletester import TempAbstractPickleTests as XXXTemp
 from test.pickletester import AbstractPickleModuleTests
 from test.pickletester import AbstractPersistentPicklerTests
 
-class PickleTests(AbstractPickleTests, AbstractPickleModuleTests, XXXTemp):
+class PickleTests(AbstractPickleTests, AbstractPickleModuleTests):
 
     def dumps(self, arg, proto=0, fast=0):
         # Ignore fast
diff --git a/Modules/cPickle.c b/Modules/cPickle.c
index f09e502..3f3d82f 100644
--- a/Modules/cPickle.c
+++ b/Modules/cPickle.c
@@ -4309,43 +4309,93 @@ load_setitems(Unpicklerobject *self)
 static int
 load_build(Unpicklerobject *self)
 {
-	PyObject *value = 0, *inst = 0, *instdict = 0, *d_key = 0, *d_value = 0,
-		*junk = 0, *__setstate__ = 0;
-	int i, r = 0;
+	PyObject *state, *inst, *slotstate;
+	PyObject *__setstate__;
+	PyObject *d_key, *d_value;
+	int i;
+	int res = -1;
 
-	if (self->stack->length < 2) return stackUnderflow();
-	PDATA_POP(self->stack, value);
-	if (! value) return -1;
-	inst=self->stack->data[self->stack->length-1];
+	/* Stack is ... instance, state.  We want to leave instance at
+	 * the stack top, possibly mutated via instance.__setstate__(state).
+	 */
+	if (self->stack->length < 2)
+		return stackUnderflow();
+	PDATA_POP(self->stack, state);
+	if (state == NULL)
+		return -1;
+	inst = self->stack->data[self->stack->length - 1];
+
+	__setstate__ = PyObject_GetAttr(inst, __setstate___str);
+	if (__setstate__ != NULL) {
+		PyObject *junk = NULL;
 
-	if ((__setstate__ = PyObject_GetAttr(inst, __setstate___str))) {
-		ARG_TUP(self, value);
+		/* The explicit __setstate__ is responsible for everything. */
+		ARG_TUP(self, state);
 		if (self->arg) {
 			junk = PyObject_Call(__setstate__, self->arg, NULL);
 			FREE_ARG_TUP(self);
 		}
 		Py_DECREF(__setstate__);
-		if (! junk) return -1;
+		if (junk == NULL)
+			return -1;
 		Py_DECREF(junk);
 		return 0;
 	}
-
 	PyErr_Clear();
-	if ((instdict = PyObject_GetAttr(inst, __dict___str))) {
+
+	/* A default __setstate__.  First see whether state embeds a
+	 * slot state dict too (a proto 2 addition).
+	 */
+	if (PyTuple_Check(state) && PyTuple_Size(state) == 2) {
+		PyObject *temp = state;
+		state = PyTuple_GET_ITEM(temp, 0);
+		slotstate = PyTuple_GET_ITEM(temp, 1);
+		Py_INCREF(state);
+		Py_INCREF(slotstate);
+		Py_DECREF(temp);
+	}
+	else
+		slotstate = NULL;
+
+	/* Set inst.__dict__ from the state dict (if any). */
+	if (state != Py_None) {
+		PyObject *dict;
+		if (! PyDict_Check(state)) {
+			PyErr_SetString(UnpicklingError, "state is not a "
+					"dictionary");
+			goto finally;
+		}
+		dict = PyObject_GetAttr(inst, __dict___str);
+		if (dict == NULL)
+			goto finally;
+
 		i = 0;
-		while (PyDict_Next(value, &i, &d_key, &d_value)) {
-			if (PyObject_SetItem(instdict, d_key, d_value) < 0) {
-				r=-1;
-				break;
-			}
+		while (PyDict_Next(state, &i, &d_key, &d_value)) {
+			if (PyObject_SetItem(dict, d_key, d_value) < 0)
+				goto finally;
 		}
-		Py_DECREF(instdict);
+		Py_DECREF(dict);
 	}
-	else r=-1;
 
-	Py_XDECREF(value);
+	/* Also set instance attributes from the slotstate dict (if any). */
+	if (slotstate != NULL) {
+		if (! PyDict_Check(slotstate)) {
+			PyErr_SetString(UnpicklingError, "slot state is not "
+					"a dictionary");
+			goto finally;
+		}
+		i = 0;
+		while (PyDict_Next(slotstate, &i, &d_key, &d_value)) {
+			if (PyObject_SetAttr(inst, d_key, d_value) < 0)
+				goto finally;
+		}
+	}
+	res = 0;
 
-	return r;
+  finally:
+	Py_DECREF(state);
+	Py_XDECREF(slotstate);
+	return res;
 }
 
 
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
index de7a18e..9ae7185 100644
--- a/Objects/dictobject.c
+++ b/Objects/dictobject.c
@@ -642,7 +642,17 @@ PyDict_Clear(PyObject *op)
 		PyMem_DEL(table);
 }
 
-/* CAUTION:  In general, it isn't safe to use PyDict_Next in a loop that
+/*
+ * Iterate over a dict.  Use like so:
+ *
+ *     int i;
+ *     PyObject *key, *value;
+ *     i = 0;   # important!  i should not otherwise be changed by you
+ *     while (PyDict_Next(yourdict, &i, &key, &value) {
+ *              Refer to borrowed references in key and value.
+ *     }
+ *
+ * CAUTION:  In general, it isn't safe to use PyDict_Next in a loop that
  * mutates the dict.  One exception:  it is safe if the loop merely changes
  * the values associated with the keys (but doesn't insert new keys or
  * delete keys), via PyDict_SetItem().