diff --git a/CHANGES.md b/CHANGES.md index 160e0d68f..9658a676b 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,3 +1,10 @@ +0.8.1 +===== + +- Fix a bug (already present before 0.5.3 and re-introduced in 0.8.0) + affecting relative import instructions inside depickled functions + ([issue #254](https://github.com/cloudpipe/cloudpickle/pull/254)) + 0.8.0 ===== diff --git a/cloudpickle/cloudpickle.py b/cloudpickle/cloudpickle.py index 75c29977e..b5f414ac2 100644 --- a/cloudpickle/cloudpickle.py +++ b/cloudpickle/cloudpickle.py @@ -666,6 +666,15 @@ def extract_func_data(self, func): # multiple invokations are bound to the same Cloudpickler. base_globals = self.globals_ref.setdefault(id(func.__globals__), {}) + if base_globals == {}: + # Add module attributes used to resolve relative imports + # instructions inside func. + for k in ["__package__", "__name__", "__path__", "__file__"]: + # Some built-in functions/methods such as object.__new__ have + # their __globals__ set to None in PyPy + if func.__globals__ is not None and k in func.__globals__: + base_globals[k] = func.__globals__[k] + return (code, f_globals, defaults, closure, dct, base_globals) def save_builtin_function(self, obj): diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index 6ffe2ad0c..caa222ed3 100644 --- a/tests/cloudpickle_test.py +++ b/tests/cloudpickle_test.py @@ -1368,6 +1368,30 @@ def test_dataclass(self): pickle_depickle(DataClass, protocol=self.protocol) assert data.x == pickle_depickle(data, protocol=self.protocol).x == 42 + def test_relative_import_inside_function(self): + # Make sure relative imports inside round-tripped functions is not + # broken.This was a bug in cloudpickle versions <= 0.5.3 and was + # re-introduced in 0.8.0. + + # Both functions living inside modules and packages are tested. + def f(): + # module_function belongs to mypkg.mod1, which is a module + from .mypkg import module_function + return module_function() + + def g(): + # package_function belongs to mypkg, which is a package + from .mypkg import package_function + return package_function() + + for func, source in zip([f, g], ["module", "package"]): + # Make sure relative imports are initially working + assert func() == "hello from a {}!".format(source) + + # Make sure relative imports still work after round-tripping + cloned_func = pickle_depickle(func, protocol=self.protocol) + assert cloned_func() == "hello from a {}!".format(source) + class Protocol2CloudPickleTest(CloudPickleTest): diff --git a/tests/mypkg/__init__.py b/tests/mypkg/__init__.py new file mode 100644 index 000000000..fe3cc6b1d --- /dev/null +++ b/tests/mypkg/__init__.py @@ -0,0 +1,6 @@ +from .mod import module_function + + +def package_function(): + """Function living inside a package, not a simple module""" + return "hello from a package!" diff --git a/tests/mypkg/mod.py b/tests/mypkg/mod.py new file mode 100644 index 000000000..a703a6aa5 --- /dev/null +++ b/tests/mypkg/mod.py @@ -0,0 +1,2 @@ +def module_function(): + return "hello from a module!"