diff -r d4b68ac6c4ea Lib/pickle.py --- a/Lib/pickle.py Tue Nov 03 22:43:31 2015 -0800 +++ b/Lib/pickle.py Thu Nov 05 09:55:25 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 d4b68ac6c4ea Lib/test/pickletester.py --- a/Lib/test/pickletester.py Tue Nov 03 22:43:31 2015 -0800 +++ b/Lib/test/pickletester.py Thu Nov 05 09:55:25 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 @@ -670,33 +676,91 @@ class AbstractPickleTests(unittest.TestC got = filelike.getvalue() self.assertEqual(expected, got) + def recursive_collection_test(self, cls): + h = H() + y = cls([h]) + h.attr = y + for proto in protocols: + s = self.dumps(y, proto) + x = self.loads(s) + self.assertIs(type(x), type(y)) + self.assertEqual(len(x), 1) + self.assertIs(type(list(x)[0]), H) + self.assertIs(list(x)[0].attr, x) + def test_recursive_list(self): + self.recursive_collection_test(list) l = [] l.append(l) for proto in protocols: s = self.dumps(l, proto) x = self.loads(s) + self.assertIs(type(x), list) self.assertEqual(len(x), 1) - self.assertTrue(x is x[0]) + self.assertIs(x[0], x) def test_recursive_tuple(self): + self.recursive_collection_test(tuple) t = ([],) t[0].append(t) for proto in protocols: s = self.dumps(t, proto) x = self.loads(s) + self.assertIs(type(x), tuple) self.assertEqual(len(x), 1) + self.assertIs(type(x[0]), list) self.assertEqual(len(x[0]), 1) - self.assertTrue(x is x[0][0]) + self.assertIs(x[0][0], x) def test_recursive_dict(self): + self.recursive_collection_test(dict.fromkeys) d = {} d[1] = d for proto in protocols: s = self.dumps(d, proto) x = self.loads(s) + self.assertIs(type(x), dict) self.assertEqual(x.keys(), [1]) - self.assertTrue(x[1] is x) + self.assertIs(x[1], x) + + def test_recursive_set(self): + if self.fast: + self.skipTest("can't pickle cyclic objects in fast mode") + self.recursive_collection_test(set) + + def test_recursive_frozenset(self): + if self.fast: + self.skipTest("can't pickle cyclic objects in fast mode") + self.recursive_collection_test(frozenset) + + def test_recursive_list_subclass(self): + if self.fast: + self.skipTest("can't pickle cyclic objects in fast mode") + self.recursive_collection_test(MyList) + y = MyList() + y.append(y) + s = self.dumps(y, 2) + x = self.loads(s) + self.assertIs(type(x), MyList) + self.assertEqual(len(x), 1) + self.assertIs(x[0], x) + + def test_recursive_tuple_subclass(self): + if self.fast: + self.skipTest("can't pickle cyclic objects in fast mode") + self.recursive_collection_test(MyTuple) + + def test_recursive_dict_subclass(self): + if self.fast: + self.skipTest("can't pickle cyclic objects in fast mode") + self.recursive_collection_test(MyDict.fromkeys) + d = MyDict() + d[1] = d + s = self.dumps(d, 2) + x = self.loads(s) + self.assertIs(type(x), MyDict) + self.assertEqual(x.keys(), [1]) + self.assertIs(x[1], x) def test_recursive_inst(self): i = C() diff -r d4b68ac6c4ea Lib/test/test_cpickle.py --- a/Lib/test/test_cpickle.py Tue Nov 03 22:43:31 2015 -0800 +++ b/Lib/test/test_cpickle.py Thu Nov 05 09:55:25 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 d4b68ac6c4ea Modules/cPickle.c --- a/Modules/cPickle.c Tue Nov 03 22:43:31 2015 -0800 +++ b/Modules/cPickle.c Thu Nov 05 09:55:25 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;