Fix module locals has no builtins

CHANGES.md

@@ -1,3 +1,10 @@
+0.6.1
+=====
+
+- Fix regression in 0.6.0 which breaks the pickling of local function defined    in a module, making it impossible to access builtins functions.
+  ([issue #211](https://github.com/cloudpipe/cloudpickle/issues/211))
+
+
 0.6.0
 =====
 

cloudpickle/cloudpickle.py

@@ -42,19 +42,20 @@
 """
 from __future__ import print_function
 
-import dis
-from functools import partial
 import io
-import itertools
-import logging
+import dis
+import sys
+import types
 import opcode
-import operator
 import pickle
 import struct
-import sys
-import traceback
-import types
+import logging
 import weakref
+import operator
+import importlib
+import itertools
+import traceback
+from functools import partial
 
 
 # cloudpickle is meant for inter process communication: we expect all
@@ -1103,11 +1104,12 @@ def _make_skel_func(code, cell_count, base_globals=None):
         base_globals = {}
     elif isinstance(base_globals, str):
         base_globals_name = base_globals
-        if sys.modules.get(base_globals, None) is not None:
-            # This checks if we can import the previous environment the object
-            # lived in
-            base_globals = vars(sys.modules[base_globals])
-        else:
+        try:
+            # First try to reuse the globals from the module containing the
+            # function. If it is not possible to retrieve it, fallback to an
+            # empty dictionary.
+            base_globals = vars(importlib.import_module(base_globals))
+        except ImportError:
             base_globals = _dynamic_modules_globals.get(
                     base_globals_name, None)
             if base_globals is None:

tests/cloudpickle_test.py

@@ -290,6 +290,8 @@ def test_locally_defined_function_and_class(self):
         LOCAL_CONSTANT = 42
 
         def some_function(x, y):
+            # Make sure the __builtins__ are not broken (see #211)
+            sum(range(10))
             return (x + y) / LOCAL_CONSTANT
 
         # pickle the function definition
@@ -494,6 +496,38 @@ def func():
         finally:
             os.unlink(pickled_module_path)
 
+    def test_module_locals_behavior(self):
+        # Makes sure that a local function defined in another module is
+        # correctly serialized. This notably checks that the globals are
+        # accessible and that there is no issue with the builtins (see #211)
+
+        pickled_func_path = 'local_func_g.pkl'
+
+        child_process_script = '''
+        import pickle
+        import gc
+        with open("{pickled_func_path}", 'rb') as f:
+            func = pickle.load(f)
+
+        assert func(range(10)) == 45
+        '''
+
+        child_process_script = child_process_script.format(
+                pickled_func_path=pickled_func_path)
+
+        try:
+
+            from .testutils import make_local_function
+
+            g = make_local_function()
+            with open(pickled_func_path, 'wb') as f:
+                cloudpickle.dump(g, f)
+
+            assert_run_python_script(textwrap.dedent(child_process_script))
+
+        finally:
+            os.unlink(pickled_func_path)
+
     def test_load_dynamic_module_in_grandchild_process(self):
         # Make sure that when loaded, a dynamic module preserves its dynamic
         # property. Otherwise, this will lead to an ImportError if pickled in

tests/testutils.py

@@ -17,6 +17,19 @@ class TimeoutExpired(Exception):
     timeout_supported = False
 
 
+TEST_GLOBALS = "a test value"
+
+
+def make_local_function():
+    def g(x):
+        # this function checks that the globals are correctly handled and that
+        # the builtins are available
+        assert TEST_GLOBALS == "a test value"
+        return sum(range(10))
+
+    return g
+
+
 def subprocess_pickle_echo(input_data, protocol=None):
     """Echo function with a child Python process