runcmd.py

import os
import subprocess
import select
from termcolor import colored, cprint
import time


# from http://stackoverflow.com/questions/7729336/how-can-i-print-and-display-subprocess-stdout-and-stderr-output-without-distorti
def make_async(fd):
    '''Helper function to add the O_NONBLOCK flag to a file descriptor'''
    fcntl.fcntl(fd, fcntl.F_SETFL, fcntl.fcntl(fd, fcntl.F_GETFL) | os.O_NONBLOCK)

def read_async(fd):
    '''Helper function to read some data from a file descriptor, ignoring EAGAIN errors'''
    try:
        s= fd.read()
        if s is not None:
            return s.decode('utf-8', errors='replace')
        else:
            return ''
    except IOError as e:
        if e.errno != errno.EAGAIN:
            raise e
        else:
            return ''


def runcmd(cmd, workdir=None, log= False, disp=True, wait=True):
    '''Run command (cmd) and capture both stdout and stderr to a log.

    If wait=False, the command will be launched and then the Python script will continue executing.
    '''
    odir = os.path.realpath('.')    
    workdir = workdir or '.'
    cmdstr= ''
    for c in cmd:
        cmdstr += c + ' '
    if disp:
        cprint(' running: "' + cmdstr + '"', 'blue' )
    try:
        os.chdir(workdir)
        
        tool= subprocess.Popen(cmd, stderr= subprocess.PIPE,
                           stdout= subprocess.PIPE)

        make_async(tool.stdout)
        make_async(tool.stderr)
        stdout = str()
        stderr = str()
        returnCode = None

        if not wait:
            return tool
    
        while returnCode is None:
            # Wait for data to become available 
            select.select([tool.stdout, tool.stderr], [], [])

            # Try reading some data from each
            stdoutPiece = read_async(tool.stdout)
            stderrPiece = read_async(tool.stderr)
    
            if disp:
                if stdoutPiece:
                    print(stdoutPiece, end="")
                if stderrPiece:
                    cprint(stderrPiece, 'red', end="")
    
            stdout += stdoutPiece
            stderr += stderrPiece
            returnCode = tool.poll()
        if disp:
            cprint('{} returned {}'.format(cmd[0], tool.returncode), 'blue')
        if tool.returncode == 0:
            if disp:
                cprint("{} sucessful!".format(cmd[0]), 'blue')            
        else:
            raise ExternalProgramError("Problem running {}!!\n return code: {}".format(cmd, tool.returncode))
    finally:
        os.chdir(odir)
    if log:
        return returnCode, stdout, stderr
    return returnCode