Back to index

python-cliapp  1.20120630
runcmd.py
Go to the documentation of this file.
00001 # Copyright (C) 2011  Lars Wirzenius
00002 # 
00003 # This program is free software; you can redistribute it and/or modify
00004 # it under the terms of the GNU General Public License as published by
00005 # the Free Software Foundation; either version 2 of the License, or
00006 # (at your option) any later version.
00007 # 
00008 # This program is distributed in the hope that it will be useful,
00009 # but WITHOUT ANY WARRANTY; without even the implied warranty of
00010 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
00011 # GNU General Public License for more details.
00012 # 
00013 # You should have received a copy of the GNU General Public License along
00014 # with this program; if not, write to the Free Software Foundation, Inc.,
00015 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
00016 
00017 
00018 import fcntl
00019 import logging
00020 import os
00021 import select
00022 import subprocess
00023 
00024 import cliapp
00025 
00026 
00027 def runcmd(argv, *args, **kwargs):
00028     '''Run external command or pipeline.
00029 
00030     Example: ``runcmd(['grep', 'foo'], ['wc', '-l'], 
00031                       feed_stdin='foo\nbar\n')``
00032 
00033     Return the standard output of the command.
00034     
00035     Raise ``cliapp.AppException`` if external command returns
00036     non-zero exit code. ``*args`` and ``**kwargs`` are passed
00037     onto ``subprocess.Popen``.
00038     
00039     '''
00040 
00041     if 'ignore_fail' in kwargs:
00042         ignore_fail = kwargs['ignore_fail']
00043         del kwargs['ignore_fail']
00044     else:
00045         ignore_fail = False
00046 
00047     exit, out, err = runcmd_unchecked(argv, *args, **kwargs)
00048     if exit != 0:
00049         msg = 'Command failed: %s\n%s' % (' '.join(argv), err)
00050         if ignore_fail:
00051             logging.info(msg)
00052         else:
00053             logging.error(msg)
00054             raise cliapp.AppException(msg)
00055     return out
00056     
00057 def runcmd_unchecked(argv, *argvs, **kwargs):
00058     '''Run external command or pipeline.
00059 
00060     Return the exit code, and contents of standard output and error
00061     of the command.
00062     
00063     See also ``runcmd``.
00064     
00065     '''
00066 
00067     argvs = [argv] + list(argvs)
00068     logging.debug('run external command: %s' % repr(argvs))
00069 
00070     def pop_kwarg(name, default):
00071         if name in kwargs:
00072             value = kwargs[name]
00073             del kwargs[name]
00074             return value
00075         else:
00076             return default
00077 
00078     feed_stdin = pop_kwarg('feed_stdin', '')
00079     pipe_stdin = pop_kwarg('stdin', subprocess.PIPE)
00080     pipe_stdout = pop_kwarg('stdout', subprocess.PIPE)
00081     pipe_stderr = pop_kwarg('stderr', subprocess.PIPE)
00082 
00083     try:
00084         pipeline = _build_pipeline(argvs,
00085                                    pipe_stdin,
00086                                    pipe_stdout,
00087                                    pipe_stderr,
00088                                    kwargs)
00089         return _run_pipeline(pipeline, feed_stdin, pipe_stdin,
00090                               pipe_stdout, pipe_stderr)
00091     except OSError, e: # pragma: no cover
00092         if e.errno == errno.ENOENT and e.filename is None:
00093             e.filename = argv[0]
00094             raise e
00095         else:
00096             raise
00097 
00098 def _build_pipeline(argvs, pipe_stdin, pipe_stdout, pipe_stderr, kwargs):
00099     procs = []
00100     for i, argv in enumerate(argvs):
00101         if i == 0 and i == len(argvs) - 1:
00102             stdin = pipe_stdin
00103             stdout = pipe_stdout
00104             stderr = pipe_stderr
00105         elif i == 0:
00106             stdin = pipe_stdin
00107             stdout = subprocess.PIPE
00108             stderr = pipe_stderr
00109         elif i == len(argv) - 1:
00110             stdin = procs[-1].stdout
00111             stdout = pipe_stdout
00112             stderr = pipe_stderr
00113         else:
00114             stdin = procs[-1].stdout
00115             stdout = subprocess.PIPE
00116             stderr = pipe_stderr
00117         p = subprocess.Popen(argv, stdin=stdin, stdout=stdout, 
00118                              stderr=stderr, close_fds=True, **kwargs)
00119         procs.append(p)
00120 
00121     return procs
00122 
00123 def _run_pipeline(procs, feed_stdin, pipe_stdin, pipe_stdout, pipe_stderr):
00124     stdout_eof = False
00125     stderr_eof = False
00126     out = []
00127     err = []
00128     pos = 0
00129     io_size = 1024
00130     
00131     def set_nonblocking(fd):
00132         flags = fcntl.fcntl(fd, fcntl.F_GETFL, 0)
00133         flags = flags | os.O_NONBLOCK
00134         fcntl.fcntl(fd, fcntl.F_SETFL, flags)
00135 
00136     if feed_stdin and pipe_stdin == subprocess.PIPE:
00137         set_nonblocking(procs[0].stdin.fileno())
00138     if pipe_stdout == subprocess.PIPE:
00139         set_nonblocking(procs[-1].stdout.fileno())
00140     if pipe_stderr == subprocess.PIPE:
00141         set_nonblocking(procs[-1].stderr.fileno())
00142     
00143     def still_running():
00144         for p in procs:
00145             p.poll()
00146         for p in procs:
00147             if p.returncode is None:
00148                 return True
00149         if pipe_stdout == subprocess.PIPE and not stdout_eof:
00150             return True
00151         if pipe_stderr == subprocess.PIPE and not stderr_eof:
00152             return True # pragma: no cover
00153         return False
00154 
00155     while still_running():
00156         rlist = []
00157         if not stdout_eof and pipe_stdout == subprocess.PIPE:
00158             rlist.append(procs[-1].stdout)
00159         if not stderr_eof and pipe_stderr == subprocess.PIPE:
00160             rlist.append(procs[-1].stderr)
00161         
00162         wlist = []
00163         if pipe_stdin == subprocess.PIPE and pos < len(feed_stdin):
00164             wlist.append(procs[0].stdin)
00165 
00166         if rlist or wlist:
00167             r, w, x = select.select(rlist, wlist, [])
00168         else:
00169             r = w = [] # pragma: no cover
00170 
00171         if procs[0].stdin in w and pos < len(feed_stdin):
00172             data = feed_stdin[pos : pos+io_size]
00173             procs[0].stdin.write(data)
00174             pos += len(data)
00175             if pos >= len(feed_stdin):
00176                 procs[0].stdin.close()
00177 
00178         if procs[-1].stdout in r:
00179             data = procs[-1].stdout.read(io_size)
00180             if data:
00181                 out.append(data)
00182             else:
00183                 stdout_eof = True
00184 
00185         if procs[-1].stderr in r:
00186             data = procs[-1].stderr.read(io_size)
00187             if data:
00188                 err.append(data)
00189             else:
00190                 stderr_eof = True
00191 
00192     return procs[-1].returncode, ''.join(out), ''.join(err)
00193