Home | History | Annotate | Download | only in Lib
      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):
    101     loader = get_loader(mod_name)
    102     if loader is None:
    103         raise ImportError("No module named %s" % mod_name)
    104     if loader.is_package(mod_name):
    105         if mod_name == "__main__" or mod_name.endswith(".__main__"):
    106             raise ImportError("Cannot use package as __main__ module")
    107         try:
    108             pkg_main_name = mod_name + ".__main__"
    109             return _get_module_details(pkg_main_name)
    110         except ImportError, e:
    111             raise ImportError(("%s; %r is a package and cannot " +
    112                                "be directly executed") %(e, mod_name))
    113     code = loader.get_code(mod_name)
    114     if code is None:
    115         raise ImportError("No code object available for %s" % mod_name)
    116     filename = _get_filename(loader, mod_name)
    117     return mod_name, loader, code, filename
    118 
    119 
    120 def _get_main_module_details():
    121     # Helper that gives a nicer error message when attempting to

    122     # execute a zipfile or directory by invoking __main__.py

    123     main_name = "__main__"
    124     try:
    125         return _get_module_details(main_name)
    126     except ImportError as exc:
    127         if main_name in str(exc):
    128             raise ImportError("can't find %r module in %r" %
    129                               (main_name, sys.path[0]))
    130         raise
    131 
    132 # This function is the actual implementation of the -m switch and direct

    133 # execution of zipfiles and directories and is deliberately kept private.

    134 # This avoids a repeat of the situation where run_module() no longer met the

    135 # needs of mainmodule.c, but couldn't be changed because it was public

    136 def _run_module_as_main(mod_name, alter_argv=True):
    137     """Runs the designated module in the __main__ namespace
    138 
    139        Note that the executed module will have full access to the
    140        __main__ namespace. If this is not desirable, the run_module()
    141        function should be used to run the module code in a fresh namespace.
    142 
    143        At the very least, these variables in __main__ will be overwritten:
    144            __name__
    145            __file__
    146            __loader__
    147            __package__
    148     """
    149     try:
    150         if alter_argv or mod_name != "__main__": # i.e. -m switch

    151             mod_name, loader, code, fname = _get_module_details(mod_name)
    152         else:          # i.e. directory or zipfile execution

    153             mod_name, loader, code, fname = _get_main_module_details()
    154     except ImportError as exc:
    155         msg = "%s: %s" % (sys.executable, str(exc))
    156         sys.exit(msg)
    157     pkg_name = mod_name.rpartition('.')[0]
    158     main_globals = sys.modules["__main__"].__dict__
    159     if alter_argv:
    160         sys.argv[0] = fname
    161     return _run_code(code, main_globals, None,
    162                      "__main__", fname, loader, pkg_name)
    163 
    164 def run_module(mod_name, init_globals=None,
    165                run_name=None, alter_sys=False):
    166     """Execute a module's code without importing it
    167 
    168        Returns the resulting top level namespace dictionary
    169     """
    170     mod_name, loader, code, fname = _get_module_details(mod_name)
    171     if run_name is None:
    172         run_name = mod_name
    173     pkg_name = mod_name.rpartition('.')[0]
    174     if alter_sys:
    175         return _run_module_code(code, init_globals, run_name,
    176                                 fname, loader, pkg_name)
    177     else:
    178         # Leave the sys module alone

    179         return _run_code(code, {}, init_globals, run_name,
    180                          fname, loader, pkg_name)
    181 
    182 
    183 # XXX (ncoghlan): Perhaps expose the C API function

    184 # as imp.get_importer instead of reimplementing it in Python?

    185 def _get_importer(path_name):
    186     """Python version of PyImport_GetImporter C API function"""
    187     cache = sys.path_importer_cache
    188     try:
    189         importer = cache[path_name]
    190     except KeyError:
    191         # Not yet cached. Flag as using the

    192         # standard machinery until we finish

    193         # checking the hooks

    194         cache[path_name] = None
    195         for hook in sys.path_hooks:
    196             try:
    197                 importer = hook(path_name)
    198                 break
    199             except ImportError:
    200                 pass
    201         else:
    202             # The following check looks a bit odd. The trick is that

    203             # NullImporter throws ImportError if the supplied path is a

    204             # *valid* directory entry (and hence able to be handled

    205             # by the standard import machinery)

    206             try:
    207                 importer = imp.NullImporter(path_name)
    208             except ImportError:
    209                 return None
    210         cache[path_name] = importer
    211     return importer
    212 
    213 def _get_code_from_file(fname):
    214     # Check for a compiled file first

    215     with open(fname, "rb") as f:
    216         code = read_code(f)
    217     if code is None:
    218         # That didn't work, so try it as normal source code

    219         with open(fname, "rU") as f:
    220             code = compile(f.read(), fname, 'exec')
    221     return code
    222 
    223 def run_path(path_name, init_globals=None, run_name=None):
    224     """Execute code located at the specified filesystem location
    225 
    226        Returns the resulting top level namespace dictionary
    227 
    228        The file path may refer directly to a Python script (i.e.
    229        one that could be directly executed with execfile) or else
    230        it may refer to a zipfile or directory containing a top
    231        level __main__.py script.
    232     """
    233     if run_name is None:
    234         run_name = "<run_path>"
    235     importer = _get_importer(path_name)
    236     if isinstance(importer, imp.NullImporter):
    237         # Not a valid sys.path entry, so run the code directly

    238         # execfile() doesn't help as we want to allow compiled files

    239         code = _get_code_from_file(path_name)
    240         return _run_module_code(code, init_globals, run_name, path_name)
    241     else:
    242         # Importer is defined for path, so add it to

    243         # the start of sys.path

    244         sys.path.insert(0, path_name)
    245         try:
    246             # Here's where things are a little different from the run_module

    247             # case. There, we only had to replace the module in sys while the

    248             # code was running and doing so was somewhat optional. Here, we

    249             # have no choice and we have to remove it even while we read the

    250             # code. If we don't do this, a __loader__ attribute in the

    251             # existing __main__ module may prevent location of the new module.

    252             main_name = "__main__"
    253             saved_main = sys.modules[main_name]
    254             del sys.modules[main_name]
    255             try:
    256                 mod_name, loader, code, fname = _get_main_module_details()
    257             finally:
    258                 sys.modules[main_name] = saved_main
    259             pkg_name = ""
    260             with _TempModule(run_name) as temp_module, \
    261                  _ModifiedArgv0(path_name):
    262                 mod_globals = temp_module.module.__dict__
    263                 return _run_code(code, mod_globals, init_globals,
    264                                     run_name, fname, loader, pkg_name).copy()
    265         finally:
    266             try:
    267                 sys.path.remove(path_name)
    268             except ValueError:
    269                 pass
    270 
    271 
    272 if __name__ == "__main__":
    273     # Run the module specified as the next command line argument

    274     if len(sys.argv) < 2:
    275         print >> sys.stderr, "No module specified for execution"
    276     else:
    277         del sys.argv[0] # Make the requested module sys.argv[0]

    278         _run_module_as_main(sys.argv[0])
    279