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

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,36 @@ def __reduce__(self):
                     expected = "changed size during iteration"
                     self.assertIn(expected, str(e))
 
+    def test_fast_save_enter_leave(self):
+        # gh-146059: Check that fast_save_leave() is called when
+        # fast_save_enter() is called.
+        if not hasattr(self, "pickler"):
+            self.skipTest("need Pickler class")
+
+        limit = FAST_NESTING_LIMIT * 2
+        for proto in protocols:
+            for tested_type in (frozenset, list, dict):
+                with self.subTest(proto=proto, tested_type=tested_type):
+                    buf = io.BytesIO()
+                    pickler = self.pickler(buf, protocol=proto)
+                    # Enable fast mode (disables memo, enables cycle detection)
+                    pickler.fast = 1
+
+                    if tested_type == frozenset:
+                        data = [frozenset([i]) for i in range(limit)]
+                    elif tested_type == list:
+                        data = [[i] for i in range(limit)]
+                    elif tested_type == dict:
+                        data = [{"key": 123} for i in range(limit)]
+                    else:
+                        self.fail("unknown tested_type")
+                    data = {"key": data}
+                    pickler.dump(data)
+
+                    buf.seek(0)
+                    data2 = self.unpickler(buf).load()
+                    self.assertEqual(data2, 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)
 {