A useful wrapper for Python subprocess

Python comes with a subprocess module that can be used to deal with running subprocesses, collecting output from them, checking exit status, etc. Particularly subprocess.Popen is powerful. It is lowlevel and quite handy.

But I often find myself in the following situation:

  • Run a program and collect its output
  • As the process is running read the output and continuously print it
  • Check the exit status of the process

The second requirement is the trickiest one. I had to do a lot of digging around and come up with a near-perfect solution. The solution below is battle tested. I am sharing it so that someone else does not have to waste time.

#!/usr/bin/env python3

import os
import sys
import subprocess
import signal
import shlex
import shutil
import selectors
import io

def run_cmd(cmd, tee=False):
    """This function will run the given command.  If tee=True, the stdout 
    of the subprocess is read and printed in the stdout of the current process. 
    (like tee program prints to a file and stdout).

    You can specify the command as a string or a list. When you specify the 
    command as a string, it will be properly split into program name and 
    arguments using shlex module.

    When the subprocess is done, exit status and all the output collected
    so far (immaterial of already printed with tee=True or not) is returned.
    """
    if type(cmd) == str:
        args = shlex.split(cmd)
    elif type(cmd) == list:
        args = cmd
    else:
        raise Exception(f"*** Expected args to be a str or list. Found: {cmd}")

    # Replace the first argument with its absolute path.
    if not "/" in args[0]:
        args[0] = shutil.which(args[0])

    # When bufsize=1, the subprocess' stdout is set to be line buffered.
    # In that case, we also need to specify universal_newlines=True.
    p = subprocess.Popen(args, bufsize=1, universal_newlines=True, 
                            stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, 
                            stderr=subprocess.STDOUT, text=True)
    selector = selectors.DefaultSelector()
    proc_output = io.StringIO()
    def read_stdout():
        while p.stdout.readable():
            output = p.stdout.readline()
            if len(output) == 0:
                break
            if tee:
                print(f"{output}", end="")
            proc_output.write(output)

    selector.register(p.stdout, selectors.EVENT_READ, read_stdout)

    # Keep going until the process has not exited. At that point, p.poll() will 
    # return non-None.
    while p.poll() is None:
        events = selector.select()
        for key, mask in events:
            # print(f"event.key: {key}, event.mask: {mask}")
            callback = key.data
            callback()
    proc_exit_status = p.wait()
    return proc_exit_status, proc_output.getvalue()

Limitations:

  • The subprocess cannot accept input from stdin. It is set to be /dev/null.
  • The program does not run shell commands. Though that feature can be added with a simple tweak.
  • This utility always redirects stderr to stdout.

links

social