commit 2bf6e446a571865bc7f2dc3a1051dd12f026cf10 Author: Mike Kaplinskiy Date: Sat Feb 20 18:18:41 2016 -0800 Add argv argument to runpy. diff --git a/Lib/runpy.py b/Lib/runpy.py index af6205d..274d109 100644 --- a/Lib/runpy.py +++ b/Lib/runpy.py @@ -17,9 +17,11 @@ import types from pkgutil import read_code, get_importer __all__ = [ - "run_module", "run_path", + "run_module", "run_path", "INHERIT" ] +INHERIT = object() # special constant to inherit argv from sys.argv + class _TempModule(object): """Temporarily replace a module in sys.modules with an empty namespace""" def __init__(self, mod_name): @@ -43,20 +45,33 @@ class _TempModule(object): del sys.modules[self.mod_name] self._saved_module = [] -class _ModifiedArgv0(object): - def __init__(self, value): - self.value = value +class _ModifiedArgv(object): + def __init__(self, value, mode): + if mode is None: + self.replacement = None + else: + self.replacement = [value] + if mode is INHERIT: + self.replacement.extend(sys.argv[1:]) + else: + self.replacement.extend(mode) + 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 + if self.replacement is None: + return + + sys.argv[:] = self._saved_value # TODO: Replace these helpers with importlib._bootstrap_external functions. def _run_code(code, run_globals, init_globals=None, @@ -86,11 +101,12 @@ def _run_code(code, run_globals, init_globals=None, return run_globals def _run_module_code(code, init_globals=None, - mod_name=None, mod_spec=None, - pkg_name=None, script_name=None): + mod_name=None, mod_spec=None, + pkg_name=None, script_name=None, + argv=INHERIT): """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): + with _TempModule(mod_name) as temp_module, _ModifiedArgv(fname, argv): mod_globals = temp_module.module.__dict__ _run_code(code, mod_globals, init_globals, mod_name, mod_spec, pkg_name, script_name) @@ -184,7 +200,7 @@ def _run_module_as_main(mod_name, alter_argv=True): "__main__", mod_spec) def run_module(mod_name, init_globals=None, - run_name=None, alter_sys=False): + run_name=None, alter_sys=False, argv=INHERIT): """Execute a module's code without importing it Returns the resulting top level namespace dictionary @@ -193,7 +209,8 @@ def run_module(mod_name, init_globals=None, if run_name is None: run_name = mod_name if alter_sys: - return _run_module_code(code, init_globals, run_name, mod_spec) + return _run_module_code(code, init_globals, run_name, mod_spec, + argv=argv) else: # Leave the sys module alone return _run_code(code, {}, init_globals, run_name, mod_spec) @@ -227,7 +244,7 @@ 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): +def run_path(path_name, init_globals=None, run_name=None, argv=INHERIT): """Execute code located at the specified filesystem location Returns the resulting top level namespace dictionary @@ -251,7 +268,8 @@ def run_path(path_name, init_globals=None, run_name=None): # 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) + pkg_name=pkg_name, script_name=fname, + argv=argv) else: # Importer is defined for path, so add it to # the start of sys.path @@ -265,10 +283,10 @@ def run_path(path_name, init_globals=None, run_name=None): # existing __main__ module may prevent location of the new module. mod_name, mod_spec, code = _get_main_module_details() with _TempModule(run_name) as temp_module, \ - _ModifiedArgv0(path_name): + _ModifiedArgv(path_name, argv): mod_globals = temp_module.module.__dict__ return _run_code(code, mod_globals, init_globals, - run_name, mod_spec, pkg_name).copy() + run_name, mod_spec, pkg_name).copy() finally: try: sys.path.remove(path_name) diff --git a/Lib/test/test_runpy.py b/Lib/test/test_runpy.py index 87c83ec..fa6d05d 100644 --- a/Lib/test/test_runpy.py +++ b/Lib/test/test_runpy.py @@ -32,7 +32,7 @@ f() del f # Check the sys module import sys -run_argv0 = sys.argv[0] +run_argv = sys.argv[:] run_name_in_sys_modules = __name__ in sys.modules module_in_sys_modules = (run_name_in_sys_modules and globals() is sys.modules[__name__].__dict__) @@ -53,7 +53,7 @@ example_namespace = { "sys": sys, "runpy": runpy, "result": ["Top level assignment", "Lower level reference"], - "run_argv0": sys.argv[0], + "run_argv": sys.argv[:], "run_name_in_sys_modules": False, "module_in_sys_modules": False, "nested": dict(implicit_namespace, @@ -123,12 +123,14 @@ class CodeExecutionMixin: sentinel = object() expected_ns = expected_namespace.copy() run_name = expected_ns["__name__"] - saved_argv0 = sys.argv[0] + saved_argv = sys.argv[:] saved_mod = sys.modules.get(run_name, sentinel) # Check without initial globals result_ns = create_namespace(None) self.assertNamespaceMatches(result_ns, expected_ns) - self.assertIs(sys.argv[0], saved_argv0) + self.assertEqual(sys.argv, saved_argv) + for i in range(len(sys.argv)): + self.assertIs(sys.argv[i], saved_argv[i]) self.assertIs(sys.modules.get(run_name, sentinel), saved_mod) # And then with initial globals initial_ns = {"sentinel": sentinel} @@ -136,9 +138,19 @@ class CodeExecutionMixin: result_ns = create_namespace(initial_ns) self.assertIsNot(result_ns, initial_ns) self.assertNamespaceMatches(result_ns, expected_ns) - self.assertIs(sys.argv[0], saved_argv0) + self.assertEqual(sys.argv, saved_argv) + for i in range(len(sys.argv)): + self.assertIs(sys.argv[i], saved_argv[i]) self.assertIs(sys.modules.get(run_name, sentinel), saved_mod) + def _get_expected_argv(self, mod_fname, argv): + if argv is runpy.INHERIT: + return [mod_fname] + sys.argv[1:] + elif argv is None: + return sys.argv[:] + else: + return [mod_fname] + list(argv) + class ExecutionLayerTestCase(unittest.TestCase, CodeExecutionMixin): """Unit tests for runpy._run_code and runpy._run_module_code""" @@ -167,7 +179,7 @@ class ExecutionLayerTestCase(unittest.TestCase, CodeExecutionMixin): "__loader__": mod_loader, "__package__": mod_package, "__spec__": mod_spec, - "run_argv0": mod_fname, + "run_argv": [mod_fname] + sys.argv[1:], "run_name_in_sys_modules": True, "module_in_sys_modules": True, }) @@ -271,7 +283,7 @@ class RunModuleTestCase(unittest.TestCase, CodeExecutionMixin): except OSError as ex: if verbose > 1: print(ex) # Persist with cleaning up - def _fix_ns_for_legacy_pyc(self, ns, alter_sys): + def _fix_ns_for_legacy_pyc(self, ns, alter_sys, argv): char_to_add = "c" ns["__file__"] += char_to_add ns["__cached__"] = ns["__file__"] @@ -279,11 +291,10 @@ class RunModuleTestCase(unittest.TestCase, CodeExecutionMixin): new_spec = importlib.util.spec_from_file_location(spec.name, ns["__file__"]) ns["__spec__"] = new_spec - if alter_sys: - ns["run_argv0"] += char_to_add - + if alter_sys and argv is not None: + ns["run_argv"][0] += char_to_add - def _check_module(self, depth, alter_sys=False, + def _check_module(self, depth, alter_sys=False, argv=runpy.INHERIT, *, namespace=False, parent_namespaces=False): pkg_dir, mod_fname, mod_name, mod_spec = ( self._make_pkg(example_source, depth, @@ -300,12 +311,13 @@ class RunModuleTestCase(unittest.TestCase, CodeExecutionMixin): }) if alter_sys: expected_ns.update({ - "run_argv0": mod_fname, + "run_argv": self._get_expected_argv(mod_fname, argv), "run_name_in_sys_modules": True, "module_in_sys_modules": True, }) def create_ns(init_globals): - return run_module(mod_name, init_globals, alter_sys=alter_sys) + return run_module(mod_name, init_globals, alter_sys=alter_sys, + argv=argv) try: if verbose > 1: print("Running from source:", mod_name) self.check_code_execution(create_ns, expected_ns) @@ -317,13 +329,13 @@ class RunModuleTestCase(unittest.TestCase, CodeExecutionMixin): unload(mod_name) # In case loader caches paths importlib.invalidate_caches() if verbose > 1: print("Running from compiled:", mod_name) - self._fix_ns_for_legacy_pyc(expected_ns, alter_sys) + self._fix_ns_for_legacy_pyc(expected_ns, alter_sys, argv) self.check_code_execution(create_ns, expected_ns) finally: self._del_pkg(pkg_dir, depth, mod_name) if verbose > 1: print("Module executed successfully") - def _check_package(self, depth, alter_sys=False, + def _check_package(self, depth, alter_sys=False, argv=runpy.INHERIT, *, namespace=False, parent_namespaces=False): pkg_dir, mod_fname, mod_name, mod_spec = ( self._make_pkg(example_source, depth, "__main__", @@ -341,12 +353,13 @@ class RunModuleTestCase(unittest.TestCase, CodeExecutionMixin): }) if alter_sys: expected_ns.update({ - "run_argv0": mod_fname, + "run_argv": self._get_expected_argv(mod_fname, argv), "run_name_in_sys_modules": True, "module_in_sys_modules": True, }) def create_ns(init_globals): - return run_module(pkg_name, init_globals, alter_sys=alter_sys) + return run_module(pkg_name, init_globals, alter_sys=alter_sys, + argv=argv) try: if verbose > 1: print("Running from source:", pkg_name) self.check_code_execution(create_ns, expected_ns) @@ -358,7 +371,7 @@ class RunModuleTestCase(unittest.TestCase, CodeExecutionMixin): unload(mod_name) # In case loader caches paths if verbose > 1: print("Running from compiled:", pkg_name) importlib.invalidate_caches() - self._fix_ns_for_legacy_pyc(expected_ns, alter_sys) + self._fix_ns_for_legacy_pyc(expected_ns, alter_sys, argv) self.check_code_execution(create_ns, expected_ns) finally: self._del_pkg(pkg_dir, depth, pkg_name) @@ -490,11 +503,31 @@ from ..uncle.cousin import nephew if verbose > 1: print("Testing package depth:", depth) self._check_module(depth, alter_sys=True) + def test_run_module_alter_sys_argv(self): + for depth in range(4): + if verbose > 1: print("Testing package depth:", depth) + self._check_module(depth, alter_sys=True, argv=('hi', 'there')) + + def test_run_module_alter_sys_argv_None(self): + for depth in range(4): + if verbose > 1: print("Testing package depth:", depth) + self._check_module(depth, alter_sys=True, argv=None) + def test_run_package_alter_sys(self): for depth in range(1, 4): if verbose > 1: print("Testing package depth:", depth) self._check_package(depth, alter_sys=True) + def test_run_package_alter_sys_argv(self): + for depth in range(1, 4): + if verbose > 1: print("Testing package depth:", depth) + self._check_package(depth, alter_sys=True, argv=('hi', 'there')) + + def test_run_package_alter_sys_argv_None(self): + for depth in range(1, 4): + if verbose > 1: print("Testing package depth:", depth) + self._check_package(depth, alter_sys=True, argv=None) + def test_explicit_relative_import(self): for depth in range(2, 5): if verbose > 1: print("Testing relative imports at depth:", depth) @@ -572,11 +605,11 @@ class RunPathTestCase(unittest.TestCase, CodeExecutionMixin): source, omit_suffix) def _check_script(self, script_name, expected_name, expected_file, - expected_argv0, mod_name=None, + expected_argv0, mod_name=None, argv=runpy.INHERIT, expect_spec=True, check_loader=True): # First check is without run_name def create_ns(init_globals): - return run_path(script_name, init_globals) + return run_path(script_name, init_globals, argv=argv) expected_ns = example_namespace.copy() if mod_name is None: spec_name = expected_name @@ -597,7 +630,7 @@ class RunPathTestCase(unittest.TestCase, CodeExecutionMixin): "__cached__": mod_cached, "__package__": "", "__spec__": mod_spec, - "run_argv0": expected_argv0, + "run_argv": self._get_expected_argv(expected_argv0, argv), "run_name_in_sys_modules": True, "module_in_sys_modules": True, }) @@ -605,7 +638,7 @@ class RunPathTestCase(unittest.TestCase, CodeExecutionMixin): # Second check makes sure run_name works in all cases run_name = "prove.issue15230.is.fixed" def create_ns(init_globals): - return run_path(script_name, init_globals, run_name) + return run_path(script_name, init_globals, run_name, argv=argv) if expect_spec and mod_name is None: mod_spec = importlib.util.spec_from_file_location(run_name, expected_file) @@ -627,6 +660,21 @@ class RunPathTestCase(unittest.TestCase, CodeExecutionMixin): self._check_script(script_name, "", script_name, script_name, expect_spec=False) + def test_basic_script_argv(self): + with temp_dir() as script_dir: + mod_name = 'script' + script_name = self._make_test_script(script_dir, mod_name) + self._check_script(script_name, "", script_name, + script_name, expect_spec=False, + argv=('hi', 'there')) + + def test_basic_script_argv_None(self): + with temp_dir() as script_dir: + mod_name = 'script' + script_name = self._make_test_script(script_dir, mod_name) + self._check_script(script_name, "", script_name, + script_name, expect_spec=False, argv=None) + def test_basic_script_no_suffix(self): with temp_dir() as script_dir: mod_name = 'script'