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):
  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,
  32                        __loader__ = mod_loader)
  33     exec code in run_globals
  34     return run_globals
  35
  36 def _run_module_code(code, init_globals=None,
  37                     mod_name=None, mod_fname=None,
  38                     mod_loader=None):
  39     """Helper for run_module"""
  40     # Set up the top level namespace dictionary
  41     temp_module = imp.new_module(mod_name)
  42     mod_globals = temp_module.__dict__
  43     # Modify sys.argv[0] and sys.module[mod_name]
  44     saved_argv0 = sys.argv[0]
  45     restore_module = mod_name in sys.modules
  46     if restore_module:
  47         saved_module = sys.modules[mod_name]
  48     sys.argv[0] = mod_fname
  49     sys.modules[mod_name] = temp_module
  50     try:
  51         _run_code(code, mod_globals, init_globals,
  52                     mod_name, mod_fname, mod_loader)
  53     finally:
  54         sys.argv[0] = saved_argv0
  55         if restore_module:
  56             sys.modules[mod_name] = saved_module
  57         else:
  58             del sys.modules[mod_name]
  59     # Copy the globals of the temporary module, as they
  60     # may be cleared when the temporary module goes away
  61     return mod_globals.copy()
  62
  63
  64 # This helper is needed due to a missing component in the PEP 302
  65 # loader protocol (specifically, "get_filename" is non-standard)
  66 def _get_filename(loader, mod_name):
  67     try:
  68         get_filename = loader.get_filename
  69     except AttributeError:
  70         return None
  71     else:
  72         return get_filename(mod_name)
  73
  74 # Helper to get the loader, code and filename for a module
  75 def _get_module_details(mod_name):
  76     loader = get_loader(mod_name)
  77     if loader is None:
  78         raise ImportError("No module named %s" % mod_name)
  79     if loader.is_package(mod_name):
  80         raise ImportError(("%s is a package and cannot " +
  81                           "be directly executed") % mod_name)
  82     code = loader.get_code(mod_name)
  83     if code is None:
  84         raise ImportError("No code object available for %s" % mod_name)
  85     filename = _get_filename(loader, mod_name)
  86     return loader, code, filename
  87
  88
  89 # XXX ncoghlan: Should this be documented and made public?
  90 def _run_module_as_main(mod_name, set_argv0=True):
  91     """Runs the designated module in the __main__ namespace
  92
  93        These __*__ magic variables will be overwritten:
  94            __file__
  95            __loader__
  96     """
  97     loader, code, fname = _get_module_details(mod_name)
  98     main_globals = sys.modules["__main__"].__dict__
  99     if set_argv0:
 100         sys.argv[0] = fname
 101     return _run_code(code, main_globals, None,
 102                      "__main__", fname, loader)
 103
 104 def run_module(mod_name, init_globals=None,
 105                run_name=None, alter_sys=False):
 106     """Execute a module's code without importing it
 107
 108        Returns the resulting top level namespace dictionary
 109     """
 110     loader, code, fname = _get_module_details(mod_name)
 111     if run_name is None:
 112         run_name = mod_name
 113     if alter_sys:
 114         return _run_module_code(code, init_globals, run_name,
 115                                 fname, loader)
 116     else:
 117         # Leave the sys module alone
 118         return _run_code(code, {}, init_globals,
 119                          run_name, fname, loader)
 120
 121
 122 if __name__ == "__main__":
 123     # Run the module specified as the next command line argument
 124     if len(sys.argv) < 2:
 125         print >> sys.stderr, "No module specified for execution"
 126     else:
 127         del sys.argv[0] # Make the requested module sys.argv[0]
 128         _run_module_as_main(sys.argv[0])