Last active
February 6, 2022 00:57
-
-
Save nathan815/ee35edc3b977c27dfff44a77763bb62f to your computer and use it in GitHub Desktop.
Python run a command with TTYs for stdin, stdout, stderr with output streaming (generator function)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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='') |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment