diff -r 51dddfead80a Lib/multiprocessing/spawn.py --- a/Lib/multiprocessing/spawn.py Sun Dec 15 20:33:02 2013 +1000 +++ b/Lib/multiprocessing/spawn.py Sun Dec 15 22:35:38 2013 +1000 @@ -11,6 +11,8 @@ import os import pickle import sys +import runpy +import types from . import get_start_method, set_start_method from . import process @@ -148,6 +150,9 @@ else: sys_path[i] = process.ORIGINAL_DIR + main_module = sys.modules['__main__'] + main_mod_name = getattr(main_module.__spec__, "name", None) + d.update( name=name, sys_path=sys_path, @@ -157,14 +162,27 @@ start_method=get_start_method(), ) - if sys.platform != 'win32' or (not WINEXE and not WINSERVICE): - main_path = getattr(sys.modules['__main__'], '__file__', None) + if main_mod_name is not None: + d['main_mod_name'] = main_mod_name + elif sys.platform != 'win32' or (not WINEXE and not WINSERVICE): + main_path = getattr(main_module, '__file__', None) if not main_path and sys.argv[0] not in ('', '-c'): + # This can happen if (for example), the -i command line switch + # is used to open an interactive prompt after directly + # executing a script main_path = sys.argv[0] if main_path is not None: if (not os.path.isabs(main_path) and process.ORIGINAL_DIR is not None): main_path = os.path.join(process.ORIGINAL_DIR, main_path) + # Running the test suite through -m, should never get here... + # ... but we do in child processes, so there's currently + # something going wrong with getting __main__.__spec__ + # populated. This guard at least makes those cases fail fast, + # rather than spawning grandchild processes that then misbehave + main_name = os.path.splitext(os.path.basename(main_path))[0] + if main_name == '__main__': + raise RuntimeError('%r looks like a module path' % main_path) d['main_path'] = os.path.normpath(main_path) return d @@ -206,55 +224,64 @@ if 'start_method' in data: set_start_method(data['start_method']) - if 'main_path' in data: - import_main_path(data['main_path']) + if 'main_mod_name' in data: + _fixup_main_from_name(data['main_mod_name']) + elif 'main_path' in data: + _fixup_main_from_path(data['main_path']) + +# Multiprocessing module helpers to fix up the main module in +# spawned subprocesses +def _fixup_main_from_name(mod_name): + # __main__.py files for packages, directories, zip archives, etc, run + # their "main only" code unconditionally, so we don't even try to + # populate anything in __main__, nor do we make any changes to + # __main__ attributes + if mod_name == "__main__" or mod_name.endswith(".__main__"): + return + + # If this process was forked, __main__ may already be populated + current_main = sys.modules['__main__'] + if getattr(current_main.__spec__, "name", None) == mod_name: + return + + # Otherwise, __main__ may contain some non-main code where we need to + # support unpickling it properly. We rerun it as __mp_main__ and make + # the normal __main__ an alias to that + old_main_modules.append(current_main) + main_module = types.ModuleType("__mp_main__") + main_content = runpy.run_module(mod_name, + run_name="__mp_main__", + alter_sys=True) + main_module.__dict__.update(main_content) + sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module + + +def _fixup_main_from_path(main_path): + # If this process was forked, __main__ may already be populated + current_main = sys.modules['__main__'] + if getattr(current_main, '__file__', None) == main_path: + return + + # Unfortunately, the main ipython launch script currently has no + # "if __name__ == '__main__'" guard, so we work around that here + # See https://github.com/ipython/ipython/issues/4698 + main_name = os.path.splitext(os.path.basename(main_path))[0] + if main_name == 'ipython': + return + + # If the parent process has sent a path through rather than a module + # name we assume it is an executable script that may contain + # non-main code that needs to be executed + old_main_modules.append(current_main) + main_module = types.ModuleType("__mp_main__") + main_content = runpy.run_path(main_path, + run_name="__mp_main__") + main_module.__dict__.update(main_content) + sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module def import_main_path(main_path): ''' Set sys.modules['__main__'] to module at main_path ''' - # XXX (ncoghlan): The following code makes several bogus - # assumptions regarding the relationship between __file__ - # and a module's real name. See PEP 302 and issue #10845 - if getattr(sys.modules['__main__'], '__file__', None) == main_path: - return - - main_name = os.path.splitext(os.path.basename(main_path))[0] - if main_name == '__init__': - main_name = os.path.basename(os.path.dirname(main_path)) - - if main_name == '__main__': - main_module = sys.modules['__main__'] - main_module.__file__ = main_path - elif main_name != 'ipython': - # Main modules not actually called __main__.py may - # contain additional code that should still be executed - import importlib - import types - - if main_path is None: - dirs = None - elif os.path.basename(main_path).startswith('__init__.py'): - dirs = [os.path.dirname(os.path.dirname(main_path))] - else: - dirs = [os.path.dirname(main_path)] - - assert main_name not in sys.modules, main_name - sys.modules.pop('__mp_main__', None) - # We should not try to load __main__ - # since that would execute 'if __name__ == "__main__"' - # clauses, potentially causing a psuedo fork bomb. - main_module = types.ModuleType(main_name) - # XXX Use a target of main_module? - spec = importlib.find_spec(main_name, path=dirs) - if spec is None: - raise ImportError(name=main_name) - methods = importlib._bootstrap._SpecMethods(spec) - methods.init_module_attrs(main_module) - main_module.__name__ = '__mp_main__' - code = spec.loader.get_code(main_name) - exec(code, main_module.__dict__) - - old_main_modules.append(sys.modules['__main__']) - sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module + _fixup_main_from_path(main_path)