| 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
 | 
| 14 | from pkgutil import read_code
 | 
| 15 | try:
 | 
| 16 |     from imp import get_loader
 | 
| 17 | except ImportError:
 | 
| 18 |     from pkgutil import get_loader
 | 
| 19 | 
 | 
| 20 | __all__ = [
 | 
| 21 |     "run_module", "run_path",
 | 
| 22 | ]
 | 
| 23 | 
 | 
| 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
 | 
| 61 | 
 | 
| 62 | def _run_code(code, run_globals, init_globals=None,
 | 
| 63 |               mod_name=None, mod_fname=None,
 | 
| 64 |               mod_loader=None, pkg_name=None):
 | 
| 65 |     """Helper to run code in nominated namespace"""
 | 
| 66 |     if init_globals is not None:
 | 
| 67 |         run_globals.update(init_globals)
 | 
| 68 |     run_globals.update(__name__ = mod_name,
 | 
| 69 |                        __file__ = mod_fname,
 | 
| 70 |                        __loader__ = mod_loader,
 | 
| 71 |                        __package__ = pkg_name)
 | 
| 72 |     exec code in run_globals
 | 
| 73 |     return run_globals
 | 
| 74 | 
 | 
| 75 | def _run_module_code(code, init_globals=None,
 | 
| 76 |                     mod_name=None, mod_fname=None,
 | 
| 77 |                     mod_loader=None, pkg_name=None):
 | 
| 78 |     """Helper to run code in new namespace with sys modified"""
 | 
| 79 |     with _TempModule(mod_name) as temp_module, _ModifiedArgv0(mod_fname):
 | 
| 80 |         mod_globals = temp_module.module.__dict__
 | 
| 81 |         _run_code(code, mod_globals, init_globals,
 | 
| 82 |                   mod_name, mod_fname, mod_loader, pkg_name)
 | 
| 83 |     # Copy the globals of the temporary module, as they
 | 
| 84 |     # may be cleared when the temporary module goes away
 | 
| 85 |     return mod_globals.copy()
 | 
| 86 | 
 | 
| 87 | 
 | 
| 88 | # This helper is needed due to a missing component in the PEP 302
 | 
| 89 | # loader protocol (specifically, "get_filename" is non-standard)
 | 
| 90 | # Since we can't introduce new features in maintenance releases,
 | 
| 91 | # support was added to zipimporter under the name '_get_filename'
 | 
| 92 | def _get_filename(loader, mod_name):
 | 
| 93 |     for attr in ("get_filename", "_get_filename"):
 | 
| 94 |         meth = getattr(loader, attr, None)
 | 
| 95 |         if meth is not None:
 | 
| 96 |             return meth(mod_name)
 | 
| 97 |     return None
 | 
| 98 | 
 | 
| 99 | # Helper to get the loader, code and filename for a module
 | 
| 100 | def _get_module_details(mod_name, error=ImportError):
 | 
| 101 |     try:
 | 
| 102 |         loader = get_loader(mod_name)
 | 
| 103 |         if loader is None:
 | 
| 104 |             raise error("No module named %s" % mod_name)
 | 
| 105 |         ispkg = loader.is_package(mod_name)
 | 
| 106 |     except ImportError as e:
 | 
| 107 |         raise error(format(e))
 | 
| 108 |     if ispkg:
 | 
| 109 |         if mod_name == "__main__" or mod_name.endswith(".__main__"):
 | 
| 110 |             raise error("Cannot use package as __main__ module")
 | 
| 111 |         __import__(mod_name)  # Do not catch exceptions initializing package
 | 
| 112 |         try:
 | 
| 113 |             pkg_main_name = mod_name + ".__main__"
 | 
| 114 |             return _get_module_details(pkg_main_name)
 | 
| 115 |         except ImportError, e:
 | 
| 116 |             raise error(("%s; %r is a package and cannot " +
 | 
| 117 |                                "be directly executed") %(e, mod_name))
 | 
| 118 |     try:
 | 
| 119 |         code = loader.get_code(mod_name)
 | 
| 120 |     except ImportError as e:
 | 
| 121 |         raise error(format(e))
 | 
| 122 |     if code is None:
 | 
| 123 |         raise error("No code object available for %s" % mod_name)
 | 
| 124 |     filename = _get_filename(loader, mod_name)
 | 
| 125 |     return mod_name, loader, code, filename
 | 
| 126 | 
 | 
| 127 | 
 | 
| 128 | def _get_main_module_details(error=ImportError):
 | 
| 129 |     # Helper that gives a nicer error message when attempting to
 | 
