diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py index 75e2afa7ae68df0..dafd2e9717b5b0c 100644 --- a/Lib/test/pickletester.py +++ b/Lib/test/pickletester.py @@ -4906,6 +4906,14 @@ def test_clear_pickler_memo(self): self.assertNotEqual(first_pickled, second_pickled) self.assertEqual(first_pickled, third_pickled) + def test_clear_unpickler_memo(self): + value = [b'repeated'] * 2 + data = pickle.dumps(value, protocol=4) + unpickler = self.unpickler_class(io.BytesIO(data * 2)) + self.assertEqual(unpickler.load(), value) + unpickler.memo.clear() + self.assertEqual(unpickler.load(), value) + def test_priming_pickler_memo(self): # Verify that we can set the Pickler's memo attribute. data = ["abcdefg", "abcdefg", 44] diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py index 9d86c1b050304db..475c52cd6c1f477 100644 --- a/Lib/test/test_pickle.py +++ b/Lib/test/test_pickle.py @@ -17,6 +17,7 @@ from test import support from test.support import cpython_only, import_helper, os_helper from test.support.import_helper import ensure_lazy_imports +from test.support.script_helper import assert_python_ok from test.pickletester import AbstractHookTests from test.pickletester import AbstractUnpickleTests @@ -455,6 +456,70 @@ class CPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests, unittest pickler_class = _pickle.Pickler unpickler_class = _pickle.Unpickler + @support.requires_subprocess() + def test_clear_unpickler_memo_memory_error(self): + import_helper.import_module('_testcapi') + code = dedent(''' + import _testcapi + import io + import pickle + from test.support import SuppressCrashReport + + value = [b'repeated'] * 2 + data = pickle.dumps(value, protocol=4) + unpickler = pickle.Unpickler(io.BytesIO(data * 2)) + assert unpickler.load() == value + memo = unpickler.memo.copy() + clear = unpickler.memo.clear + reset = _testcapi.remove_mem_hooks + failed = False + with SuppressCrashReport(): + # Fail the replacement memo allocation, with all callable + # lookups performed before enabling the allocator hook. + _testcapi.set_nomemory(0, 1) + try: + clear() + except MemoryError: + failed = True + finally: + reset() + assert failed + assert unpickler.memo.copy() == memo + clear() + assert unpickler.memo.copy() == {} + assert unpickler.load() == value + ''') + assert_python_ok('-c', code) + + @support.requires_subprocess() + def test_clear_unpickler_memo_reentrant(self): + code = dedent(r''' + import io + import pickle + from test.support import gc_collect, SuppressCrashReport + + snapshots = [] + class Entry: + def __del__(self): + snapshots.append(unpickler.memo.copy()) + + value = [b'repeated'] * 2 + # Create the entry through the C unpickler so clearing its + # sole reference runs the finalizer even without the GIL. + data = b'\x80\x02c__main__\nEntry\n)\x81q\x00.' + unpickler = pickle.Unpickler( + io.BytesIO(data + pickle.dumps(value, protocol=4))) + unpickler.load() + gc_collect() + with SuppressCrashReport(): + unpickler.memo.clear() + gc_collect() + assert snapshots == [{}] + assert unpickler.memo.copy() == {} + assert unpickler.load() == value + ''') + assert_python_ok('-c', code) + def test_issue18339(self): unpickler = self.unpickler_class(io.BytesIO()) with self.assertRaises(TypeError): diff --git a/Misc/NEWS.d/next/Library/2026-10-22-03-55-33.gh-issue-159739.a8bf30.rst b/Misc/NEWS.d/next/Library/2026-10-22-03-55-33.gh-issue-159739.a8bf30.rst new file mode 100644 index 000000000000000..79aea931e201c6a --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-10-22-03-55-33.gh-issue-159739.a8bf30.rst @@ -0,0 +1,3 @@ +Fix clearing a :class:`pickle.Unpickler` memo so that the unpickler can be +reused for independent records. Preserve a valid memo if allocation fails or +releasing an old entry runs a destructor that accesses the memo. diff --git a/Modules/_pickle.c b/Modules/_pickle.c index cba47396a7c2806..ecfbc19749cbe26 100644 --- a/Modules/_pickle.c +++ b/Modules/_pickle.c @@ -1675,16 +1675,16 @@ _Unpickler_NewMemo(Py_ssize_t new_size) return memo; } -/* Free the unpickler's memo, taking care to decref any items left in it. */ +/* Replace the unpickler's memo and release the old entries. */ static void -_Unpickler_MemoCleanup(UnpicklerObject *self) +_Unpickler_MemoCleanup(UnpicklerObject *self, PyObject **replacement) { Py_ssize_t i; PyObject **memo = self->memo; - if (self->memo == NULL) + self->memo = replacement; + if (memo == NULL) return; - self->memo = NULL; i = self->memo_size; while (--i >= 0) { Py_XDECREF(memo[i]); @@ -7465,7 +7465,7 @@ Unpickler_clear(PyObject *op) self->buffer.buf = NULL; } - _Unpickler_MemoCleanup(self); + _Unpickler_MemoCleanup(self, NULL); PyMem_Free(self->marks); self->marks = NULL; PyMem_Free(self->input_line); @@ -7611,10 +7611,13 @@ static PyObject * _pickle_UnpicklerMemoProxy_clear_impl(UnpicklerMemoProxyObject *self) /*[clinic end generated code: output=d20cd43f4ba1fb1f input=b1df7c52e7afd9bd]*/ { - _Unpickler_MemoCleanup(self->unpickler); - self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size); - if (self->unpickler->memo == NULL) + UnpicklerObject *unpickler = self->unpickler; + PyObject **memo = _Unpickler_NewMemo(unpickler->memo_size); + if (memo == NULL) { return NULL; + } + unpickler->memo_len = 0; + _Unpickler_MemoCleanup(unpickler, memo); Py_RETURN_NONE; } @@ -7827,7 +7830,7 @@ Unpickler_set_memo(PyObject *op, PyObject *obj, void *Py_UNUSED(closure)) return -1; } - _Unpickler_MemoCleanup(self); + _Unpickler_MemoCleanup(self, NULL); self->memo_size = new_memo_size; self->memo = new_memo;