X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=exec_utils.py;h=016310793152ddeb8872f5ec279c26ae1994655e;hb=36fea7f15ed17150691b5b3ead75450e575229ef;hp=b52f52f0dc785033fc437a638c734d79cf8d5aa8;hpb=b29be4f1750fd20bd2eada88e751dfae85817882;p=python_utils.git diff --git a/exec_utils.py b/exec_utils.py index b52f52f..0163107 100644 --- a/exec_utils.py +++ b/exec_utils.py @@ -2,31 +2,62 @@ import atexit import logging +import selectors import shlex import subprocess +import sys from typing import List, Optional logger = logging.getLogger(__file__) -def cmd_showing_output(command: str) -> None: +def cmd_showing_output( + command: str, +) -> int: + """Kick off a child process. Capture and print all output that it + produces on stdout and stderr. Wait for the subprocess to exit + and return the exit value as the return code of this function. + + """ + line_enders = set([b'\n', b'\r']) p = subprocess.Popen( command, shell=True, bufsize=0, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + universal_newlines=False, ) - for line in iter(p.stdout.readline, b''): - print(line.decode('utf-8'), end='') - p.stdout.close() + sel = selectors.DefaultSelector() + sel.register(p.stdout, selectors.EVENT_READ) + sel.register(p.stderr, selectors.EVENT_READ) + stream_ends = 0 + while stream_ends < 2: + for key, _ in sel.select(): + char = key.fileobj.read(1) + if not char: + stream_ends += 1 + continue + if key.fileobj is p.stdout: + sys.stdout.buffer.write(char) + if char in line_enders: + sys.stdout.flush() + else: + sys.stderr.buffer.write(char) + if char in line_enders: + sys.stderr.flush() p.wait() + sys.stdout.flush() + sys.stderr.flush() + return p.returncode def cmd_with_timeout(command: str, timeout_seconds: Optional[float]) -> int: - """ - Run a command but do not let it run for more than timeout seconds. + """Run a command but do not let it run for more than timeout seconds. + Doesn't capture or rebroadcast command output. Function returns + the exit value of the command or raises a TimeoutExpired exception + if the deadline is exceeded. >>> cmd_with_timeout('/bin/echo foo', 10.0) 0 @@ -43,8 +74,9 @@ def cmd_with_timeout(command: str, timeout_seconds: Optional[float]) -> int: def cmd(command: str, timeout_seconds: Optional[float] = None) -> str: - """Run a command with everything encased in a string and return - the output text as a string. Raises subprocess.CalledProcessError. + """Run a command and capture its output to stdout (only) in a string. + Return that string as this function's output. Raises + subprocess.CalledProcessError or TimeoutExpired on error. >>> cmd('/bin/echo foo')[:-1] 'foo' @@ -56,12 +88,16 @@ def cmd(command: str, timeout_seconds: Optional[float] = None) -> str: """ ret = subprocess.run( - command, shell=True, capture_output=True, check=True, timeout=timeout_seconds, + command, + shell=True, + capture_output=True, + check=True, + timeout=timeout_seconds, ).stdout return ret.decode("utf-8") -def run_silently(command: str) -> None: +def run_silently(command: str, timeout_seconds: Optional[float] = None) -> None: """Run a command silently but raise subprocess.CalledProcessError if it fails. @@ -74,31 +110,41 @@ def run_silently(command: str) -> None: """ subprocess.run( - command, shell=True, stderr=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, capture_output=False, check=True + command, + shell=True, + stderr=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + capture_output=False, + check=True, + timeout=timeout_seconds, ) def cmd_in_background( - command: str, *, silent: bool = False + command: str, *, silent: bool = False ) -> subprocess.Popen: args = shlex.split(command) if silent: - subproc = subprocess.Popen(args, - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL) + subproc = subprocess.Popen( + args, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) else: subproc = subprocess.Popen(args, stdin=subprocess.DEVNULL) def kill_subproc() -> None: try: if subproc.poll() is None: - logger.info("At exit handler: killing {}: {}".format(subproc, command)) + logger.info( + "At exit handler: killing {}: {}".format(subproc, command) + ) subproc.terminate() subproc.wait(timeout=10.0) except BaseException as be: logger.exception(be) + atexit.register(kill_subproc) return subproc @@ -113,4 +159,5 @@ def cmd_list(command: List[str]) -> str: if __name__ == '__main__': import doctest + doctest.testmod()