FIX Handling of global variables by locally defined functions (#198)

Author: pierreglaser

CHANGES.md

@@ -1,7 +1,10 @@
-
 master
 ======
 
+  - Ensure that unpickling a locally defined function that accesses the global variables
+  of a module does not reset the values of the global variables if they are already initialized.
+  ([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

@@ -635,11 +635,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 +935,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 +1090,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):
@@ -436,6 +440,74 @@ def method(self, x):
         mod1, mod2 = pickle_depickle([mod, mod])
         self.assertEqual(id(mod1), id(mod2))
 
+    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
+        # the child process and reloaded in another one.
+
+        # We create a new dynamic module
+        mod = imp.new_module('mod')
+        code = '''
+        x = 1
+        '''
+        exec(textwrap.dedent(code), mod.__dict__)
+
+        # This script will be ran in a separate child process. It will import
+        # the pickled dynamic module, and then re-pickle it under a new name.
+        # Finally, it will create a child process that will load the re-pickled
+        # dynamic module.
+        parent_process_module_file = 'dynamic_module_from_parent_process.pkl'
+        child_process_module_file = 'dynamic_module_from_child_process.pkl'
+        child_process_script = '''
+            import pickle
+            import textwrap
+
+            import cloudpickle
+            from testutils import assert_run_python_script
+
+
+            child_of_child_process_script = {child_of_child_process_script}
+
+            with open('{parent_process_module_file}', 'rb') as f:
+                mod = pickle.load(f)
+
+            with open('{child_process_module_file}', 'wb') as f:
+                cloudpickle.dump(mod, f)
+
+            assert_run_python_script(textwrap.dedent(child_of_child_process_script))
+            '''
+
+        # The script ran by the process created by the child process
+        child_of_child_process_script = """ '''
+                import pickle
+                with open('{child_process_module_file}','rb') as fid:
+                    mod = pickle.load(fid)
+                ''' """
+
+        # Filling the two scripts with the pickled modules filepaths and,
+        # for the first child process, the script to be executed by its
+        # own child process.
+        child_of_child_process_script = child_of_child_process_script.format(
+                child_process_module_file=child_process_module_file)
+
+        child_process_script = child_process_script.format(
+                parent_process_module_file=parent_process_module_file,
+                child_process_module_file=child_process_module_file,
+                child_of_child_process_script=child_of_child_process_script)
+
+        try:
+            with open(parent_process_module_file, 'wb') as fid:
+                cloudpickle.dump(mod, fid)
+
+            assert_run_python_script(textwrap.dedent(child_process_script))
+
+        finally:
+            # Remove temporary created files
+            if os.path.exists(parent_process_module_file):
+                os.unlink(parent_process_module_file)
+            if os.path.exists(child_process_module_file):
+                os.unlink(child_process_module_file)
+
     def test_find_module(self):
         import pickle  # ensure this test is decoupled from global imports
         _find_module('pickle')
@@ -887,6 +959,40 @@ 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
+        assert _TEST_GLOBAL_VARIABLE == "default_value"
+        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
+
     @pytest.mark.skipif(sys.version_info >= (3, 0),
                         reason="hardcoded pickle bytes for 2.7")
     def test_function_pickle_compat_0_4_0(self):