Fix relative imports inside function body (#254)

Author: pierreglaser

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

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

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

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

tests/mypkg/mod.py

@@ -0,0 +1,2 @@
+def module_function():
+    return "hello from a module!"