commit 8ba58f31987e9f8ed1d2a5b5e7f4f7746dedbb82 Author: Mike Kaplinskiy Date: Sat Mar 5 12:00:24 2016 -0800 Add a runpy.RunnableModule abstraction that separates loading modules and running them. Adds two extra public functions: load_module & load_path. These return a RunnableModule that can be tweaked before calling `.run()`. The intended path for user code setting/changing argv is to use the load_* family, e.g. try: runnable = runpy.load_module('foo') catch Exception: runnable = runpy.load_path('foo.py', run_name='foo') runnable.argv = ['-m', runnable.name] + sys.argv[1:] runnable.run() diff --git a/Lib/runpy.py b/Lib/runpy.py index d86f0e4..ff3ebbd 100644 --- a/Lib/runpy.py +++ b/Lib/runpy.py @@ -17,7 +17,8 @@ import types from pkgutil import read_code, get_importer __all__ = [ - "run_module", "run_path", + "run_module", "run_path", "load_module", "load_path", "RunnableModule", + "ModifiedArgv", ] class _TempModule(object): @@ -43,60 +44,118 @@ class _TempModule(object): del sys.modules[self.mod_name] self._saved_module = [] -class _ModifiedArgv0(object): + +class ModifiedArgv(object): def __init__(self, value): - self.value = value + self.replacement = value self._saved_value = self._sentinel = object() def __enter__(self): + if self.replacement is None: + return + if self._saved_value is not self._sentinel: raise RuntimeError("Already preserving saved value") - self._saved_value = sys.argv[0] - sys.argv[0] = self.value + self._saved_value = sys.argv[:] + sys.argv[:] = self.replacement def __exit__(self, *args): - self.value = self._sentinel - sys.argv[0] = self._saved_value - -# TODO: Replace these helpers with importlib._bootstrap_external functions. -def _run_code(code, run_globals, init_globals=None, - mod_name=None, mod_spec=None, - pkg_name=None, script_name=None): - """Helper to run code in nominated namespace""" - if init_globals is not None: - run_globals.update(init_globals) - if mod_spec is None: - loader = None - fname = script_name - cached = None - else: - loader = mod_spec.loader - fname = mod_spec.origin - cached = mod_spec.cached - if pkg_name is None: - pkg_name = mod_spec.parent - run_globals.update(__name__ = mod_name, - __file__ = fname, - __cached__ = cached, - __doc__ = None, - __loader__ = loader, - __package__ = pkg_name, - __spec__ = mod_spec) - exec(code, run_globals) - return run_globals - -def _run_module_code(code, init_globals=None, - mod_name=None, mod_spec=None, - pkg_name=None, script_name=None): - """Helper to run code in new namespace with sys modified""" - fname = script_name if mod_spec is None else mod_spec.origin - with _TempModule(mod_name) as temp_module, _ModifiedArgv0(fname): - mod_globals = temp_module.module.__dict__ - _run_code(code, mod_globals, init_globals, - mod_name, mod_spec, pkg_name, script_name) - # Copy the globals of the temporary module, as they - # may be cleared when the temporary module goes away - return mod_globals.copy() + if self.replacement is None: + return + + sys.argv[:] = self._saved_value + + +class _AddedPath(object): + def __init__(self, value): + self.value = value + + def __enter__(self): + if self.value is None: + return + + sys.path.insert(0, self.value) + + def __exit__(self, *args): + if self.value is None: + return + + try: + sys.path.remove(self.value) + except ValueError: + pass + + +class RunnableModule(object): + def __init__(self, code, mod_name=None, mod_spec=None, script_name=None, pkg_name=None, + use_globals_from_sys_modules=False): + assert code is not None + self.globals = sys.modules[mod_name].__dict__ if use_globals_from_sys_modules else {} + self.code = code + self.module_spec = mod_spec + + self.name = mod_name + self.file = script_name if mod_spec is None else mod_spec.origin + self.package = pkg_name + if mod_spec is not None and pkg_name is None: + self.package = mod_spec.parent + + self.globals.update( + __cached__ = mod_spec.cached if mod_spec is not None else None, + __doc__ = None, + __loader__ = mod_spec.loader if mod_spec is not None else None, + __spec__ = mod_spec, + ) + + self.add_to_sys_modules = True + self.sys_path_addition = None + + @property + def name(self): + return self.globals.get('__name__') + + @name.setter + def name(self, value): + self.globals['__name__'] = value + + @property + def file(self): + return self.globals.get('__file__') + + @file.setter + def file(self, value): + self.globals['__file__'] = value + + @property + def package(self): + return self.globals.get('__package__') + + @package.setter + def package(self, value): + self.globals['__package__'] = value + + def run(self): + with _AddedPath(self.sys_path_addition): + if self.add_to_sys_modules: + if not self.name: + raise ValueError( + "Can't add a module without a name to sys.modules. " + "Either set .name or set .add_to_sys_modules = False") + with _TempModule(self.name) as temp_module: + mod_globals = temp_module.module.__dict__ + return self._exec_module(mod_globals) + else: + return self._exec_module(self.globals) + + def _exec_module(self, globals): + globals.update(self.globals) + exec(self.code, globals) + # Copy the globals of the temporary module, as they + # may be cleared when the temporary module goes away + self.globals = globals.copy() + + return self.globals + # Helper to get the full name, spec and code for a module def _get_module_details(mod_name, error=ImportError): @@ -157,6 +216,7 @@ def _get_module_details(mod_name, error=ImportError): raise error("No code object available for %s" % mod_name) return mod_name, spec, code + class _Error(Exception): """Error that _run_module_as_main() should report without a traceback""" @@ -186,26 +246,48 @@ def _run_module_as_main(mod_name, alter_argv=True): except _Error as exc: msg = "%s: %s" % (sys.executable, exc) sys.exit(msg) - main_globals = sys.modules["__main__"].__dict__ + if alter_argv: sys.argv[0] = mod_spec.origin - return _run_code(code, main_globals, None, - "__main__", mod_spec) -def run_module(mod_name, init_globals=None, - run_name=None, alter_sys=False): + runner = RunnableModule( + code=code, mod_name='__main__', mod_spec=mod_spec, + use_globals_from_sys_modules=True) + runner.add_to_sys_modules = False + runner.name = '__main__' + + return runner.run() + + +def _set_init_globals(runnable, init_globals): + """Sets globals without overwriting what's already there.""" + if init_globals is not None: + for k, v in init_globals.items(): + if k not in runnable.globals: + runnable.globals[k] = v + + +def load_module(mod_name): + mod_name, mod_spec, code = _get_module_details(mod_name) + return RunnableModule(code=code, mod_name=mod_name, mod_spec=mod_spec) + + +def run_module(mod_name, init_globals=None, run_name=None, alter_sys=False): """Execute a module's code without importing it Returns the resulting top level namespace dictionary """ - mod_name, mod_spec, code = _get_module_details(mod_name) - if run_name is None: - run_name = mod_name - if alter_sys: - return _run_module_code(code, init_globals, run_name, mod_spec) - else: - # Leave the sys module alone - return _run_code(code, {}, init_globals, run_name, mod_spec) + runnable = load_module(mod_name) + if run_name is not None: + runnable.name = run_name + + runnable.add_to_sys_modules = alter_sys + _set_init_globals(runnable, init_globals) + + argv = [runnable.file] + sys.argv[1:] if alter_sys else None + with ModifiedArgv(argv): + return runnable.run() + def _get_main_module_details(error=ImportError): # Helper that gives a nicer error message when attempting to @@ -236,16 +318,8 @@ def _get_code_from_file(run_name, fname): code = compile(f.read(), fname, 'exec') return code, fname -def run_path(path_name, init_globals=None, run_name=None): - """Execute code located at the specified filesystem location - - Returns the resulting top level namespace dictionary - The file path may refer directly to a Python script (i.e. - one that could be directly executed with execfile) or else - it may refer to a zipfile or directory containing a top - level __main__.py script. - """ +def load_path(path_name, run_name=None): if run_name is None: run_name = "" pkg_name = run_name.rpartition(".")[0] @@ -259,30 +333,33 @@ def run_path(path_name, init_globals=None, run_name=None): # Not a valid sys.path entry, so run the code directly # execfile() doesn't help as we want to allow compiled files code, fname = _get_code_from_file(run_name, path_name) - return _run_module_code(code, init_globals, run_name, - pkg_name=pkg_name, script_name=fname) + + runnable = RunnableModule(code=code, pkg_name=pkg_name, script_name=fname) + runnable.name = run_name + return runnable else: - # Finder is defined for path, so add it to - # the start of sys.path - sys.path.insert(0, path_name) - try: - # Here's where things are a little different from the run_module - # case. There, we only had to replace the module in sys while the - # code was running and doing so was somewhat optional. Here, we - # have no choice and we have to remove it even while we read the - # code. If we don't do this, a __loader__ attribute in the - # existing __main__ module may prevent location of the new module. + with _AddedPath(path_name): mod_name, mod_spec, code = _get_main_module_details() - with _TempModule(run_name) as temp_module, \ - _ModifiedArgv0(path_name): - mod_globals = temp_module.module.__dict__ - return _run_code(code, mod_globals, init_globals, - run_name, mod_spec, pkg_name).copy() - finally: - try: - sys.path.remove(path_name) - except ValueError: - pass + runnable = RunnableModule(code=code, mod_name=mod_name, mod_spec=mod_spec, pkg_name=pkg_name) + runnable.sys_path_addition = path_name + runnable.name = run_name + return runnable + + +def run_path(path_name, init_globals=None, run_name=None): + """Execute code located at the specified filesystem location + + Returns the resulting top level namespace dictionary + + The file path may refer directly to a Python script (i.e. + one that could be directly executed with execfile) or else + it may refer to a zipfile or directory containing a top + level __main__.py script. + """ + runnable = load_path(path_name, run_name=run_name) + _set_init_globals(runnable, init_globals) + with ModifiedArgv([path_name] + sys.argv[1:]): + return runnable.run() if __name__ == "__main__": diff --git a/Lib/test/test_runpy.py b/Lib/test/test_runpy.py index 02b4d62..ca9a124 100644 --- a/Lib/test/test_runpy.py +++ b/Lib/test/test_runpy.py @@ -16,7 +16,7 @@ from test.support.script_helper import ( import runpy -from runpy import _run_code, _run_module_code, run_module, run_path +from runpy import run_module, run_path, RunnableModule, ModifiedArgv # Note: This module can't safely test _run_module_as_main as it # runs its tests in the current process, which would mess with the # real __main__ module (usually test.regrtest) @@ -39,7 +39,7 @@ module_in_sys_modules = (run_name_in_sys_modules and globals() is sys.modules[__name__].__dict__) # Check nested operation import runpy -nested = runpy._run_module_code('x=1\\n', mod_name='') +nested = runpy.RunnableModule(code='x=1\\n', mod_name='').run() """ implicit_namespace = { @@ -112,7 +112,7 @@ class CodeExecutionMixin: for k in result_ns: actual = (k, result_ns[k]) expected = (k, expected_ns[k]) - self.assertEqual(actual, expected) + self.assertEqual(actual, expected, k) def check_code_execution(self, create_namespace, expected_namespace): """Check that an interface runs the example code correctly @@ -150,7 +150,11 @@ class ExecutionLayerTestCase(unittest.TestCase, CodeExecutionMixin): "__loader__": None, }) def create_ns(init_globals): - return _run_code(example_source, {}, init_globals) + runnable = RunnableModule(code=example_source) + runnable.add_to_sys_modules = False + if init_globals is not None: + runnable.globals.update(init_globals) + return runnable.run() self.check_code_execution(create_ns, expected_ns) def test_run_module_code(self): @@ -173,10 +177,14 @@ class ExecutionLayerTestCase(unittest.TestCase, CodeExecutionMixin): "module_in_sys_modules": True, }) def create_ns(init_globals): - return _run_module_code(example_source, - init_globals, - mod_name, - mod_spec) + runnable = RunnableModule(code=example_source, + mod_name=mod_name, + mod_spec=mod_spec) + argv = [mod_fname] + sys.argv[1:] + if init_globals is not None: + runnable.globals.update(init_globals) + with ModifiedArgv(argv): + return runnable.run() self.check_code_execution(create_ns, expected_ns) # TODO: Use self.addCleanup to get rid of a lot of try-finally blocks