Pickle NewType objects defined in __main__ by value

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

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

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