diff -r acc453391c5b Lib/pickle.py --- a/Lib/pickle.py Sat Oct 24 17:39:36 2015 +0300 +++ b/Lib/pickle.py Sun Oct 25 21:44:55 2015 +0200 @@ -402,7 +402,13 @@ class Pickler: write(REDUCE) if obj is not None: - self.memoize(obj) + # If the object is already in the memo, this means it is + # recursive. In this case, throw away everything we put on the + # stack, and fetch the object back from the memo. + if id(obj) in self.memo: + write(POP + self.get(self.memo[id(obj)][0])) + else: + self.memoize(obj) # More new special cases (that work with older protocols as # well): when __reduce__ returns a tuple with 4 or 5 items, diff -r acc453391c5b Lib/test/pickletester.py --- a/Lib/test/pickletester.py Sat Oct 24 17:39:36 2015 +0300 +++ b/Lib/test/pickletester.py Sun Oct 25 21:44:55 2015 +0200 @@ -117,6 +117,9 @@ class E(C): def __getinitargs__(self): return () +class H(object): + pass + import __main__ __main__.C = C C.__module__ = "__main__" @@ -124,6 +127,8 @@ C.__module__ = "__main__" D.__module__ = "__main__" __main__.E = E E.__module__ = "__main__" +__main__.H = H +H.__module__ = "__main__" class myint(int): def __init__(self, x): @@ -621,6 +626,7 @@ class AbstractUnpickleTests(unittest.Tes class AbstractPickleTests(unittest.TestCase): # Subclass must define self.dumps, self.loads. + fast = 0 _testdata = AbstractUnpickleTests._testdata @@ -698,6 +704,32 @@ class AbstractPickleTests(unittest.TestC self.assertEqual(x.keys(), [1]) self.assertTrue(x[1] is x) + def test_recursive_set(self): + if self.fast: + self.skipTest("can't pickle cyclic objects in fast mode") + h = H() + y = set({h}) + h.attr = y + for proto in protocols: + s = self.dumps(y, proto) + x = self.loads(s) + self.assertIsInstance(x, set) + self.assertIs(list(x)[0].attr, x) + self.assertEqual(len(x), 1) + + def test_recursive_frozenset(self): + if self.fast: + self.skipTest("can't pickle cyclic objects in fast mode") + h = H() + y = frozenset({h}) + h.attr = y + for proto in protocols: + s = self.dumps(y, proto) + x = self.loads(s) + self.assertIsInstance(x, frozenset) + self.assertIs(list(x)[0].attr, x) + self.assertEqual(len(x), 1) + def test_recursive_inst(self): i = C() i.attr = i @@ -721,6 +753,32 @@ class AbstractPickleTests(unittest.TestC self.assertEqual(x[0].attr.keys(), [1]) self.assertTrue(x[0].attr[1] is x) + def test_recursive_list_subclass(self): + if self.fast: + self.skipTest("can't pickle cyclic objects in fast mode") + h = H() + y = MyList([h]) + h.attr = y + for proto in protocols: + s = self.dumps(y, proto) + x = self.loads(s) + self.assertIsInstance(x, MyList) + self.assertIs(list(x)[0].attr, x) + self.assertEqual(len(x), 1) + + def test_recursive_tuple_subclass(self): + if self.fast: + self.skipTest("can't pickle cyclic objects in fast mode") + h = H() + y = MyTuple([h]) + h.attr = y + for proto in protocols: + s = self.dumps(y, proto) + x = self.loads(s) + self.assertIsInstance(x, MyTuple) + self.assertIs(list(x)[0].attr, x) + self.assertEqual(len(x), 1) + if have_unicode: def test_unicode(self): endcases = [u'', u'<\\u>', u'<\\\u1234>', u'<\n>', diff -r acc453391c5b Lib/test/test_cpickle.py --- a/Lib/test/test_cpickle.py Sat Oct 24 17:39:36 2015 +0300 +++ b/Lib/test/test_cpickle.py Sun Oct 25 21:44:55 2015 +0200 @@ -131,6 +131,7 @@ class FileIOCPicklerListTests(FileIOMixi class cPickleFastPicklerTests(AbstractPickleTests): + fast = 1 def dumps(self, arg, proto=0): f = self.output() diff -r acc453391c5b Modules/cPickle.c --- a/Modules/cPickle.c Sat Oct 24 17:39:36 2015 +0300 +++ b/Modules/cPickle.c Sun Oct 25 21:44:55 2015 +0200 @@ -2533,6 +2533,27 @@ save_reduce(Picklerobject *self, PyObjec /* Memoize. */ /* XXX How can ob be NULL? */ if (ob != NULL) { + /* If the object is already in the memo, this means it is + recursive. In this case, throw away everything we put on the + stack, and fetch the object back from the memo. */ + if (Py_REFCNT(ob) > 1 && !self->fast) { + PyObject *py_ob_id = PyLong_FromVoidPtr(ob); + if (!py_ob_id) + return -1; + if (PyDict_GetItem(self->memo, py_ob_id)) { + const char pop_op = POP; + if (self->write_func(self, &pop_op, 1) < 0 || + get(self, py_ob_id) < 0) { + Py_DECREF(py_ob_id); + return -1; + } + Py_DECREF(py_ob_id); + return 0; + } + Py_DECREF(py_ob_id); + if (PyErr_Occurred()) + return -1; + } if (state && !PyDict_Check(state)) { if (put2(self, ob) < 0) return -1;