Lib/runpy.py

   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 try:
  15     from imp import get_loader
  16 except ImportError:
  17     from pkgutil import get_loader
  18
  19 __all__ = [
  20     "run_module",
  21 ]
  22
  23
  24 def _run_code(code, run_globals, init_globals=None,
  25               mod_name=None, mod_fname=None,
  26               mod_loader=None, pkg_name=None):
  27     """Helper to run code in nominated namespace"""
  28     if init_globals is not None:
  29         run_globals.update(init_globals)
  30     run_globals.update(__name__ = mod_name,
  31                        __file__ = mod_fname,
  32                        __loader__ = mod_loader,
  33                        __package__ = pkg_name)
  34     exec code in run_globals
  35     return run_globals
  36
  37 def _run_module_code(code, init_globals=None,
  38                     mod_name=None, mod_fname=None,
  39                     mod_loader=None, pkg_name=None):
  40     """Helper to run code in new namespace with sys modified"""
  41     # Set up the top level namespace dictionary
  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,
  53                     mod_name, mod_fname,
  54                     mod_loader, pkg_name)
  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()
  64
  65
  66 # This helper is needed due to a missing component in the PEP 302
  67 # loader protocol (specifically, "get_filename" is non-standard)
  68 # Since we can't introduce new features in maintenance releases,
  69 # support was added to zipimporter under the name '_get_filename'
  70 def _get_filename(loader, mod_name):
  71     for attr in ("get_filename", "_get_filename"):
  72         meth = getattr(loader, attr, None)
  73         if meth is not None:
  74             return meth(mod_name)
  75     return None
  76
  77 # Helper to get the loader, code and filename for a module
  78 def _get_module_details(mod_name):
  79     loader = get_loader(mod_name)
  80     if loader is None:
  81         raise ImportError("No module named %s" % mod_name)
  82     if loader.is_package(mod_name):
  83         if mod_name == "__main__" or mod_name.endswith(".__main__"):
  84             raise ImportError("Cannot use package as __main__ module")
  85         try:
  86             pkg_main_name = mod_name + ".__main__"
  87             return _get_module_details(pkg_main_name)
  88         except ImportError, e:
  89             raise ImportError(("%s; %r is a package and cannot " +
  90                                "be directly executed") %(e, mod_name))
  91     code = loader.get_code(mod_name)
  92     if code is None:
  93         raise ImportError("No code object available for %s" % mod_name)
  94     filename = _get_filename(loader, mod_name)
  95     return mod_name, loader, code, filename
  96
  97
  98 # XXX ncoghlan: Should this be documented and made public?
  99 # (Current thoughts: don't repeat the mistake that lead to its
 100 # creation when run_module() no longer met the needs of
 101 # mainmodule.c, but couldn't be changed because it was public)
 102 def _run_module_as_main(mod_name, alter_argv=True):
 103     """Runs the designated module in the __main__ namespace
 104
 105        Note that the executed module will have full access to the
 106        __main__ namespace. If this is not desirable, the run_module()
 107        function sbould be used to run the module code in a fresh namespace.
 108
 109        At the very least, these variables in __main__ will be overwritten:
 110            __name__
 111            __file__
 112            __loader__
 113            __package__
 114     """
 115     try:
 116         mod_name, loader, code, fname = _get_module_details(mod_name)
 117     except ImportError as exc:
 118         # Try to provide a good error message
 119         # for directories, zip files and the -m switch
 120         if alter_argv:
 121             # For -m switch, just display the exception
 122             info = str(exc)
 123         else:
 124             # For directories/zipfiles, let the user
 125             # know what the code was looking for
 126             info = "can't find '__main__.py' in %r" % sys.argv[0]
 127         msg = "%s: %s" % (sys.executable, info)
 128         sys.exit(msg)
 129     pkg_name = mod_name.rpartition('.')[0]
 130     main_globals = sys.modules["__main__"].__dict__
 131     if alter_argv:
 132         sys.argv[0] = fname
 133     return _run_code(code, main_globals, None,
 134                      "__main__", fname, loader, pkg_name)
 135
 136 def run_module(mod_name, init_globals=None,
 137                run_name=None, alter_sys=False):
 138     """Execute a module's code without importing it
 139
 140        Returns the resulting top level namespace dictionary
 141     """
 142     mod_name, loader, code, fname = _get_module_details(mod_name)
 143     if run_name is None:
 144         run_name = mod_name
 145     pkg_name = mod_name.rpartition('.')[0]
 146     if alter_sys:
 147         return _run_module_code(code, init_globals, run_name,
 148                                 fname, loader, pkg_name)
 149     else:
 150         # Leave the sys module alone
 151         return _run_code(code, {}, init_globals, run_name,
 152                          fname, loader, pkg_name)
 153
 154
 155 if __name__ == "__main__":
 156     # Run the module specified as the next command line argument
 157     if len(sys.argv) < 2:
 158         print >> sys.stderr, "No module specified for execution"
 159     else:
 160         del sys.argv[0] # Make the requested module sys.argv[0]
 161         _run_module_as_main(sys.argv[0])