diff --git a/CHANGES.md b/CHANGES.md index a6b0b443..d586240e 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,6 +1,10 @@ In development ============== +- Allow importable module subclasses with default reduction hooks to be + referenced when pickling functions on CPython. Custom reducers retain their + existing behavior; by-value module subclass reconstruction is unchanged. + - 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..0eb582e9 100644 --- a/cloudpickle/cloudpickle.py +++ b/cloudpickle/cloudpickle.py @@ -1367,7 +1367,7 @@ def __init__(self, file, protocol=None, buffer_callback=None): # not be the case anymore when pypy implements protocol 5. def reducer_override(self, obj): - """Type-agnostic reducing callback for function and classes. + """Type-agnostic callback for functions, classes and module subclasses. For performance reasons, subclasses of the C `pickle.Pickler` class cannot register custom reducers for functions and classes in the @@ -1407,6 +1407,24 @@ def reducer_override(self, obj): return _class_reduce(obj) elif isinstance(obj, types.FunctionType): return self._function_reduce(obj) + elif isinstance(obj, types.ModuleType) and t not in self.dispatch_table: + # Module subclasses miss the exact-type dispatch entry (#397). + # Only fill in their default, unsupported reduction: custom + # dispatch entries and reduction hooks must retain precedence. + reduce_ex = getattr(obj, "__reduce_ex__", None) + if not ( + isinstance(reduce_ex, types.BuiltinMethodType) + and reduce_ex == object.__reduce_ex__.__get__(obj) + ): + return NotImplemented + reduce = getattr(obj, "__reduce__", None) + if ( + isinstance(reduce, types.BuiltinMethodType) + and reduce == object.__reduce__.__get__(obj) + and _should_pickle_by_reference(obj) + ): + return _module_reduce(obj) + return NotImplemented else: # fallback to save_global, including the Pickler's # dispatch_table diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index e2097d1c..cbdad547 100644 --- a/tests/cloudpickle_test.py +++ b/tests/cloudpickle_test.py @@ -87,6 +87,210 @@ def _maybe_remove(list_, item): return list_ +@pytest.fixture +def module_facade(): + class Facade(types.ModuleType): + pass + + module = Facade("_cloudpickle_module_facade") + module.foo = 42 + return module + + +def _module_facade_roundtrip( + module, protocol, shape, pickler_class=cloudpickle.Pickler +): + obj = module + if shape == "function": + namespace = {"mod": module, "__name__": "__main__"} + exec("def f():\n return mod.foo\n", namespace) + obj = namespace["f"] + buffer = io.BytesIO() + pickler_class(buffer, protocol=protocol).dump(obj) + cloned = pickle.loads(buffer.getvalue()) + if shape == "function": + assert cloned is not obj + return cloned.__globals__["mod"] + return cloned + + +@pytest.mark.skipif( + platform.python_implementation() != "CPython", reason="CPython hook" +) +@pytest.mark.parametrize("protocol", [2, 4, 5]) +@pytest.mark.parametrize("shape", ["module", "function"]) +def test_importable_module_facade(module_facade, monkeypatch, protocol, shape): + # #397: a module subclass must not prevent by-value function serialization. + module = module_facade + monkeypatch.setitem(sys.modules, module.__name__, module) + state = module.__dict__.copy() + cloned = _module_facade_roundtrip(module, protocol, shape) + assert cloned is module + assert cloned.foo == 42 + assert module.__dict__ == state + + +@pytest.mark.skipif( + platform.python_implementation() != "CPython", reason="CPython hook" +) +@pytest.mark.parametrize("protocol", [2, 4, 5]) +@pytest.mark.parametrize("shape", ["module", "function"]) +@pytest.mark.parametrize("by_value", [False, True]) +def test_module_facade_by_value_is_not_flattened( + module_facade, monkeypatch, protocol, shape, by_value +): + # By-value subclass reconstruction is a separate, unsupported contract. + module = module_facade + state = module.__dict__.copy() + if by_value: + monkeypatch.setitem(sys.modules, module.__name__, module) + register_pickle_by_value(module) + try: + with pytest.raises(TypeError, match="cannot pickle.*Facade"): + _module_facade_roundtrip(module, protocol, shape) + assert module.__dict__ == state + finally: + if by_value: + unregister_pickle_by_value(module) + + +@pytest.mark.parametrize("protocol", [2, 4, 5]) +@pytest.mark.parametrize("shape", ["module", "function"]) +@pytest.mark.parametrize( + "reducer", + [ + "copyreg", + "dispatch", + "reduce", + "reduce_ex", + "instance_reduce", + "instance_reduce_ex", + "getattribute", + ], +) +def test_module_facade_custom_reducer(monkeypatch, protocol, shape, reducer): + import copyreg + + def reduce_module(module): + return str, (reducer,) + + class Facade(types.ModuleType): + pass + + if reducer == "reduce": + Facade.__reduce__ = reduce_module + elif reducer == "reduce_ex": + Facade.__reduce_ex__ = lambda self, protocol: reduce_module(self) + elif reducer == "getattribute": + + def getattribute(self, name): + if name == "__reduce_ex__": + return lambda protocol: reduce_module(self) + return types.ModuleType.__getattribute__(self, name) + + Facade.__getattribute__ = getattribute + + module = Facade("_cloudpickle_custom_module_facade") + module.foo = 42 + monkeypatch.setitem(sys.modules, module.__name__, module) + pickler_class = cloudpickle.Pickler + if reducer == "copyreg": + monkeypatch.setitem(copyreg.dispatch_table, Facade, reduce_module) + elif reducer == "dispatch": + + class CustomPickler(cloudpickle.Pickler): + dispatch_table = cloudpickle.Pickler.dispatch_table.new_child( + {Facade: reduce_module} + ) + + pickler_class = CustomPickler + elif reducer == "instance_reduce": + module.__reduce__ = lambda: reduce_module(module) + elif reducer == "instance_reduce_ex": + module.__reduce_ex__ = lambda protocol: reduce_module(module) + if reducer == "instance_reduce": + # ModuleType's inherited __reduce_ex__ does not delegate to an instance + # __reduce__ override. Keep that existing failure instead of ignoring + # the override and accidentally importing the module by reference. + with pytest.raises(TypeError, match="cannot pickle.*Facade"): + _module_facade_roundtrip(module, protocol, shape, pickler_class) + else: + assert ( + _module_facade_roundtrip(module, protocol, shape, pickler_class) == reducer + ) + + +@pytest.mark.skipif( + platform.python_implementation() != "CPython", reason="CPython hook" +) +@pytest.mark.parametrize("protocol", [2, 4, 5]) +def test_module_facade_fresh_process(tmp_path, monkeypatch, protocol): + module_name = "_cloudpickle_importable_module_facade" + (tmp_path / (module_name + ".py")).write_text( + "import sys\nfrom types import ModuleType\n" + "class Facade(ModuleType):\n pass\n" + "sys.modules[__name__].__class__ = Facade\nfoo = 42\n", + encoding="utf-8", + ) + monkeypatch.syspath_prepend(str(tmp_path)) + module = __import__(module_name) + try: + namespace = {"mod": module, "__name__": "__main__"} + exec("def f():\n return mod.foo\n", namespace) + payload = cloudpickle.dumps(namespace["f"], protocol=protocol) + script = ( + "import sys, pickle\n" + "sys.path.insert(0, sys.argv[1])\n" + "f = pickle.loads(sys.stdin.buffer.read())\n" + "assert f() == 42\n" + f"assert f.__globals__['mod'] is __import__({module_name!r})\n" + ) + result = subprocess.run( + [sys.executable, "-I", "-c", script, str(tmp_path)], + input=payload, + capture_output=True, + timeout=60, + ) + assert result.returncode == 0, result.stderr.decode("utf-8") + missing = subprocess.run( + [ + sys.executable, + "-I", + "-c", + "import sys, pickle; pickle.loads(sys.stdin.buffer.read())", + ], + input=payload, + capture_output=True, + timeout=60, + ) + assert missing.returncode != 0 + assert b"ModuleNotFoundError" in missing.stderr + assert module_name.encode() in missing.stderr + finally: + sys.modules.pop(module_name) + + +@pytest.mark.parametrize("protocol", [2, 4, 5]) +@pytest.mark.parametrize("shape", ["module", "function"]) +@pytest.mark.parametrize("hidden", ["__name__", "__reduce__"]) +def test_module_facade_custom_reducer_no_introspection( + monkeypatch, protocol, shape, hidden +): + class Facade(types.ModuleType): + def __reduce_ex__(self, protocol): + return str, ("custom",) + + def __getattribute__(self, name): + if name == hidden: + raise ValueError("must not inspect custom reducer module") + return types.ModuleType.__getattribute__(self, name) + + module = Facade("_cloudpickle_hidden_module_facade") + module.foo = 42 + monkeypatch.setitem(sys.modules, "_cloudpickle_hidden_module_facade", module) + assert _module_facade_roundtrip(module, protocol, shape) == "custom" + + def test_extract_class_dict(): class A(int): """A docstring"""