Skip to content

Commit

Permalink
Enable tracking for dynamically defined TypeVar instances (#353)
Browse files Browse the repository at this point in the history
* Ensure that dynamic typevars are tracked and reconciled as dynamic classes and enums

* Rename _ensure_tracking to _get_or_create_tracker_id
  • Loading branch information
ogrisel authored Mar 15, 2020
1 parent 3e80b26 commit d8452cc
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 8 deletions.
13 changes: 8 additions & 5 deletions cloudpickle/cloudpickle.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
)

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),
)


Expand Down
6 changes: 3 additions & 3 deletions cloudpickle/cloudpickle_fast.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down Expand Up @@ -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
Expand Down
5 changes: 5 additions & 0 deletions tests/cloudpickle_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down

0 comments on commit d8452cc

Please sign in to comment.