diff --git a/CHANGES.md b/CHANGES.md index a6b0b443..6b23200f 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,6 +1,10 @@ In development ============== +- Pickle `typing.NewType` objects defined in `__main__` by value on Python + 3.10 and later. + ([issue #520](https://github.com/cloudpipe/cloudpickle/issues/520)) + - Make pickling of functions depending on globals in notebook more deterministic. ([PR#560](https://github.com/cloudpipe/cloudpickle/pull/560)) diff --git a/cloudpickle/cloudpickle.py b/cloudpickle/cloudpickle.py index 08882306..dbfbf581 100644 --- a/cloudpickle/cloudpickle.py +++ b/cloudpickle/cloudpickle.py @@ -6,7 +6,7 @@ In particular this pickler adds the following features: - serialize interactively-defined or locally-defined functions, classes, - enums, typevars, lambdas and nested functions to compiled byte code; + enums, typevars, newtypes, lambdas and nested functions to compiled byte code; - deal with some other non-serializable objects in an ad-hoc manner where applicable. @@ -630,6 +630,35 @@ def _typevar_reduce(obj): return (getattr, module_and_name) +def _make_newtype(name, supertype, module, class_tracker_id): + newtype = typing.NewType(name, supertype) + # NewType.__init__ records the caller module. Reconstruction runs inside + # cloudpickle, so put the original module back before tracking the object. + newtype.__module__ = module + return _lookup_class_or_track(class_tracker_id, newtype) + + +def _decompose_newtype(obj): + return ( + obj.__qualname__, + obj.__supertype__, + obj.__module__, + _get_or_create_tracker_id(obj), + ) + + +def _newtype_reduce(obj): + # Same constraint as TypeVar: module information decides reference vs value. + module_and_name = _lookup_module_and_qualname(obj, name=obj.__name__) + + if module_and_name is None: + return (_make_newtype, _decompose_newtype(obj)) + elif _is_registered_pickle_by_value(module_and_name[0]): + return (_make_newtype, _decompose_newtype(obj)) + + return (getattr, module_and_name) + + def _get_bases(typ): if "__orig_bases__" in getattr(typ, "__dict__", {}): # For generic types (see PEP 560) @@ -1253,6 +1282,10 @@ class Pickler(pickle.Pickler): _dispatch_table[types.MappingProxyType] = _mappingproxy_reduce _dispatch_table[weakref.WeakSet] = _weakset_reduce _dispatch_table[typing.TypeVar] = _typevar_reduce + if isinstance(typing.NewType, type): + # Python 3.8 and 3.9 implement NewType as a function, which the + # function reducer already pickles by value when it is dynamic. + _dispatch_table[typing.NewType] = _newtype_reduce _dispatch_table[_collections_abc.dict_keys] = _dict_keys_reduce _dispatch_table[_collections_abc.dict_values] = _dict_values_reduce _dispatch_table[_collections_abc.dict_items] = _dict_items_reduce diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index e2097d1c..9616c842 100644 --- a/tests/cloudpickle_test.py +++ b/tests/cloudpickle_test.py @@ -2591,6 +2591,116 @@ def test_pickle_importable_typevar(self): assert AnyStr is pickle_depickle(AnyStr, protocol=self.protocol) + def _skip_unless_newtype_is_class(self): + if not isinstance(typing.NewType, type): + self.skipTest("typing.NewType is a function before Python 3.10") + + def test_pickle_dynamic_newtype(self): + self._skip_unless_newtype_is_class() + run_dict = typing.NewType("RunDict", dict) + restored = pickle_depickle(run_dict, protocol=self.protocol) + assert restored is run_dict + assert restored({"a": 1}) == {"a": 1} + + echoed = subprocess_pickle_echo(run_dict, protocol=self.protocol) + assert echoed is run_dict + assert echoed({"a": 1}) == {"a": 1} + + def test_pickle_dynamic_newtype_memoization(self): + self._skip_unless_newtype_is_class() + run_dict = typing.NewType("RunDict", dict) + first, second = pickle_depickle((run_dict, run_dict), protocol=self.protocol) + assert first is second + + def test_pickle_nested_newtype(self): + self._skip_unless_newtype_is_class() + inner = typing.NewType("Inner", int) + outer = typing.NewType("Outer", inner) + restored_outer = pickle_depickle(outer, protocol=self.protocol) + assert restored_outer is outer + assert restored_outer.__supertype__ is inner + assert restored_outer.__qualname__ == "Outer" + assert restored_outer.__module__ == outer.__module__ + assert restored_outer(5) == 5 + + def test_pickle_newtype_dotted_qualname(self): + self._skip_unless_newtype_is_class() + alias = typing.NewType("Foo.Bar", int) + restored = pickle_depickle(alias, protocol=self.protocol) + assert restored is alias + assert restored.__qualname__ == "Foo.Bar" + assert restored.__name__ == "Bar" + assert restored.__supertype__ is int + + def test_pickle_importable_newtype_by_reference(self): + self._skip_unless_newtype_is_class() + module_name = "cloudpickle_newtype_ref_mod" + module = types.ModuleType(module_name) + user_id = typing.NewType("UserId", int) + user_id.__module__ = module_name + module.UserId = user_id + sys.modules[module_name] = module + try: + payload = cloudpickle.dumps(user_id, protocol=self.protocol) + assert b"_make_newtype" not in payload + assert pickle.loads(payload) is user_id + finally: + sys.modules.pop(module_name, None) + + def test_pickle_registered_newtype_by_value(self): + self._skip_unless_newtype_is_class() + module_name = "cloudpickle_newtype_value_mod" + module = types.ModuleType(module_name) + user_id = typing.NewType("UserId", int) + user_id.__module__ = module_name + module.UserId = user_id + sys.modules[module_name] = module + registered = False + try: + register_pickle_by_value(module) + registered = True + payload = cloudpickle.dumps(user_id, protocol=self.protocol) + assert b"_make_newtype" in payload + echoed = subprocess_pickle_echo(user_id, protocol=self.protocol) + assert echoed.__name__ == "UserId" + assert echoed.__supertype__ is int + assert echoed(3) == 3 + finally: + if registered: + unregister_pickle_by_value(module) + sys.modules.pop(module_name, None) + + def test_newtype_defined_in_main(self): + # https://github.com/cloudpipe/cloudpickle/issues/520 + self._skip_unless_newtype_is_class() + folder = Path(self.tmpdir) + pickle_path = folder / "tmp.pickle" + (folder / "a.py").write_text( + "import cloudpickle\n" + "from typing import Any, NewType\n" + "RunDict = NewType('RunDict', dict[str, Any])\n" + f"open(r'{pickle_path}', 'wb').write(cloudpickle.dumps(RunDict))\n", + encoding="utf-8", + ) + (folder / "b.py").write_text( + "import cloudpickle\n" + f"value = cloudpickle.load(open(r'{pickle_path}', 'rb'))\n" + "assert value.__name__ == 'RunDict'\n" + "assert value({'a': 1}) == {'a': 1}\n", + encoding="utf-8", + ) + env = os.environ.copy() + repo = Path(__file__).resolve().parents[1] + env["PYTHONPATH"] = os.pathsep.join( + [str(repo), env.get("PYTHONPATH", "")] + ).rstrip(os.pathsep) + subprocess.check_call( + [sys.executable, str(folder / "a.py")], cwd=folder, env=env + ) + subprocess.check_call( + [sys.executable, str(folder / "b.py")], cwd=folder, env=env + ) + def test_generic_type(self): T = typing.TypeVar("T")