Enable tracking for dynamically defined TypeVar instances (#353)

* Ensure that dynamic typevars are tracked and reconciled as dynamic classes and enums

* Rename _ensure_tracking to _get_or_create_tracker_id

Author: ogrisel

cloudpickle/cloudpickle.py

@@ -89,7 +89,7 @@
 _extract_code_globals_cache = weakref.WeakKeyDictionary()
 
 
-def _ensure_tracking(class_def):
+def _get_or_create_tracker_id(class_def):
     with _DYNAMIC_CLASS_TRACKER_LOCK:
         class_tracker_id = _DYNAMIC_CLASS_TRACKER_BY_CLASS.get(class_def)
         if class_tracker_id is None:
@@ -544,7 +544,7 @@ def _save_dynamic_enum(self, obj, clsdict):
         self.save_reduce(
                 _make_skeleton_enum,
                 (obj.__bases__, obj.__name__, obj.__qualname__,
-                 members, obj.__module__, _ensure_tracking(obj), None),
+                 members, obj.__module__, _get_or_create_tracker_id(obj), None),
                 obj=obj
          )
 
@@ -646,7 +646,7 @@ def save_dynamic_class(self, obj):
             tp = type(obj)
             self.save_reduce(_make_skeleton_class,
                              (tp, obj.__name__, _get_bases(obj), type_kwargs,
-                              _ensure_tracking(obj), None),
+                              _get_or_create_tracker_id(obj), None),
                              obj=obj)
 
         # Now save the rest of obj's __dict__. Any references to obj
@@ -1251,17 +1251,20 @@ def _is_dynamic(module):
     return _find_spec(module.__name__, pkgpath, module) is None
 
 
-def _make_typevar(name, bound, constraints, covariant, contravariant):
-    return typing.TypeVar(
+def _make_typevar(name, bound, constraints, covariant, contravariant,
+                  class_tracker_id):
+    tv = typing.TypeVar(
         name, *constraints, bound=bound,
         covariant=covariant, contravariant=contravariant
     )
+    return _lookup_class_or_track(class_tracker_id, tv)
 
 
 def _decompose_typevar(obj):
     return (
         obj.__name__, obj.__bound__, obj.__constraints__,
         obj.__covariant__, obj.__contravariant__,
+        _get_or_create_tracker_id(obj),
     )
 
 

cloudpickle/cloudpickle_fast.py

@@ -27,7 +27,7 @@
 from .cloudpickle import (
     _is_dynamic, _extract_code_globals, _BUILTIN_TYPE_NAMES, DEFAULT_PROTOCOL,
     _find_imported_submodules, _get_cell_contents, _is_importable_by_name, _builtin_type,
-    Enum, _ensure_tracking,  _make_skeleton_class, _make_skeleton_enum,
+    Enum, _get_or_create_tracker_id,  _make_skeleton_class, _make_skeleton_enum,
     _extract_class_dict, dynamic_subimport, subimport, _typevar_reduce, _get_bases,
 )
 
@@ -77,13 +77,13 @@ def _class_getnewargs(obj):
         type_kwargs['__dict__'] = __dict__
 
     return (type(obj), obj.__name__, _get_bases(obj), type_kwargs,
-            _ensure_tracking(obj), None)
+            _get_or_create_tracker_id(obj), None)
 
 
 def _enum_getnewargs(obj):
     members = dict((e.name, e.value) for e in obj)
     return (obj.__bases__, obj.__name__, obj.__qualname__, members,
-            obj.__module__, _ensure_tracking(obj), None)
+            obj.__module__, _get_or_create_tracker_id(obj), None)
 
 
 # COLLECTION OF OBJECTS RECONSTRUCTORS

tests/cloudpickle_test.py

@@ -2122,6 +2122,11 @@ def test_pickle_dynamic_typevar(self):
         for attr in attr_list:
             assert getattr(T, attr) == getattr(depickled_T, attr)
 
+    def test_pickle_dynamic_typevar_tracking(self):
+        T = typing.TypeVar("T")
+        T2 = subprocess_pickle_echo(T, protocol=self.protocol)
+        assert T is T2
+
     def test_pickle_dynamic_typevar_memoization(self):
         T = typing.TypeVar('T')
         depickled_T1, depickled_T2 = pickle_depickle((T, T),