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 |
Nick Coghlan | 260bd3e | 2009-11-16 06:49:25 +0000 | [diff] [blame] | 14 | from pkgutil import read_code |
Thomas Wouters | 49fd7fa | 2006-04-21 10:40:58 +0000 | [diff] [blame] | 15 | try: |
| 16 | from imp import get_loader |
| 17 | except ImportError: |
| 18 | from pkgutil import get_loader |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 19 | |
| 20 | __all__ = [ |
Nick Coghlan | 260bd3e | 2009-11-16 06:49:25 +0000 | [diff] [blame] | 21 | "run_module", "run_path", |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 22 | ] |
| 23 | |
Nick Coghlan | 260bd3e | 2009-11-16 06:49:25 +0000 | [diff] [blame] | 24 | class _TempModule(object): |
| 25 | """Temporarily replace a module in sys.modules with an empty namespace""" |
| 26 | def __init__(self, mod_name): |
| 27 | self.mod_name = mod_name |
| 28 | self.module = imp.new_module(mod_name) |
| 29 | self._saved_module = [] |
| 30 | |
| 31 | def __enter__(self): |
| 32 | mod_name = self.mod_name |
| 33 | try: |
| 34 | self._saved_module.append(sys.modules[mod_name]) |
| 35 | except KeyError: |
| 36 | pass |
| 37 | sys.modules[mod_name] = self.module |
| 38 | return self |
| 39 | |
| 40 | def __exit__(self, *args): |
| 41 | if self._saved_module: |
| 42 | sys.modules[self.mod_name] = self._saved_module[0] |
| 43 | else: |
| 44 | del sys.modules[self.mod_name] |
| 45 | self._saved_module = [] |
| 46 | |
| 47 | class _ModifiedArgv0(object): |
| 48 | def __init__(self, value): |
| 49 | self.value = value |
| 50 | self._saved_value = self._sentinel = object() |
| 51 | |
| 52 | def __enter__(self): |
| 53 | if self._saved_value is not self._sentinel: |
| 54 | raise RuntimeError("Already preserving saved value") |
| 55 | self._saved_value = sys.argv[0] |
| 56 | sys.argv[0] = self.value |
| 57 | |
| 58 | def __exit__(self, *args): |
| 59 | self.value = self._sentinel |
| 60 | sys.argv[0] = self._saved_value |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 61 | |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 62 | def _run_code(code, run_globals, init_globals=None, |
| 63 | mod_name=None, mod_fname=None, |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 64 | mod_loader=None, pkg_name=None): |
Benjamin Peterson | f6489f9 | 2009-11-25 17:46:26 +0000 | [diff] [blame] | 65 | """Helper to run code in nominated namespace""" |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 66 | if init_globals is not None: |
| 67 | run_globals.update(init_globals) |
| 68 | run_globals.update(__name__ = mod_name, |
| 69 | __file__ = mod_fname, |
Barry Warsaw | 28a691b | 2010-04-17 00:19:56 +0000 | [diff] [blame] | 70 | __cached__ = None, |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 71 | __loader__ = mod_loader, |
| 72 | __package__ = pkg_name) |
Georg Brandl | 7cae87c | 2006-09-06 06:51:57 +0000 | [diff] [blame] | 73 | exec(code, run_globals) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 74 | return run_globals |
| 75 | |
| 76 | def _run_module_code(code, init_globals=None, |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 77 | mod_name=None, mod_fname=None, |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 78 | mod_loader=None, pkg_name=None): |
Benjamin Peterson | f6489f9 | 2009-11-25 17:46:26 +0000 | [diff] [blame] | 79 | """Helper to run code in new namespace with sys modified""" |
Nick Coghlan | 260bd3e | 2009-11-16 06:49:25 +0000 | [diff] [blame] | 80 | with _TempModule(mod_name) as temp_module, _ModifiedArgv0(mod_fname): |
| 81 | mod_globals = temp_module.module.__dict__ |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 82 | _run_code(code, mod_globals, init_globals, |
Nick Coghlan | 260bd3e | 2009-11-16 06:49:25 +0000 | [diff] [blame] | 83 | mod_name, mod_fname, mod_loader, pkg_name) |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 84 | # Copy the globals of the temporary module, as they |
| 85 | # may be cleared when the temporary module goes away |
| 86 | return mod_globals.copy() |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 87 | |
| 88 | |
Thomas Wouters | 49fd7fa | 2006-04-21 10:40:58 +0000 | [diff] [blame] | 89 | # This helper is needed due to a missing component in the PEP 302 |
| 90 | # loader protocol (specifically, "get_filename" is non-standard) |
Nick Coghlan | f088e5e | 2008-12-14 11:50:48 +0000 | [diff] [blame] | 91 | # Since we can't introduce new features in maintenance releases, |
| 92 | # support was added to zipimporter under the name '_get_filename' |
Thomas Wouters | 49fd7fa | 2006-04-21 10:40:58 +0000 | [diff] [blame] | 93 | def _get_filename(loader, mod_name): |
Nick Coghlan | f088e5e | 2008-12-14 11:50:48 +0000 | [diff] [blame] | 94 | for attr in ("get_filename", "_get_filename"): |
| 95 | meth = getattr(loader, attr, None) |
| 96 | if meth is not None: |
| 97 | return meth(mod_name) |
| 98 | return None |
Thomas Wouters | 49fd7fa | 2006-04-21 10:40:58 +0000 | [diff] [blame] | 99 | |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 100 | # Helper to get the loader, code and filename for a module |
| 101 | def _get_module_details(mod_name): |
Thomas Wouters | 49fd7fa | 2006-04-21 10:40:58 +0000 | [diff] [blame] | 102 | loader = get_loader(mod_name) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 103 | if loader is None: |
Guido van Rossum | 806c246 | 2007-08-06 23:33:07 +0000 | [diff] [blame] | 104 | raise ImportError("No module named %s" % mod_name) |
| 105 | if loader.is_package(mod_name): |
Nick Coghlan | 3f48ae3 | 2009-02-08 01:58:26 +0000 | [diff] [blame] | 106 | if mod_name == "__main__" or mod_name.endswith(".__main__"): |
Benjamin Peterson | f6489f9 | 2009-11-25 17:46:26 +0000 | [diff] [blame] | 107 | raise ImportError("Cannot use package as __main__ module") |
Nick Coghlan | 3f48ae3 | 2009-02-08 01:58:26 +0000 | [diff] [blame] | 108 | try: |
| 109 | pkg_main_name = mod_name + ".__main__" |
| 110 | return _get_module_details(pkg_main_name) |
| 111 | except ImportError as e: |
| 112 | raise ImportError(("%s; %r is a package and cannot " + |
| 113 | "be directly executed") %(e, mod_name)) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 114 | code = loader.get_code(mod_name) |
| 115 | if code is None: |
Guido van Rossum | 806c246 | 2007-08-06 23:33:07 +0000 | [diff] [blame] | 116 | raise ImportError("No code object available for %s" % mod_name) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 117 | filename = _get_filename(loader, mod_name) |
Nick Coghlan | 3f48ae3 | 2009-02-08 01:58:26 +0000 | [diff] [blame] | 118 | return mod_name, loader, code, filename |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 119 | |
Benjamin Peterson | f6489f9 | 2009-11-25 17:46:26 +0000 | [diff] [blame] | 120 | # XXX ncoghlan: Should this be documented and made public? |
| 121 | # (Current thoughts: don't repeat the mistake that lead to its |
| 122 | # creation when run_module() no longer met the needs of |
| 123 | # mainmodule.c, but couldn't be changed because it was public) |
Nick Coghlan | 260bd3e | 2009-11-16 06:49:25 +0000 | [diff] [blame] | 124 | def _run_module_as_main(mod_name, alter_argv=True): |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 125 | """Runs the designated module in the __main__ namespace |
| 126 | |
Benjamin Peterson | f6489f9 | 2009-11-25 17:46:26 +0000 | [diff] [blame] | 127 | Note that the executed module will have full access to the |
| 128 | __main__ namespace. If this is not desirable, the run_module() |
R. David Murray | 445448c | 2009-12-20 17:28:31 +0000 | [diff] [blame] | 129 | function should be used to run the module code in a fresh namespace. |
Benjamin Peterson | f6489f9 | 2009-11-25 17:46:26 +0000 | [diff] [blame] | 130 | |
| 131 | At the very least, these variables in __main__ will be overwritten: |
| 132 | __name__ |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 133 | __file__ |
Barry Warsaw | 28a691b | 2010-04-17 00:19:56 +0000 | [diff] [blame] | 134 | __cached__ |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 135 | __loader__ |
Benjamin Peterson | f6489f9 | 2009-11-25 17:46:26 +0000 | [diff] [blame] | 136 | __package__ |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 137 | """ |
Christian Heimes | c3f30c4 | 2008-02-22 16:37:40 +0000 | [diff] [blame] | 138 | try: |
Nick Coghlan | 260bd3e | 2009-11-16 06:49:25 +0000 | [diff] [blame] | 139 | if alter_argv or mod_name != "__main__": # i.e. -m switch |
| 140 | mod_name, loader, code, fname = _get_module_details(mod_name) |
| 141 | else: # i.e. directory or zipfile execution |
| 142 | mod_name, loader, code, fname = _get_main_module_details() |
Christian Heimes | c3f30c4 | 2008-02-22 16:37:40 +0000 | [diff] [blame] | 143 | except ImportError as exc: |
Benjamin Peterson | f6489f9 | 2009-11-25 17:46:26 +0000 | [diff] [blame] | 144 | # Try to provide a good error message |
| 145 | # for directories, zip files and the -m switch |
| 146 | if alter_argv: |
| 147 | # For -m switch, just display the exception |
| 148 | info = str(exc) |
| 149 | else: |
| 150 | # For directories/zipfiles, let the user |
| 151 | # know what the code was looking for |
Benjamin Peterson | e360795 | 2009-11-25 18:38:11 +0000 | [diff] [blame] | 152 | info = "can't find '__main__' module in %r" % sys.argv[0] |
Benjamin Peterson | f6489f9 | 2009-11-25 17:46:26 +0000 | [diff] [blame] | 153 | msg = "%s: %s" % (sys.executable, info) |
Christian Heimes | c3f30c4 | 2008-02-22 16:37:40 +0000 | [diff] [blame] | 154 | sys.exit(msg) |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 155 | pkg_name = mod_name.rpartition('.')[0] |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 156 | main_globals = sys.modules["__main__"].__dict__ |
Nick Coghlan | 260bd3e | 2009-11-16 06:49:25 +0000 | [diff] [blame] | 157 | if alter_argv: |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 158 | sys.argv[0] = fname |
| 159 | return _run_code(code, main_globals, None, |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 160 | "__main__", fname, loader, pkg_name) |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 161 | |
| 162 | def run_module(mod_name, init_globals=None, |
| 163 | run_name=None, alter_sys=False): |
| 164 | """Execute a module's code without importing it |
| 165 | |
| 166 | Returns the resulting top level namespace dictionary |
| 167 | """ |
Nick Coghlan | 3f48ae3 | 2009-02-08 01:58:26 +0000 | [diff] [blame] | 168 | mod_name, loader, code, fname = _get_module_details(mod_name) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 169 | if run_name is None: |
| 170 | run_name = mod_name |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 171 | pkg_name = mod_name.rpartition('.')[0] |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 172 | if alter_sys: |
| 173 | return _run_module_code(code, init_globals, run_name, |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 174 | fname, loader, pkg_name) |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 175 | else: |
| 176 | # Leave the sys module alone |
Christian Heimes | cbf3b5c | 2007-12-03 21:02:03 +0000 | [diff] [blame] | 177 | return _run_code(code, {}, init_globals, run_name, |
| 178 | fname, loader, pkg_name) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 179 | |
Benjamin Peterson | f6489f9 | 2009-11-25 17:46:26 +0000 | [diff] [blame] | 180 | def _get_main_module_details(): |
| 181 | # Helper that gives a nicer error message when attempting to |
| 182 | # execute a zipfile or directory by invoking __main__.py |
| 183 | main_name = "__main__" |
| 184 | try: |
| 185 | return _get_module_details(main_name) |
| 186 | except ImportError as exc: |
| 187 | if main_name in str(exc): |
| 188 | raise ImportError("can't find %r module in %r" % |
| 189 | (main_name, sys.path[0])) |
| 190 | raise |
| 191 | |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 192 | |
Nick Coghlan | 260bd3e | 2009-11-16 06:49:25 +0000 | [diff] [blame] | 193 | # XXX (ncoghlan): Perhaps expose the C API function |
| 194 | # as imp.get_importer instead of reimplementing it in Python? |
| 195 | def _get_importer(path_name): |
| 196 | """Python version of PyImport_GetImporter C API function""" |
| 197 | cache = sys.path_importer_cache |
| 198 | try: |
| 199 | importer = cache[path_name] |
| 200 | except KeyError: |
| 201 | # Not yet cached. Flag as using the |
| 202 | # standard machinery until we finish |
| 203 | # checking the hooks |
| 204 | cache[path_name] = None |
| 205 | for hook in sys.path_hooks: |
| 206 | try: |
| 207 | importer = hook(path_name) |
| 208 | break |
| 209 | except ImportError: |
| 210 | pass |
| 211 | else: |
| 212 | # The following check looks a bit odd. The trick is that |
| 213 | # NullImporter throws ImportError if the supplied path is a |
| 214 | # *valid* directory entry (and hence able to be handled |
| 215 | # by the standard import machinery) |
| 216 | try: |
| 217 | importer = imp.NullImporter(path_name) |
| 218 | except ImportError: |
| 219 | return None |
| 220 | cache[path_name] = importer |
| 221 | return importer |
| 222 | |
| 223 | def _get_code_from_file(fname): |
| 224 | # Check for a compiled file first |
| 225 | with open(fname, "rb") as f: |
| 226 | code = read_code(f) |
| 227 | if code is None: |
| 228 | # That didn't work, so try it as normal source code |
| 229 | with open(fname, "rU") as f: |
| 230 | code = compile(f.read(), fname, 'exec') |
| 231 | return code |
| 232 | |
| 233 | def run_path(path_name, init_globals=None, run_name=None): |
| 234 | """Execute code located at the specified filesystem location |
| 235 | |
| 236 | Returns the resulting top level namespace dictionary |
| 237 | |
| 238 | The file path may refer directly to a Python script (i.e. |
| 239 | one that could be directly executed with execfile) or else |
| 240 | it may refer to a zipfile or directory containing a top |
| 241 | level __main__.py script. |
| 242 | """ |
| 243 | if run_name is None: |
| 244 | run_name = "<run_path>" |
| 245 | importer = _get_importer(path_name) |
| 246 | if isinstance(importer, imp.NullImporter): |
| 247 | # Not a valid sys.path entry, so run the code directly |
| 248 | # execfile() doesn't help as we want to allow compiled files |
| 249 | code = _get_code_from_file(path_name) |
| 250 | return _run_module_code(code, init_globals, run_name, path_name) |
| 251 | else: |
| 252 | # Importer is defined for path, so add it to |
| 253 | # the start of sys.path |
| 254 | sys.path.insert(0, path_name) |
| 255 | try: |
| 256 | # Here's where things are a little different from the run_module |
| 257 | # case. There, we only had to replace the module in sys while the |
| 258 | # code was running and doing so was somewhat optional. Here, we |
| 259 | # have no choice and we have to remove it even while we read the |
| 260 | # code. If we don't do this, a __loader__ attribute in the |
| 261 | # existing __main__ module may prevent location of the new module. |
| 262 | main_name = "__main__" |
| 263 | saved_main = sys.modules[main_name] |
| 264 | del sys.modules[main_name] |
| 265 | try: |
| 266 | mod_name, loader, code, fname = _get_main_module_details() |
| 267 | finally: |
| 268 | sys.modules[main_name] = saved_main |
| 269 | pkg_name = "" |
| 270 | with _TempModule(run_name) as temp_module, \ |
| 271 | _ModifiedArgv0(path_name): |
| 272 | mod_globals = temp_module.module.__dict__ |
| 273 | return _run_code(code, mod_globals, init_globals, |
| 274 | run_name, fname, loader, pkg_name) |
| 275 | finally: |
| 276 | try: |
| 277 | sys.path.remove(path_name) |
| 278 | except ValueError: |
| 279 | pass |
| 280 | |
| 281 | |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 282 | if __name__ == "__main__": |
| 283 | # Run the module specified as the next command line argument |
| 284 | if len(sys.argv) < 2: |
Guido van Rossum | be19ed7 | 2007-02-09 05:37:30 +0000 | [diff] [blame] | 285 | print("No module specified for execution", file=sys.stderr) |
Thomas Wouters | a977329 | 2006-04-21 09:43:23 +0000 | [diff] [blame] | 286 | else: |
| 287 | del sys.argv[0] # Make the requested module sys.argv[0] |
Thomas Wouters | ed03b41 | 2007-08-28 21:37:11 +0000 | [diff] [blame] | 288 | _run_module_as_main(sys.argv[0]) |