X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=exec_utils.py;h=49484c61e40e4bcf9332e213c8afce2a00795e90;hb=e46158e49121b8a955bb07b73f5bcf9928b79c90;hp=a52e206ad95496ebf4de9eb08f19df8073339d4a;hpb=f83dffe5e1c358ebfee2583f950a42cf7e909969;p=python_utils.git diff --git a/exec_utils.py b/exec_utils.py index a52e206..49484c6 100644 --- a/exec_utils.py +++ b/exec_utils.py @@ -1,91 +1,148 @@ #!/usr/bin/env python3 +# © Copyright 2021-2022, Scott Gasch + +"""Helper methods concerned with executing subprocesses.""" + import atexit import logging +import os import selectors import shlex import subprocess import sys from typing import List, Optional - logger = logging.getLogger(__file__) -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. - +def cmd_showing_output( + command: str, + *, + timeout_seconds: Optional[float] = None, +) -> int: + """Kick off a child process. Capture and emit all output that it + produces on stdout and stderr in a character by character manner + so that we don't have to wait on newlines. This was done to + capture the output of a subprocess that created dots to show + incremental progress on a task and render it correctly. + + Args: + command: the command to execute + timeout_seconds: terminate the subprocess if it takes longer + than N seconds; None means to wait as long as it takes. + + Returns: + the exit status of the subprocess once the subprocess has + exited. Raises TimeoutExpired after killing the subprocess + if the timeout expires. + + Side effects: + prints all output of the child process (stdout or stderr) """ + + def timer_expired(p): + p.kill() + raise subprocess.TimeoutExpired(command, timeout_seconds) + line_enders = set([b'\n', b'\r']) - p = subprocess.Popen( + sel = selectors.DefaultSelector() + with subprocess.Popen( command, shell=True, bufsize=0, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=False, - ) - sel = selectors.DefaultSelector() - sel.register(p.stdout, selectors.EVENT_READ) - sel.register(p.stderr, selectors.EVENT_READ) - should_exit = False - while not should_exit: - for key, _ in sel.select(): - char = key.fileobj.read(1) - if not char: - should_exit = True - 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() - 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. - 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) + ) as p: + timer = None + if timeout_seconds: + import threading + + timer = threading.Timer(timeout_seconds, timer_expired(p)) + timer.start() + try: + sel.register(p.stdout, selectors.EVENT_READ) # type: ignore + sel.register(p.stderr, selectors.EVENT_READ) # type: ignore + done = False + while not done: + for key, _ in sel.select(): + char = key.fileobj.read(1) # type: ignore + if not char: + sel.unregister(key.fileobj) + if len(sel.get_map()) == 0: + sys.stdout.flush() + sys.stderr.flush() + sel.close() + done = True + if key.fileobj is p.stdout: + os.write(sys.stdout.fileno(), char) + if char in line_enders: + sys.stdout.flush() + else: + os.write(sys.stderr.fileno(), char) + if char in line_enders: + sys.stderr.flush() + p.wait() + finally: + if timer: + timer.cancel() + return p.returncode + + +def cmd_exitcode(command: str, timeout_seconds: Optional[float] = None) -> int: + """Run a command silently and return its exit code once it has + finished. If timeout_seconds is provided and the command runs too + long it will raise a TimeoutExpired exception. + + Args: + command: the command to run + timeout_seconds: the max number of seconds to allow the subprocess + to execute or None to indicate no timeout + + Returns: + the exit status of the subprocess once the subprocess has + exited + + >>> cmd_exitcode('/bin/echo foo', 10.0) 0 - >>> cmd_with_timeout('/bin/sleep 2', 0.1) + >>> cmd_exitcode('/bin/sleep 2', 0.01) Traceback (most recent call last): ... - subprocess.TimeoutExpired: Command '['/bin/bash', '-c', '/bin/sleep 2']' timed out after 0.1 seconds + subprocess.TimeoutExpired: Command '['/bin/bash', '-c', '/bin/sleep 2']' timed out after 0.01 seconds """ - return subprocess.check_call( - ["/bin/bash", "-c", command], timeout=timeout_seconds - ) + return subprocess.check_call(["/bin/bash", "-c", command], timeout=timeout_seconds) def cmd(command: str, timeout_seconds: Optional[float] = None) -> str: - """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. + """Run a command and capture its output to stdout and stderr into a + string buffer. Return that string as this function's output. + Raises subprocess.CalledProcessError or TimeoutExpired on error. + + Args: + command: the command to run + timeout_seconds: the max number of seconds to allow the subprocess + to execute or None to indicate no timeout + + Returns: + The captured output of the subprocess' stdout as a string buffer >>> cmd('/bin/echo foo')[:-1] 'foo' - >>> cmd('/bin/sleep 2', 0.1) + >>> cmd('/bin/sleep 2', 0.01) Traceback (most recent call last): ... - subprocess.TimeoutExpired: Command '/bin/sleep 2' timed out after 0.1 seconds + subprocess.TimeoutExpired: Command '/bin/sleep 2' timed out after 0.01 seconds """ ret = subprocess.run( command, shell=True, - capture_output=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, check=True, timeout=timeout_seconds, ).stdout @@ -94,7 +151,16 @@ def cmd(command: str, timeout_seconds: Optional[float] = None) -> str: def run_silently(command: str, timeout_seconds: Optional[float] = None) -> None: """Run a command silently but raise subprocess.CalledProcessError if - it fails. + it fails and raise a TimeoutExpired if it runs too long. + + Args: + command: the command to run + timeout_seconds: the max number of seconds to allow the subprocess + to execute or None to indicate no timeout + + Returns: + No return value; error conditions (including non-zero child process + exits) produce exceptions. >>> run_silently("/usr/bin/true") @@ -115,26 +181,40 @@ def run_silently(command: str, timeout_seconds: Optional[float] = None) -> None: ) -def cmd_in_background( - command: str, *, silent: bool = False -) -> subprocess.Popen: +def cmd_in_background(command: str, *, silent: bool = False) -> subprocess.Popen: + """Spawns a child process in the background and registers an exit + handler to make sure we kill it if the parent process (us) is + terminated. + + Args: + command: the command to run + silent: do not allow any output from the child process to be displayed + in the parent process' window + + Returns: + the :class:`Popen` object that can be used to communicate + with the background process. + """ 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 %s (%s)', subproc, command) subproc.terminate() subproc.wait(timeout=10.0) except BaseException as be: logger.exception(be) + atexit.register(kill_subproc) return subproc @@ -149,4 +229,5 @@ def cmd_list(command: List[str]) -> str: if __name__ == '__main__': import doctest + doctest.testmod()