Pickle typing.TypeVars not as globals

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
 =====
 

cloudpickle/cloudpickle.py

@@ -60,6 +60,7 @@
 import weakref
 import uuid
 import threading
+import typing
 from enum import Enum
 
 from pickle import _Pickler as Pickler
@@ -118,7 +119,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
@@ -141,22 +142,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:
@@ -165,18 +190,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):
@@ -420,6 +447,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())
 
@@ -469,7 +501,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)
@@ -772,7 +804,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)
@@ -1216,3 +1248,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)

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)

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()

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')