diff --git a/cloudpickle/cloudpickle.py b/cloudpickle/cloudpickle.py index 537a00cc5..7fb6ee559 100644 --- a/cloudpickle/cloudpickle.py +++ b/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), ) diff --git a/cloudpickle/cloudpickle_fast.py b/cloudpickle/cloudpickle_fast.py index 49453d5b1..7df95f8dc 100644 --- a/cloudpickle/cloudpickle_fast.py +++ b/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 diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index eab45763b..c1a55134d 100644 --- a/tests/cloudpickle_test.py +++ b/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),