Python run a command with TTYs for stdin, stdout, stderr with output streaming (generator function)

run_cmd_tty_stream.py

def run_cmd_tty_stream(cmd, bytes_input=b''):
    """Streams the output of cmd with bytes_input to stdin,
    with stdin, stdout and stderr as TTYs.

    Each yield from this function is a tuple of (stream name [str], data [bytes])

    Adapted from https://stackoverflow.com/a/52954716/507629
    and https://gist.github.com/hayd/4f46a68fc697ba8888a7b517a414583e
    """
    # provide tty to enable line-buffering
    mo, so = pty.openpty() # stdout
    me, se = pty.openpty() # stderr
    mi, si = pty.openpty() # stdin

    p = subprocess.Popen(
        cmd,
        bufsize=1, stdin=si, stdout=so, stderr=se, 
        close_fds=True)
    for fd in [so, se, si]:
        os.close(fd)
    os.write(mi, bytes_input)

    timeout = 0.04  # seconds
    readable = [mo, me]
    fd_name = {mo: 'stdout', me: 'stderr'}
    try:
        while readable:
            ready, _, _ = select.select(readable, [], [], timeout)
            for fd in ready:
                try:
                    data = os.read(fd, 512)
                except OSError as e:
                    if e.errno != errno.EIO:
                        raise
                    # EIO means EOF on some systems
                    readable.remove(fd)
                else:
                    if not data: # EOF
                        readable.remove(fd)
                    yield (fd_name[fd], data)

    finally:
        for fd in [mo, me, mi]:
            os.close(fd)
        if p.poll() is None:
            p.kill()
        p.wait()
        return

test.py

output_lines = {'stdout': [], 'stderr': []}
current_line = {'stdout': b'', 'stderr': b''}
ansi_escape_re = re.compile(r'(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]') # for filtering ansi color codes

for name, data in run_cmd_tty_stream(cmd):
    current_line[name] += data
    eol = b'\n' in current_line[name][-2:]
    if eol:
        try:
            line = current_line[name].decode('utf-8')
        except AttributeError:
            line = current_line[name]
        output_lines[name].append(ansi_escape_re.sub('', line))
        current_line[name] = b''
        print(line, end='')