diff --git a/Lib/test/pickletester.py b/Lib/test/pickletester.py index 6ac4b19da3ea9c..881e672a76ff3f 100644 --- a/Lib/test/pickletester.py +++ b/Lib/test/pickletester.py @@ -57,6 +57,8 @@ # kind of outer loop. protocols = range(pickle.HIGHEST_PROTOCOL + 1) +FAST_NESTING_LIMIT = 50 + # Return True if opcode code appears in the pickle, else False. def opcode_in_pickle(code, pickle): @@ -4552,6 +4554,98 @@ def __reduce__(self): expected = "changed size during iteration" self.assertIn(expected, str(e)) + def fast_save_enter(self, create_data, minprotocol=0): + # gh-146059: Check that fast_save() is called when + # fast_save_enter() is called. + if not hasattr(self, "pickler"): + self.skipTest("need Pickler class") + + data = [create_data(i) for i in range(FAST_NESTING_LIMIT * 2)] + data = {"key": data} + protocols = range(minprotocol, pickle.HIGHEST_PROTOCOL + 1) + for proto in protocols: + with self.subTest(proto=proto): + buf = io.BytesIO() + pickler = self.pickler(buf, protocol=proto) + # Enable fast mode (disables memo, enables cycle detection) + pickler.fast = 1 + pickler.dump(data) + + buf.seek(0) + data2 = self.unpickler(buf).load() + self.assertEqual(data2, data) + + def test_fast_save_enter_tuple(self): + self.fast_save_enter(lambda i: (i,)) + + def test_fast_save_enter_list(self): + self.fast_save_enter(lambda i: [i]) + + def test_fast_save_enter_frozenset(self): + self.fast_save_enter(lambda i: frozenset([i])) + + def test_fast_save_enter_set(self): + self.fast_save_enter(lambda i: set([i])) + + def test_fast_save_enter_frozendict(self): + if self.py_version < (3, 15): + self.skipTest('need frozendict') + self.fast_save_enter(lambda i: frozendict(key=i), minprotocol=2) + + def test_fast_save_enter_dict(self): + self.fast_save_enter(lambda i: {"key": i}) + + def deep_nested_struct(self, seed, create_nested, + minprotocol=0, compare_equal=True, + depth=FAST_NESTING_LIMIT * 2): + # gh-146059: Check that fast_save() is called when + # fast_save_enter() is called. + if not hasattr(self, "pickler"): + self.skipTest("need Pickler class") + + data = seed + for i in range(depth): + data = create_nested(data) + data = {"key": data} + protocols = range(minprotocol, pickle.HIGHEST_PROTOCOL + 1) + for proto in protocols: + with self.subTest(proto=proto): + buf = io.BytesIO() + pickler = self.pickler(buf, protocol=proto) + # Enable fast mode (disables memo, enables cycle detection) + pickler.fast = 1 + pickler.dump(data) + + buf.seek(0) + data2 = self.unpickler(buf).load() + if compare_equal: + self.assertEqual(data2, data) + + def test_deep_nested_struct_tuple(self): + self.deep_nested_struct((1,), lambda data: (data,)) + + def test_deep_nested_struct_list(self): + self.deep_nested_struct([1], lambda data: [data]) + + def test_deep_nested_struct_frozenset(self): + self.deep_nested_struct(frozenset((1,)), + lambda data: frozenset((1, data))) + + def test_deep_nested_struct_set(self): + self.deep_nested_struct({1}, lambda data: {K(data)}, + depth=FAST_NESTING_LIMIT+1, + compare_equal=False) + + def test_deep_nested_struct_frozendict(self): + if self.py_version < (3, 15): + self.skipTest('need frozendict') + self.deep_nested_struct(frozendict(x=1), + lambda data: frozendict(x=data), + minprotocol=2) + + def test_deep_nested_struct_dict(self): + self.deep_nested_struct({'x': 1}, lambda data: {'x': data}) + class BigmemPickleTests: diff --git a/Modules/_pickle.c b/Modules/_pickle.c index 65facaa6db2036..d65e83d4463fa0 100644 --- a/Modules/_pickle.c +++ b/Modules/_pickle.c @@ -3671,16 +3671,13 @@ save_set(PickleState *state, PicklerObject *self, PyObject *obj) } static int -save_frozenset(PickleState *state, PicklerObject *self, PyObject *obj) +save_frozenset_impl(PickleState *state, PicklerObject *self, PyObject *obj) { PyObject *iter; const char mark_op = MARK; const char frozenset_op = FROZENSET; - if (self->fast && !fast_save_enter(self, obj)) - return -1; - if (self->proto < 4) { PyObject *items; PyObject *reduce_value; @@ -3751,6 +3748,19 @@ save_frozenset(PickleState *state, PicklerObject *self, PyObject *obj) return 0; } +static int +save_frozenset(PickleState *state, PicklerObject *self, PyObject *obj) +{ + if (self->fast && !fast_save_enter(self, obj)) { + return -1; + } + int status = save_frozenset_impl(state, self, obj); + if (self->fast && !fast_save_leave(self, obj)) { + return -1; + } + return status; +} + static int fix_imports(PickleState *st, PyObject **module_name, PyObject **global_name) {