| 130 |     # execute a zipfile or directory by invoking __main__.py
 | 
| 131 |     main_name = "__main__"
 | 
| 132 |     try:
 | 
| 133 |         return _get_module_details(main_name)
 | 
| 134 |     except ImportError as exc:
 | 
| 135 |         if main_name in str(exc):
 | 
| 136 |             raise error("can't find %r module in %r" %
 | 
| 137 |                               (main_name, sys.path[0]))
 | 
| 138 |         raise
 | 
| 139 | 
 | 
| 140 | class _Error(Exception):
 | 
| 141 |     """Error that _run_module_as_main() should report without a traceback"""
 | 
| 142 | 
 | 
| 143 | # This function is the actual implementation of the -m switch and direct
 | 
| 144 | # execution of zipfiles and directories and is deliberately kept private.
 | 
| 145 | # This avoids a repeat of the situation where run_module() no longer met the
 | 
| 146 | # needs of mainmodule.c, but couldn't be changed because it was public
 | 
| 147 | def _run_module_as_main(mod_name, alter_argv=True):
 | 
| 148 |     """Runs the designated module in the __main__ namespace
 | 
| 149 | 
 | 
| 150 |        Note that the executed module will have full access to the
 | 
| 151 |        __main__ namespace. If this is not desirable, the run_module()
 | 
| 152 |        function should be used to run the module code in a fresh namespace.
 | 
| 153 | 
 | 
| 154 |        At the very least, these variables in __main__ will be overwritten:
 | 
| 155 |            __name__
 | 
| 156 |            __file__
 | 
| 157 |            __loader__
 | 
| 158 |            __package__
 | 
| 159 |     """
 | 
| 160 |     try:
 | 
| 161 |         if alter_argv or mod_name != "__main__": # i.e. -m switch
 | 
| 162 |             mod_name, loader, code, fname = _get_module_details(
 | 
| 163 |                 mod_name, _Error)
 | 
| 164 |         else:          # i.e. directory or zipfile execution
 | 
| 165 |             mod_name, loader, code, fname = _get_main_module_details(_Error)
 | 
| 166 |     except _Error as exc:
 | 
| 167 |         msg = "%s: %s" % (sys.executable, exc)
 | 
| 168 |         sys.exit(msg)
 | 
| 169 |     pkg_name = mod_name.rpartition('.')[0]
 | 
| 170 |     main_globals = sys.modules["__main__"].__dict__
 | 
| 171 |     if alter_argv:
 | 
| 172 |         sys.argv[0] = fname
 | 
| 173 |     return _run_code(code, main_globals, None,
 | 
| 174 |                      "__main__", fname, loader, pkg_name)
 | 
| 175 | 
 | 
| 176 | def run_module(mod_name, init_globals=None,
 | 
| 177 |                run_name=None, alter_sys=False):
 | 
| 178 |     """Execute a module's code without importing it
 | 
| 179 | 
 | 
| 180 |        Returns the resulting top level namespace dictionary
 | 
| 181 |     """
 | 
| 182 |     mod_name, loader, code, fname = _get_module_details(mod_name)
 | 
| 183 |     if run_name is None:
 | 
| 184 |         run_name = mod_name
 | 
| 185 |     pkg_name = mod_name.rpartition('.')[0]
 | 
| 186 |     if alter_sys:
 | 
| 187 |         return _run_module_code(code, init_globals, run_name,
 | 
| 188 |                                 fname, loader, pkg_name)
 | 
| 189 |     else:
 | 
| 190 |         # Leave the sys module alone
 | 
| 191 |         return _run_code(code, {}, init_globals, run_name,
 | 
| 192 |                          fname, loader, pkg_name)
 | 
| 193 | 
 | 
| 194 | 
 | 
| 195 | # XXX (ncoghlan): Perhaps expose the C API function
 | 
| 196 | # as imp.get_importer instead of reimplementing it in Python?
 | 
| 197 | def _get_importer(path_name):
 | 
| 198 |     """Python version of PyImport_GetImporter C API function"""
 | 
| 199 |     cache = sys.path_importer_cache
 | 
| 200 |     try:
 | 
| 201 |         importer = cache[path_name]
 | 
| 202 |     except KeyError:
 | 
| 203 |         # Not yet cached. Flag as using the
 | 
| 204 |         # standard machinery until we finish
 | 
| 205 |         # checking the hooks
 | 
| 206 |         cache[path_name] = None
 | 
| 207 |         for hook in sys.path_hooks:
 | 
| 208 |             try:
 | 
