diff --git a/CHANGES.md b/CHANGES.md index 22878e500..3e0f25ee4 100644 --- a/CHANGES.md +++ b/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 ===== diff --git a/cloudpickle/cloudpickle.py b/cloudpickle/cloudpickle.py index c3635856f..bf92569c1 100644 --- a/cloudpickle/cloudpickle.py +++ b/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: diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index 7e1b128b5..4e8bc5e3b 100644 --- a/tests/cloudpickle_test.py +++ b/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 diff --git a/tests/testutils.py b/tests/testutils.py index 2a14e86c7..21cb2eafa 100644 --- a/tests/testutils.py +++ b/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