FIX More func base globals fixes

CHANGES.md

@@ -1,7 +1,10 @@
-
 master
 ======
 
+- global variables from modules referenced in sys.modules in the child process
+  now overrides the initial global variables of the pickled function.
+  ([issue #187](https://github.com/cloudpipe/cloudpickle/issues/187))
+
 
 0.5.5
 =====
@@ -19,7 +22,6 @@ master
   variables ([issue #187](
   https://github.com/cloudpipe/cloudpickle/issues/187)).
 
-
 0.5.3
 =====
 - Fixed a crash in Python 2 when serializing non-hashable instancemethods of built-in

cloudpickle/cloudpickle.py

@@ -78,6 +78,7 @@
     PY3 = True
 
 
+
 def _make_cell_set_template_code():
     """Get the Python compiler to emit LOAD_FAST(arg); STORE_DEREF
 
@@ -635,11 +636,12 @@ def extract_func_data(self, func):
 
         base_globals = self.globals_ref.get(id(func.__globals__), None)
         if base_globals is None:
-            # For functions defined in __main__, use vars(__main__) for
-            # base_global. This is necessary to share the global variables
-            # across multiple functions in this module.
-            if func.__module__ == "__main__":
-                base_globals = "__main__"
+            # For functions defined in a well behaved module use
+            # vars(func.__module__) for base_globals. This is necessary to
+            # share the global variables across multiple pickled functions from
+            # this module.
+            if hasattr(func, '__module__') and func.__module__ is not None:
+                base_globals = func.__module__
             else:
                 base_globals = {}
         self.globals_ref[id(func.__globals__)] = base_globals
@@ -934,7 +936,6 @@ def subimport(name):
 def dynamic_subimport(name, vars):
     mod = imp.new_module(name)
     mod.__dict__.update(vars)
-    sys.modules[name] = mod
     return mod
 
 
@@ -1090,7 +1091,13 @@ def _make_skel_func(code, cell_count, base_globals=None):
     if base_globals is None:
         base_globals = {}
     elif isinstance(base_globals, str):
-        base_globals = vars(sys.modules[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:
+            base_globals = {}
+
     base_globals['__builtins__'] = __builtins__
 
     closure = (

tests/cloudpickle_test.py

@@ -17,6 +17,7 @@
 import textwrap
 import unittest
 import weakref
+import os
 
 try:
     from StringIO import StringIO
@@ -47,6 +48,9 @@
 from .testutils import assert_run_python_script
 
 
+_TEST_GLOBAL_VARIABLE = "default_value"
+
+
 class RaiserOnPickle(object):
 
     def __init__(self, exc):
@@ -887,6 +891,107 @@ def f1():
                                         clone_func=clone_func)
             assert_run_python_script(textwrap.dedent(code))
 
+    def test_closure_interacting_with_a_global_variable(self):
+        global _TEST_GLOBAL_VARIABLE
+        orig_value = _TEST_GLOBAL_VARIABLE
+        try:
+            def f0():
+                global _TEST_GLOBAL_VARIABLE
+                _TEST_GLOBAL_VARIABLE = "changed_by_f0"
+
+            def f1():
+                return _TEST_GLOBAL_VARIABLE
+
+            cloned_f0 = cloudpickle.loads(cloudpickle.dumps(
+                f0, protocol=self.protocol))
+            cloned_f1 = cloudpickle.loads(cloudpickle.dumps(
+                f1, protocol=self.protocol))
+            pickled_f1 = cloudpickle.dumps(f1, protocol=self.protocol)
+
+            # Change the value of the global variable
+            cloned_f0()
+            assert _TEST_GLOBAL_VARIABLE == "changed_by_f0"
+
+            # Ensure that the global variable is the same for another function
+            result_cloned_f1 = cloned_f1()
+            assert result_cloned_f1 == "changed_by_f0", result_cloned_f1
+            assert f1() == result_cloned_f1
+
+            # Ensure that unpickling the global variable does not change its
+            # value
+            result_pickled_f1 = cloudpickle.loads(pickled_f1)()
+            assert result_pickled_f1 == "changed_by_f0", result_pickled_f1
+        finally:
+            _TEST_GLOBAL_VARIABLE = orig_value
+
+    def test_function_from_dynamic_module_with_globals_modifications(self):
+        """
+        unpickling functions coming from a dynamic module should lead to
+        new global variables being created each time.
+        Hence, a modification of the globals of one function should not
+        interact with the globals of another function.
+        This test validates this behavior.
+        """
+
+        # first, we create a dynamic module in the parent process
+        mod = imp.new_module('mod')
+        code = '''
+        x = 1
+        def func_that_relies_on_dynamic_module():
+            global x
+            return x
+        '''
+        exec(textwrap.dedent(code), mod.__dict__)
+
+        try:
+            # simple sanity check on the function's output
+            assert mod.func_that_relies_on_dynamic_module() == 1
+
+            # the function of mod is pickled two times, with two different
+            # values for the global variable x.
+
+            # a child process that sequentially unpickles the
+            # two functions is then launched, and asserts
+            # the independance of those two global variables
+
+            with open('first_function.pk', 'wb') as fid:
+                cloudpickle.dump(mod.func_that_relies_on_dynamic_module, fid)
+
+            # change the mod's global variable x
+            mod.x = 2
+
+            # at this point, mod.func_that_relies_on_dynamic_module()
+            # returns 2
+            assert mod.func_that_relies_on_dynamic_module() == 2
+            with open('function_with_modified_globals.pk', 'wb') as fid:
+                cloudpickle.dump(mod.func_that_relies_on_dynamic_module, fid)
+
+            child_process_code = """
+                import pickle
+
+                with open('first_function.pk','rb') as fid:
+                    first_f = pickle.load(fid)
+
+                assert first_f() == 1
+
+                # load a function with initial global variable x set to 2
+                with open('function_with_modified_globals.pk','rb') as fid:
+                    new_f = pickle.load(fid)
+
+                # assert that by first_f's  global variable x does not interact
+                # with new_f's global variable x
+                assert new_f()==2
+
+            """
+
+            # finally, we execute the code
+            assert_run_python_script(textwrap.dedent(child_process_code))
+
+        finally:
+            # remove the created files
+            os.unlink('first_function.pk')
+            os.unlink('function_with_modified_globals.pk')
+
     @pytest.mark.skipif(sys.version_info >= (3, 0),
                         reason="hardcoded pickle bytes for 2.7")
     def test_function_pickle_compat_0_4_0(self):