add repr and pickle tests for paramspec

python · JelleZijlstra · Apr 10, 2021 · Jan 2, 2021 · Jan 2, 2021 · Jan 3, 2021
commit a95f43597bcd9a3e0df0853f7aa490e9881c037e
diff --git a/typing_extensions/src_py3/test_typing_extensions.py b/typing_extensions/src_py3/test_typing_extensions.py
@@ -1875,14 +1875,27 @@ def test_cannot_subscript(self):
         with self.assertRaises(TypeError):
             TypeAlias[int]
 
-
+P = ParamSpec('P')
+P_co = ParamSpec('P_co', covariant=True)
+P_contra = ParamSpec('P_contra', contravariant=True)
 class ParamSpecTests(BaseTestCase):
 
     def test_basic_plain(self):
-        P = ParamSpec('P')
+        global P
         self.assertEqual(P, P)
         self.assertIsInstance(P, ParamSpec)
 
+    def test_repr(self):
+        global P, P_co, P_contra
+        P_2 = ParamSpec('P_2')
+        self.assertEqual(repr(P), '~P')
+        self.assertEqual(repr(P_2), '~P_2')
+
+        # Note: PEP 612 doesn't require these to be repr-ed correctly, but
+        # just follow CPython.
+        self.assertEqual(repr(P_co), '+P_co')
+        self.assertEqual(repr(P_contra), '-P_contra')
+
     def test_valid_uses(self):
         P = ParamSpec('P')
         T = TypeVar('T')
@@ -1904,8 +1917,20 @@ def test_valid_uses(self):
         P.args
         P.kwargs
 
-        # Note: ParamSpec doesn't work for pre-3.10 user-defined Generics due
-        # to type checks inside Generic.
+    # Note: ParamSpec doesn't work for pre-3.10 user-defined Generics due
+    # to type checks inside Generic.
+
+    def test_pickle(self):
+        global P, P_co, P_contra
+        for proto in range(pickle.HIGHEST_PROTOCOL):
+            with self.subTest(f"Pickle protocol {proto}"):
+                for paramspec in (P, P_co, P_contra):
+                    z = pickle.loads(pickle.dumps(paramspec, proto))
+                    self.assertEqual(z.__name__, paramspec.__name__)
+                    self.assertEqual(z.__covariant__, paramspec.__covariant__)
+                    self.assertEqual(z.__contravariant__, paramspec.__contravariant__)
+                    self.assertEqual(z.__bound__, paramspec.__bound__)
+
 
 class ConcatenateTests(BaseTestCase):
     def test_basics(self):

diff --git a/typing_extensions/src_py3/typing_extensions.py b/typing_extensions/src_py3/typing_extensions.py
@@ -2233,6 +2233,14 @@ def __init__(self, name, *, bound=None, covariant=False, contravariant=False):
             else:
                 self.__bound__ = None
 
+            # for pickling:
+            try:
+                def_mod = sys._getframe(1).f_globals.get('__name__', '__main__')
+            except (AttributeError, ValueError):
+                def_mod = None
+            if def_mod != 'typing_extensions':
+                self.__module__ = def_mod
+
         def __repr__(self):
             if self.__covariant__:
                 prefix = '+'
@@ -2242,6 +2250,8 @@ def __repr__(self):
                 prefix = '~'
             return prefix + self.__name__
 
+        def __reduce__(self):
+            return self.__name__
 
 # Inherits from list as a workaround for Callable checks in Python < 3.9.2.
 class _ConcatenateGenericAlias(list):