SF patch #1062279: deque pickling problems

(Contributed by Dima Dorfman.)

* Support pickling of dictionaries in instances of deque subclasses.
* Support pickling of recursive deques.

2 files changed, 48 insertions, 14 deletions

diff --git a/Lib/test/test_deque.py b/Lib/test/test_deque.py
index 19b0afc..0a6c1f9 100644
--- a/Lib/test/test_deque.py
+++ b/Lib/test/test_deque.py
@@ -297,10 +297,19 @@ class TestBasic(unittest.TestCase):
 
     def test_pickle(self):
         d = deque(xrange(200))
-        s = pickle.dumps(d)
-        e = pickle.loads(s)
-        self.assertNotEqual(id(d), id(e))
-        self.assertEqual(list(d), list(e))
+        for i in (0, 1, 2):
+            s = pickle.dumps(d, i)
+            e = pickle.loads(s)
+            self.assertNotEqual(id(d), id(e))
+            self.assertEqual(list(d), list(e))
+
+    def test_pickle_recursive(self):
+        d = deque('abc')
+        d.append(d)
+        for i in (0, 1, 2):
+            e = pickle.loads(pickle.dumps(d, i))
+            self.assertNotEqual(id(d), id(e))
+            self.assertEqual(id(e), id(e[-1]))
 
     def test_deepcopy(self):
         mut = [10]
@@ -430,6 +439,10 @@ class TestVariousIteratorArgs(unittest.TestCase):
 class Deque(deque):
     pass
 
+class DequeWithBadIter(deque):
+    def __iter__(self):
+        raise TypeError
+
 class TestSubclass(unittest.TestCase):
 
     def test_basics(self):
@@ -472,6 +485,25 @@ class TestSubclass(unittest.TestCase):
         self.assertEqual(type(d), type(e))
         self.assertEqual(list(d), list(e))
 
+    def test_pickle(self):
+        d = Deque('abc')
+        d.append(d)
+
+        e = pickle.loads(pickle.dumps(d))
+        self.assertNotEqual(id(d), id(e))
+        self.assertEqual(type(d), type(e))
+        dd = d.pop()
+        ee = e.pop()
+        self.assertEqual(id(e), id(ee))
+        self.assertEqual(d, e)
+
+        d.x = d
+        e = pickle.loads(pickle.dumps(d))
+        self.assertEqual(id(e), id(e.x))
+
+        d = DequeWithBadIter('abc')
+        self.assertRaises(TypeError, pickle.dumps, d)
+
     def test_weakref(self):
         d = deque('gallahad')
         p = proxy(d)
diff --git a/Modules/collectionsmodule.c b/Modules/collectionsmodule.c
index 10b08f9..ed10999 100644
--- a/Modules/collectionsmodule.c
+++ b/Modules/collectionsmodule.c
@@ -550,19 +550,21 @@ PyDoc_STRVAR(copy_doc, "Return a shallow copy of a deque.");
 static PyObject *
 deque_reduce(dequeobject *deque)
 {
-	PyObject *seq, *args, *result;
+	PyObject *dict, *result, *it;
 
-	seq = PySequence_Tuple((PyObject *)deque);
-	if (seq == NULL)
-		return NULL;
-	args = PyTuple_Pack(1, seq);
-	if (args == NULL) {
-		Py_DECREF(seq);
+	dict = PyObject_GetAttrString((PyObject *)deque, "__dict__");
+	if (dict == NULL) {
+		PyErr_Clear();
+		dict = Py_None;
+		Py_INCREF(dict);
+	}
+	it = PyObject_GetIter((PyObject *)deque);
+	if (it == NULL) {
+		Py_DECREF(dict);
 		return NULL;
 	}
-	result = PyTuple_Pack(2, deque->ob_type, args);
-	Py_DECREF(seq);
-	Py_DECREF(args);
+	result = Py_BuildValue("O()ON", deque->ob_type, dict, it);
+	Py_DECREF(dict);
 	return result;
 }