gh-105433: Add pickle tests for PEP695 (#105443)

sobolevn · web-flow · commit 1af8251d9ec2 · 2023-06-15T17:58:40.000-07:00
diff --git a/Lib/test/test_type_aliases.py b/Lib/test/test_type_aliases.py
@@ -228,8 +228,69 @@ def test_module(self):
         self.assertEqual(mod_generics_cache.OldStyle.__module__,
                          mod_generics_cache.__name__)
 
+
+# All these type aliases are used for pickling tests:
+T = TypeVar('T')
+type SimpleAlias = int
+type RecursiveAlias = dict[str, RecursiveAlias]
+type GenericAlias[X] = list[X]
+type GenericAliasMultipleTypes[X, Y] = dict[X, Y]
+type RecursiveGenericAlias[X] = dict[str, RecursiveAlias[X]]
+type BoundGenericAlias[X: int] = set[X]
+type ConstrainedGenericAlias[LongName: (str, bytes)] = list[LongName]
+type AllTypesAlias[A, *B, **C] = Callable[C, A] | tuple[*B]
+
+
+class TypeAliasPickleTest(unittest.TestCase):
     def test_pickling(self):
-        pickled = pickle.dumps(mod_generics_cache.Alias)
-        self.assertIs(pickle.loads(pickled), mod_generics_cache.Alias)
-        pickled = pickle.dumps(mod_generics_cache.OldStyle)
-        self.assertIs(pickle.loads(pickled), mod_generics_cache.OldStyle)
+        things_to_test = [
+            SimpleAlias,
+            RecursiveAlias,
+
+            GenericAlias,
+            GenericAlias[T],
+            GenericAlias[int],
+
+            GenericAliasMultipleTypes,
+            GenericAliasMultipleTypes[str, T],
+            GenericAliasMultipleTypes[T, str],
+            GenericAliasMultipleTypes[int, str],
+
+            RecursiveGenericAlias,
+            RecursiveGenericAlias[T],
+            RecursiveGenericAlias[int],
+
+            BoundGenericAlias,
+            BoundGenericAlias[int],
+            BoundGenericAlias[T],
+
+            ConstrainedGenericAlias,
+            ConstrainedGenericAlias[str],
+            ConstrainedGenericAlias[T],
+
+            AllTypesAlias,
+            AllTypesAlias[int, str, T, [T, object]],
+
+            # Other modules:
+            mod_generics_cache.Alias,
+            mod_generics_cache.OldStyle,
+        ]
+        for thing in things_to_test:
+            for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+                with self.subTest(thing=thing, proto=proto):
+                    pickled = pickle.dumps(thing, protocol=proto)
+                    self.assertEqual(pickle.loads(pickled), thing)
+
+    type ClassLevel = str
+
+    def test_pickling_local(self):
+        type A = int
+        things_to_test = [
+            self.ClassLevel,
+            A,
+        ]
+        for thing in things_to_test:
+            for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+                with self.subTest(thing=thing, proto=proto):
+                    with self.assertRaises(pickle.PickleError):
+                        pickle.dumps(thing, protocol=proto)
diff --git a/Lib/test/test_type_params.py b/Lib/test/test_type_params.py
@@ -2,6 +2,7 @@
 import textwrap
 import types
 import unittest
+import pickle
 from test.support import requires_working_socket, check_syntax_error, run_code
 
 from typing import Generic, Sequence, TypeVar, TypeVarTuple, ParamSpec, get_args
@@ -855,3 +856,68 @@ def func[A]():
 
         ns = run_code(code)
         self.assertEqual(ns["func"].__type_params__, ())
+
+
+
+# All these type aliases are used for pickling tests:
+T = TypeVar('T')
+def func1[X](x: X) -> X: ...
+def func2[X, Y](x: X | Y) -> X | Y: ...
+def func3[X, *Y, **Z](x: X, y: tuple[*Y], z: Z) -> X: ...
+def func4[X: int, Y: (bytes, str)](x: X, y: Y) -> X | Y: ...
+
+class Class1[X]: ...
+class Class2[X, Y]: ...
+class Class3[X, *Y, **Z]: ...
+class Class4[X: int, Y: (bytes, str)]: ...
+
+
+class TypeParamsPickleTest(unittest.TestCase):
+    def test_pickling_functions(self):
+        things_to_test = [
+            func1,
+            func2,
+            func3,
+            func4,
+        ]
+        for thing in things_to_test:
+            for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+                with self.subTest(thing=thing, proto=proto):
+                    pickled = pickle.dumps(thing, protocol=proto)
+                    self.assertEqual(pickle.loads(pickled), thing)
+
+    def test_pickling_classes(self):
+        things_to_test = [
+            Class1,
+            Class1[int],
+            Class1[T],
+
+            Class2,
+            Class2[int, T],
+            Class2[T, int],
+            Class2[int, str],
+
+            Class3,
+            Class3[int, T, str, bytes, [float, object, T]],
+
+            Class4,
+            Class4[int, bytes],
+            Class4[T, bytes],
+            Class4[int, T],
+            Class4[T, T],
+        ]
+        for thing in things_to_test:
+            for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+                with self.subTest(thing=thing, proto=proto):
+                    pickled = pickle.dumps(thing, protocol=proto)
+                    self.assertEqual(pickle.loads(pickled), thing)
+
+        for klass in things_to_test:
+            real_class = getattr(klass, '__origin__', klass)
+            thing = klass()
+            for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+                with self.subTest(thing=thing, proto=proto):
+                    pickled = pickle.dumps(thing, protocol=proto)
+                    # These instances are not equal,
+                    # but class check is good enough:
+                    self.assertIsInstance(pickle.loads(pickled), real_class)
