Fix and expand the support of typing.TypeVar for Python >=3.5 (#350)

Non-importable, dynamic `TypeVar` were pickled as importable objects in Python 3.7 -- this PR fixes it.
Also, this PR adds support to pickle `TypeVar` objects on Python 3.5 and Python 3.6.

Co-authored-by: Olivier Grisel <[email protected]>
Co-authored-by: Pierre Glaser <[email protected]>

Author: 3 people

Diff for: CHANGES.md

@@ -7,6 +7,10 @@
   (follow up on #276)
   ([issue #347](https://github.com/cloudpipe/cloudpickle/issues/347))
 
+- Fix a bug affecting the pickling of dynamic `TypeVar` instances on Python 3.7+,
+  and expand the support for pickling `TypeVar` instances (dynamic or non-dynamic)
+  to Python 3.5-3.6 ([PR #350](https://github.com/cloudpipe/cloudpickle/pull/350))
+
 1.3.0
 =====
 

Diff for: cloudpickle/cloudpickle.py

@@ -58,6 +58,7 @@
 import weakref
 import uuid
 import threading
+import typing
 from enum import Enum
 
 from pickle import _Pickler as Pickler
@@ -116,7 +117,7 @@ def _whichmodule(obj, name):
     - Errors arising during module introspection are ignored, as those errors
       are considered unwanted side effects.
     """
-    module_name = getattr(obj, '__module__', None)
+    module_name = _get_module_attr(obj)
     if module_name is not None:
         return module_name
     # Protect the iteration by using a copy of sys.modules against dynamic
@@ -139,22 +140,46 @@ def _whichmodule(obj, name):
     return None
 
 
-def _is_global(obj, name=None):
+if sys.version_info[:2] < (3, 7):  # pragma: no branch
+    # Workaround bug in old Python versions: prior to Python 3.7, T.__module__
+    # would always be set to "typing" even when the TypeVar T would be defined
+    # in a different module.
+    #
+    # For such older Python versions, we ignore the __module__ attribute of
+    # TypeVar instances and instead exhaustively lookup those instances in all
+    # currently imported modules via the _whichmodule function.
+    def _get_module_attr(obj):
+        if isinstance(obj, typing.TypeVar):
+            return None
+        return getattr(obj, '__module__', None)
+else:
+    def _get_module_attr(obj):
+        return getattr(obj, '__module__', None)
+
+
+def _is_importable_by_name(obj, name=None):
     """Determine if obj can be pickled as attribute of a file-backed module"""
+    return _lookup_module_and_qualname(obj, name=name) is not None
+
+
+def _lookup_module_and_qualname(obj, name=None):
     if name is None:
         name = getattr(obj, '__qualname__', None)
-    if name is None:
+    if name is None:  # pragma: no cover
+        # This used to be needed for Python 2.7 support but is probably not
+        # needed anymore. However we keep the __name__ introspection in case
+        # users of cloudpickle rely on this old behavior for unknown reasons.
         name = getattr(obj, '__name__', None)
 
     module_name = _whichmodule(obj, name)
 
     if module_name is None:
         # In this case, obj.__module__ is None AND obj was not found in any
         # imported module. obj is thus treated as dynamic.
-        return False
+        return None
 
     if module_name == "__main__":
-        return False
+        return None
 
     module = sys.modules.get(module_name, None)
     if module is None:
@@ -163,18 +188,20 @@ def _is_global(obj, name=None):
         # types.ModuleType. The other possibility is that module was removed
         # from sys.modules after obj was created/imported. But this case is not
         # supported, as the standard pickle does not support it either.
-        return False
+        return None
 
     # module has been added to sys.modules, but it can still be dynamic.
     if _is_dynamic(module):
-        return False
+        return None
 
     try:
         obj2, parent = _getattribute(module, name)
     except AttributeError:
         # obj was not found inside the module it points to
-        return False
-    return obj2 is obj
+        return None
+    if obj2 is not obj:
+        return None
+    return module, name
 
 
 def _extract_code_globals(co):
@@ -418,6 +445,11 @@ def dump(self, obj):
             else:
                 raise
 
+    def save_typevar(self, obj):
+        self.save_reduce(*_typevar_reduce(obj))
+
+    dispatch[typing.TypeVar] = save_typevar
+
     def save_memoryview(self, obj):
         self.save(obj.tobytes())
 
@@ -467,7 +499,7 @@ def save_function(self, obj, name=None):
         Determines what kind of function obj is (e.g. lambda, defined at
         interactive prompt, etc) and handles the pickling appropriately.
         """
-        if _is_global(obj, name=name):
+        if _is_importable_by_name(obj, name=name):
             return Pickler.save_global(self, obj, name=name)
         elif PYPY and isinstance(obj.__code__, builtin_code_type):
             return self.save_pypy_builtin_func(obj)
@@ -770,7 +802,7 @@ def save_global(self, obj, name=None, pack=struct.pack):
                 _builtin_type, (_BUILTIN_TYPE_NAMES[obj],), obj=obj)
         elif name is not None:
             Pickler.save_global(self, obj, name=name)
-        elif not _is_global(obj, name=name):
+        elif not _is_importable_by_name(obj, name=name):
             self.save_dynamic_class(obj)
         else:
             Pickler.save_global(self, obj, name=name)
@@ -1214,3 +1246,25 @@ def _is_dynamic(module):
     else:
         pkgpath = None
     return _find_spec(module.__name__, pkgpath, module) is None
+
+
+def _make_typevar(name, bound, constraints, covariant, contravariant):
+    return typing.TypeVar(
+        name, *constraints, bound=bound,
+        covariant=covariant, contravariant=contravariant
+    )
+
+
+def _decompose_typevar(obj):
+    return (
+        obj.__name__, obj.__bound__, obj.__constraints__,
+        obj.__covariant__, obj.__contravariant__,
+    )
+
+
+def _typevar_reduce(obj):
+    # TypeVar instances have no __qualname__ hence we pass the name explicitly.
+    module_and_name = _lookup_module_and_qualname(obj, name=obj.__name__)
+    if module_and_name is None:
+        return (_make_typevar, _decompose_typevar(obj))
+    return (getattr, module_and_name)

Diff for: cloudpickle/cloudpickle_fast.py

@@ -20,14 +20,15 @@
 import sys
 import types
 import weakref
+import typing
 
 from _pickle import Pickler
 
 from .cloudpickle import (
     _is_dynamic, _extract_code_globals, _BUILTIN_TYPE_NAMES, DEFAULT_PROTOCOL,
-    _find_imported_submodules, _get_cell_contents, _is_global, _builtin_type,
+    _find_imported_submodules, _get_cell_contents, _is_importable_by_name, _builtin_type,
     Enum, _ensure_tracking,  _make_skeleton_class, _make_skeleton_enum,
-    _extract_class_dict, dynamic_subimport, subimport
+    _extract_class_dict, dynamic_subimport, subimport, _typevar_reduce,
 )
 
 load, loads = _pickle.load, _pickle.loads
@@ -332,7 +333,7 @@ def _class_reduce(obj):
         return type, (NotImplemented,)
     elif obj in _BUILTIN_TYPE_NAMES:
         return _builtin_type, (_BUILTIN_TYPE_NAMES[obj],)
-    elif not _is_global(obj):
+    elif not _is_importable_by_name(obj):
         return _dynamic_class_reduce(obj)
     return NotImplemented
 
@@ -422,6 +423,7 @@ class CloudPickler(Pickler):
     dispatch[types.MethodType] = _method_reduce
     dispatch[types.MappingProxyType] = _mappingproxy_reduce
     dispatch[weakref.WeakSet] = _weakset_reduce
+    dispatch[typing.TypeVar] = _typevar_reduce
 
     def __init__(self, file, protocol=None, buffer_callback=None):
         if protocol is None:
@@ -503,7 +505,7 @@ def _function_reduce(self, obj):
         As opposed to cloudpickle.py, There no special handling for builtin
         pypy functions because cloudpickle_fast is CPython-specific.
         """
-        if _is_global(obj):
+        if _is_importable_by_name(obj):
             return NotImplemented
         else:
             return self._dynamic_function_reduce(obj)

Diff for: tests/cloudpickle_test.py

@@ -46,6 +46,7 @@
 from cloudpickle.cloudpickle import _is_dynamic
 from cloudpickle.cloudpickle import _make_empty_cell, cell_set
 from cloudpickle.cloudpickle import _extract_class_dict, _whichmodule
+from cloudpickle.cloudpickle import _lookup_module_and_qualname
 
 from .testutils import subprocess_pickle_echo
 from .testutils import assert_run_python_script
@@ -2110,11 +2111,55 @@ class LocallyDefinedClass:
         reconstructed = pickle.loads(pickle_bytes, buffers=buffers)
         np.testing.assert_allclose(reconstructed.data, data_instance.data)
 
+    def test_pickle_dynamic_typevar(self):
+        T = typing.TypeVar('T')
+        depickled_T = pickle_depickle(T, protocol=self.protocol)
+        attr_list = [
+            "__name__", "__bound__", "__constraints__", "__covariant__",
+            "__contravariant__"
+        ]
+        for attr in attr_list:
+            assert getattr(T, attr) == getattr(depickled_T, attr)
+
+    def test_pickle_importable_typevar(self):
+        from .mypkg import T
+        T1 = pickle_depickle(T, protocol=self.protocol)
+        assert T1 is T
+
+        # Standard Library TypeVar
+        from typing import AnyStr
+        assert AnyStr is pickle_depickle(AnyStr, protocol=self.protocol)
+
 
 class Protocol2CloudPickleTest(CloudPickleTest):
 
     protocol = 2
 
 
+def test_lookup_module_and_qualname_dynamic_typevar():
+    T = typing.TypeVar('T')
+    module_and_name = _lookup_module_and_qualname(T, name=T.__name__)
+    assert module_and_name is None
+
+
+def test_lookup_module_and_qualname_importable_typevar():
+    from . import mypkg
+    T = mypkg.T
+    module_and_name = _lookup_module_and_qualname(T, name=T.__name__)
+    assert module_and_name is not None
+    module, name = module_and_name
+    assert module is mypkg
+    assert name == 'T'
+
+
+def test_lookup_module_and_qualname_stdlib_typevar():
+    module_and_name = _lookup_module_and_qualname(typing.AnyStr,
+                                                  name=typing.AnyStr.__name__)
+    assert module_and_name is not None
+    module, name = module_and_name
+    assert module is typing
+    assert name == 'AnyStr'
+
+
 if __name__ == '__main__':
     unittest.main()

Diff for: tests/mypkg/__init__.py

@@ -1,3 +1,4 @@
+import typing
 from .mod import module_function
 
 
@@ -13,3 +14,4 @@ def __reduce__(self):
 
 
 some_singleton = _SingletonClass()
+T = typing.TypeVar('T')