Skip to content

Commit

Permalink
FIX globals for functions defined in __main__ (#188)
Browse files Browse the repository at this point in the history
Fix #187 

The cause of this issue is that global variables defined in `__main__` are not shared by unpickled functions and their value are overwritten at each unpickling.
  • Loading branch information
tomMoral authored and ogrisel committed Aug 23, 2018
1 parent a249c44 commit 9a52ba3
Show file tree
Hide file tree
Showing 6 changed files with 80 additions and 3 deletions.
1 change: 1 addition & 0 deletions .coveragerc
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
[run]
branch = True
parallel = True
source = cloudpickle

[report]
Expand Down
4 changes: 3 additions & 1 deletion .travis.yml
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,11 @@ before_script:
- flake8 . --count --select=E901,E999,F821,F822,F823 --show-source --statistics
# exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
- flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
- python ci/install_coverage_subprocess_pth.py
script:
- if [[ $TRAVIS_PYTHON_VERSION != 'pypy'* ]]; then source activate testenv; fi
- PYTHONPATH='.:tests' py.test -r s --cov-config .coveragerc --cov=cloudpickle
- COVERAGE_PROCESS_START="$TRAVIS_BUILD_DIR/.coveragerc" PYTHONPATH='.:tests' py.test -r s
after_success:
- if [[ $TRAVIS_PYTHON_VERSION != 'pypy'* ]]; then source activate testenv; fi
- coverage combine --append
- codecov
16 changes: 16 additions & 0 deletions ci/install_coverage_subprocess_pth.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
# Make it possible to enable test coverage reporting for Python
# code run in children processes.
# http://coverage.readthedocs.io/en/latest/subprocess.html

import os.path as op
from distutils.sysconfig import get_python_lib

FILE_CONTENT = u"""\
import coverage; coverage.process_startup()
"""

filename = op.join(get_python_lib(), 'coverage_subprocess.pth')
with open(filename, 'wb') as f:
f.write(FILE_CONTENT.encode('ascii'))

print('Installed subprocess coverage support: %s' % filename)
19 changes: 17 additions & 2 deletions cloudpickle/cloudpickle.py
Original file line number Diff line number Diff line change
Expand Up @@ -431,6 +431,7 @@ def _save_subimports(self, code, top_level_dependencies):
Ensure de-pickler imports any package child-modules that
are needed by the function
"""

# check if any known dependency is an imported package
for x in top_level_dependencies:
if isinstance(x, types.ModuleType) and hasattr(x, '__package__') and x.__package__:
Expand Down Expand Up @@ -632,7 +633,15 @@ def extract_func_data(self, func):
# save the dict
dct = func.__dict__

base_globals = self.globals_ref.get(id(func.__globals__), {})
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__"
else:
base_globals = {}
self.globals_ref[id(func.__globals__)] = base_globals

return (code, f_globals, defaults, closure, dct, base_globals)
Expand Down Expand Up @@ -1037,7 +1046,11 @@ def _fill_function(*args):
else:
raise ValueError('Unexpected _fill_value arguments: %r' % (args,))

func.__globals__.update(state['globals'])
# Only set global variables that do not exist.
for k, v in state['globals'].items():
if k not in func.__globals__:
func.__globals__[k] = v

func.__defaults__ = state['defaults']
func.__dict__ = state['dict']
if 'annotations' in state:
Expand Down Expand Up @@ -1076,6 +1089,8 @@ 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])
base_globals['__builtins__'] = __builtins__

closure = (
Expand Down
39 changes: 39 additions & 0 deletions tests/cloudpickle_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -848,6 +848,45 @@ def f4(x):
""".format(protocol=self.protocol)
assert_run_python_script(textwrap.dedent(code))

def test_interactively_defined_global_variable(self):
# Check that callables defined in the __main__ module of a Python
# script (or jupyter kernel) correctly retrieve global variables.
code_template = """\
from testutils import subprocess_pickle_echo
from cloudpickle import dumps, loads
def local_clone(obj, protocol=None):
return loads(dumps(obj, protocol=protocol))
VARIABLE = "default_value"
def f0():
global VARIABLE
VARIABLE = "changed_by_f0"
def f1():
return VARIABLE
cloned_f0 = {clone_func}(f0, protocol={protocol})
cloned_f1 = {clone_func}(f1, protocol={protocol})
pickled_f1 = dumps(f1, protocol={protocol})
# Change the value of the global variable
cloned_f0()
# Ensure that the global variable is the same for another function
result_f1 = cloned_f1()
assert result_f1 == "changed_by_f0", result_f1
# Ensure that unpickling the global variable does not change its value
result_pickled_f1 = loads(pickled_f1)()
assert result_pickled_f1 == "changed_by_f0", result_pickled_f1
"""
for clone_func in ['local_clone', 'subprocess_pickle_echo']:
code = code_template.format(protocol=self.protocol,
clone_func=clone_func)
assert_run_python_script(textwrap.dedent(code))

@pytest.mark.skipif(sys.version_info >= (3, 0),
reason="hardcoded pickle bytes for 2.7")
def test_function_pickle_compat_0_4_0(self):
Expand Down
4 changes: 4 additions & 0 deletions tests/testutils.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,10 @@ def assert_run_python_script(source_code, timeout=5):
'stderr': STDOUT,
'env': {'PYTHONPATH': pythonpath},
}
# If coverage is running, pass the config file to the subprocess
coverage_rc = os.environ.get("COVERAGE_PROCESS_START")
if coverage_rc:
kwargs['env']['COVERAGE_PROCESS_START'] = coverage_rc
if timeout_supported:
kwargs['timeout'] = timeout
try:
Expand Down

0 comments on commit 9a52ba3

Please sign in to comment.