| 209 |                 importer = hook(path_name)
 | 
| 210 |                 break
 | 
| 211 |             except ImportError:
 | 
| 212 |                 pass
 | 
| 213 |         else:
 | 
| 214 |             # The following check looks a bit odd. The trick is that
 | 
| 215 |             # NullImporter raises ImportError if the supplied path is a
 | 
| 216 |             # *valid* directory entry (and hence able to be handled
 | 
| 217 |             # by the standard import machinery)
 | 
| 218 |             try:
 | 
| 219 |                 importer = imp.NullImporter(path_name)
 | 
| 220 |             except ImportError:
 | 
| 221 |                 return None
 | 
| 222 |         cache[path_name] = importer
 | 
| 223 |     return importer
 | 
| 224 | 
 | 
| 225 | def _get_code_from_file(fname):
 | 
| 226 |     # Check for a compiled file first
 | 
| 227 |     with open(fname, "rb") as f:
 | 
| 228 |         code = read_code(f)
 | 
| 229 |     if code is None:
 | 
| 230 |         # That didn't work, so try it as normal source code
 | 
| 231 |         with open(fname, "rU") as f:
 | 
| 232 |             code = compile(f.read(), fname, 'exec')
 | 
| 233 |     return code
 | 
| 234 | 
 | 
| 235 | def run_path(path_name, init_globals=None, run_name=None):
 | 
| 236 |     """Execute code located at the specified filesystem location
 | 
| 237 | 
 | 
| 238 |        Returns the resulting top level namespace dictionary
 | 
| 239 | 
 | 
| 240 |        The file path may refer directly to a Python script (i.e.
 | 
| 241 |        one that could be directly executed with execfile) or else
 | 
| 242 |        it may refer to a zipfile or directory containing a top
 | 
| 243 |        level __main__.py script.
 | 
| 244 |     """
 | 
| 245 |     if run_name is None:
 | 
| 246 |         run_name = "<run_path>"
 | 
| 247 |     importer = _get_importer(path_name)
 | 
| 248 |     if isinstance(importer, imp.NullImporter):
 | 
| 249 |         # Not a valid sys.path entry, so run the code directly
 | 
| 250 |         # execfile() doesn't help as we want to allow compiled files
 | 
| 251 |         code = _get_code_from_file(path_name)
 | 
| 252 |         return _run_module_code(code, init_globals, run_name, path_name)
 | 
| 253 |     else:
 | 
| 254 |         # Importer is defined for path, so add it to
 | 
| 255 |         # the start of sys.path
 | 
| 256 |         sys.path.insert(0, path_name)
 | 
| 257 |         try:
 | 
| 258 |             # Here's where things are a little different from the run_module
 | 
| 259 |             # case. There, we only had to replace the module in sys while the
 | 
| 260 |             # code was running and doing so was somewhat optional. Here, we
 | 
| 261 |             # have no choice and we have to remove it even while we read the
 | 
| 262 |             # code. If we don't do this, a __loader__ attribute in the
 | 
| 263 |             # existing __main__ module may prevent location of the new module.
 | 
| 264 |             main_name = "__main__"
 | 
| 265 |             saved_main = sys.modules[main_name]
 | 
| 266 |             del sys.modules[main_name]
 | 
| 267 |             try:
 | 
| 268 |                 mod_name, loader, code, fname = _get_main_module_details()
 | 
| 269 |             finally:
 | 
| 270 |                 sys.modules[main_name] = saved_main
 | 
| 271 |             pkg_name = ""
 | 
| 272 |             with _TempModule(run_name) as temp_module, \
 | 
| 273 |                  _ModifiedArgv0(path_name):
 | 
| 274 |                 mod_globals = temp_module.module.__dict__
 | 
| 275 |                 return _run_code(code, mod_globals, init_globals,
 | 
| 276 |                                     run_name, fname, loader, pkg_name).copy()
 | 
| 277 |         finally:
 | 
| 278 |             try:
 | 
| 279 |                 sys.path.remove(path_name)
 | 
| 280 |             except ValueError:
 | 
| 281 |                 pass
 | 
| 282 | 
 | 
| 283 | 
 | 
| 284 | if __name__ == "__main__":
 | 
| 285 |     # Run the module specified as the next command line argument
 | 
| 286 |     if len(sys.argv) < 2:
 | 
| 287 |         print >> sys.stderr, "No module specified for execution"
 | 
| 288 |     else:
 | 
| 289 |         del sys.argv[0] # Make the requested module sys.argv[0]
 | 
| 290 |         _run_module_as_main(sys.argv[0])
 |