runpy.py 11 KB

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