Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 1 | """runpy.py - locating and running Python code using the module namespace |
| 2 | |
| 3 | Provides support for locating and running Python scripts using the Python |
| 4 | module namespace instead of the native filesystem. |
| 5 | |
| 6 | This allows Python code to play nicely with non-filesystem based PEP 302 |
| 7 | importers when locating support scripts as well as when importing modules. |
| 8 | """ |
| 9 | # Written by Nick Coghlan <ncoghlan at gmail.com> |
| 10 | # to implement PEP 338 (Executing Modules as Scripts) |
| 11 | |
| 12 | import sys |
| 13 | import imp |
Thomas Wouters | 49fd7fa | 2006-04-21 10:40:58 +0000 | [diff] [blame] | 14 | try: |
| 15 | from imp import get_loader |
| 16 | except ImportError: |
| 17 | from pkgutil import get_loader |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 18 | |
| 19 | __all__ = [ |
| 20 | "run_module", |
| 21 | ] |
| 22 | |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 23 | |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 24 | def _run_code(code, run_globals, init_globals=None, |
| 25 | mod_name=None, mod_fname=None, |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 26 | mod_loader=None, pkg_name=None): |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 27 | """Helper for _run_module_code""" |
| 28 | if init_globals is not None: |
| 29 | run_globals.update(init_globals) |
| 30 | run_globals.update(__name__ = mod_name, |
| 31 | __file__ = mod_fname, |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 32 | __loader__ = mod_loader, |
| 33 | __package__ = pkg_name) |
Georg Brandl | 7cae87c | 2006-09-06 06:51:57 +0000 | [diff] [blame] | 34 | exec(code, run_globals) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 35 | return run_globals |
| 36 | |
| 37 | def _run_module_code(code, init_globals=None, |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 38 | mod_name=None, mod_fname=None, |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 39 | mod_loader=None, pkg_name=None): |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 40 | """Helper for run_module""" |
| 41 | # Set up the top level namespace dictionary |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 42 | temp_module = imp.new_module(mod_name) |
| 43 | mod_globals = temp_module.__dict__ |
| 44 | # Modify sys.argv[0] and sys.module[mod_name] |
| 45 | saved_argv0 = sys.argv[0] |
| 46 | restore_module = mod_name in sys.modules |
| 47 | if restore_module: |
| 48 | saved_module = sys.modules[mod_name] |
| 49 | sys.argv[0] = mod_fname |
| 50 | sys.modules[mod_name] = temp_module |
| 51 | try: |
| 52 | _run_code(code, mod_globals, init_globals, |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 53 | mod_name, mod_fname, |
| 54 | mod_loader, pkg_name) |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 55 | finally: |
| 56 | sys.argv[0] = saved_argv0 |
| 57 | if restore_module: |
| 58 | sys.modules[mod_name] = saved_module |
| 59 | else: |
| 60 | del sys.modules[mod_name] |
| 61 | # Copy the globals of the temporary module, as they |
| 62 | # may be cleared when the temporary module goes away |
| 63 | return mod_globals.copy() |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 64 | |
| 65 | |
Thomas Wouters | 49fd7fa | 2006-04-21 10:40:58 +0000 | [diff] [blame] | 66 | # This helper is needed due to a missing component in the PEP 302 |
| 67 | # loader protocol (specifically, "get_filename" is non-standard) |
| 68 | def _get_filename(loader, mod_name): |
| 69 | try: |
| 70 | get_filename = loader.get_filename |
| 71 | except AttributeError: |
| 72 | return None |
| 73 | else: |
| 74 | return get_filename(mod_name) |
| 75 | |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 76 | # Helper to get the loader, code and filename for a module |
| 77 | def _get_module_details(mod_name): |
Thomas Wouters | 49fd7fa | 2006-04-21 10:40:58 +0000 | [diff] [blame] | 78 | loader = get_loader(mod_name) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 79 | if loader is None: |
Guido van Rossum | 806c246 | 2007-08-06 23:33:07 +0000 | [diff] [blame] | 80 | raise ImportError("No module named %s" % mod_name) |
| 81 | if loader.is_package(mod_name): |
| 82 | raise ImportError(("%s is a package and cannot " + |
| 83 | "be directly executed") % mod_name) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 84 | code = loader.get_code(mod_name) |
| 85 | if code is None: |
Guido van Rossum | 806c246 | 2007-08-06 23:33:07 +0000 | [diff] [blame] | 86 | raise ImportError("No code object available for %s" % mod_name) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 87 | filename = _get_filename(loader, mod_name) |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 88 | return loader, code, filename |
| 89 | |
| 90 | |
| 91 | # XXX ncoghlan: Should this be documented and made public? |
Christian Heimes | c3f30c4 | 2008-02-22 16:37:40 +0000 | [diff] [blame] | 92 | # (Current thoughts: don't repeat the mistake that lead to its |
| 93 | # creation when run_module() no longer met the needs of |
| 94 | # mainmodule.c, but couldn't be changed because it was public) |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 95 | def _run_module_as_main(mod_name, set_argv0=True): |
| 96 | """Runs the designated module in the __main__ namespace |
| 97 | |
| 98 | These __*__ magic variables will be overwritten: |
| 99 | __file__ |
| 100 | __loader__ |
| 101 | """ |
Christian Heimes | c3f30c4 | 2008-02-22 16:37:40 +0000 | [diff] [blame] | 102 | try: |
| 103 | loader, code, fname = _get_module_details(mod_name) |
| 104 | except ImportError as exc: |
| 105 | # Try to provide a good error message |
| 106 | # for directories, zip files and the -m switch |
| 107 | if set_argv0: |
| 108 | # For -m switch, just disply the exception |
| 109 | info = str(exc) |
| 110 | else: |
| 111 | # For directories/zipfiles, let the user |
| 112 | # know what the code was looking for |
| 113 | info = "can't find '__main__.py' in %r" % sys.argv[0] |
| 114 | msg = "%s: %s" % (sys.executable, info) |
| 115 | sys.exit(msg) |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 116 | pkg_name = mod_name.rpartition('.')[0] |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 117 | main_globals = sys.modules["__main__"].__dict__ |
| 118 | if set_argv0: |
| 119 | sys.argv[0] = fname |
| 120 | return _run_code(code, main_globals, None, |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 121 | "__main__", fname, loader, pkg_name) |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 122 | |
| 123 | def run_module(mod_name, init_globals=None, |
| 124 | run_name=None, alter_sys=False): |
| 125 | """Execute a module's code without importing it |
| 126 | |
| 127 | Returns the resulting top level namespace dictionary |
| 128 | """ |
| 129 | loader, code, fname = _get_module_details(mod_name) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 130 | if run_name is None: |
| 131 | run_name = mod_name |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 132 | pkg_name = mod_name.rpartition('.')[0] |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 133 | if alter_sys: |
| 134 | return _run_module_code(code, init_globals, run_name, |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 135 | fname, loader, pkg_name) |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 136 | else: |
| 137 | # Leave the sys module alone |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 138 | return _run_code(code, {}, init_globals, run_name, |
| 139 | fname, loader, pkg_name) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 140 | |
| 141 | |
| 142 | if __name__ == "__main__": |
| 143 | # Run the module specified as the next command line argument |
| 144 | if len(sys.argv) < 2: |
Guido van Rossum | be19ed7 | 2007-02-09 05:37:30 +0000 | [diff] [blame] | 145 | print("No module specified for execution", file=sys.stderr) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 146 | else: |
| 147 | del sys.argv[0] # Make the requested module sys.argv[0] |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 148 | _run_module_as_main(sys.argv[0]) |