gh-146059: Call fast_save_leave() in pickle save_frozenset() (#146173)

Add more pickle tests: test also nested structures.

Author: vstinner

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:
 

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)
 {