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,
  25               mod_name, mod_fname, mod_loader):
  26     """Helper for _run_module_code"""
  27     if init_globals is not None:
  28         run_globals.update(init_globals)
  29     run_globals.update(__name__ = mod_name,
  30                        __file__ = mod_fname,
  31                        __loader__ = mod_loader)
  32     exec code in run_globals
  33     return run_globals
  34
  35 def _run_module_code(code, init_globals=None,
  36                     mod_name=None, mod_fname=None,
  37                     mod_loader=None, alter_sys=False):
  38     """Helper for run_module"""
  39     # Set up the top level namespace dictionary
  40     if alter_sys:
  41         # Modify sys.argv[0] and sys.module[mod_name]
  42         temp_module = imp.new_module(mod_name)
  43         mod_globals = temp_module.__dict__
  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     else:
  63         # Leave the sys module alone
  64         return _run_code(code, {}, init_globals,
  65                          mod_name, mod_fname, mod_loader)
  66
  67
  68 # This helper is needed due to a missing component in the PEP 302
  69 # loader protocol (specifically, "get_filename" is non-standard)
  70 def _get_filename(loader, mod_name):
  71     try:
  72         get_filename = loader.get_filename
  73     except AttributeError:
  74         return None
  75     else:
  76         return get_filename(mod_name)
  77
  78
  79 def run_module(mod_name, init_globals=None,
  80                          run_name=None, alter_sys=False):
  81     """Execute a module's code without importing it
  82
  83        Returns the resulting top level namespace dictionary
  84     """
  85     loader = get_loader(mod_name)
  86     if loader is None:
  87         raise ImportError("No module named " + mod_name)
  88     code = loader.get_code(mod_name)
  89     if code is None:
  90         raise ImportError("No code object available for " + mod_name)
  91     filename = _get_filename(loader, mod_name)
  92     if run_name is None:
  93         run_name = mod_name
  94     return _run_module_code(code, init_globals, run_name,
  95                             filename, loader, alter_sys)
  96
  97
  98 if __name__ == "__main__":
  99     # Run the module specified as the next command line argument
 100     if len(sys.argv) < 2:
 101         print >> sys.stderr, "No module specified for execution"
 102     else:
 103         del sys.argv[0] # Make the requested module sys.argv[0]
 104         run_module(sys.argv[0], run_name="__main__", alter_sys=True)