Add more tests: test also nested structures

Author: vstinner

Lib/test/pickletester.py

@@ -4554,36 +4554,103 @@ 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
+    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")
 
-        limit = FAST_NESTING_LIMIT * 2
+        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:
-            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)
+            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]))
 
-                    buf.seek(0)
-                    data2 = self.unpickler(buf).load()
+    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):
+        def create_nested(data):
+            obj = Object()
+            obj.value = data
+            return {obj}
+
+        self.deep_nested_struct({1}, create_nested,
+                                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: