X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=exec_utils.py;h=df273352bb793d59f19717c9cb6806134dd5a628;hb=31c81f6539969a5eba864d3305f9fb7bf716a367;hp=a52e206ad95496ebf4de9eb08f19df8073339d4a;hpb=f83dffe5e1c358ebfee2583f950a42cf7e909969;p=python_utils.git diff --git a/exec_utils.py b/exec_utils.py index a52e206..df27335 100644 --- a/exec_utils.py +++ b/exec_utils.py @@ -2,50 +2,59 @@ 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: +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( + 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 + ) as p: + 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: + # sys.stdout.buffer.write(char) + os.write(sys.stdout.fileno(), char) + if char in line_enders: + sys.stdout.flush() + else: + # sys.stderr.buffer.write(char) + os.write(sys.stderr.fileno(), 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: @@ -63,9 +72,7 @@ def cmd_with_timeout(command: str, timeout_seconds: Optional[float]) -> int: subprocess.TimeoutExpired: Command '['/bin/bash', '-c', '/bin/sleep 2']' timed out after 0.1 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: @@ -115,26 +122,27 @@ 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: 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(f'At exit handler: killing {subproc} ({command})') subproc.terminate() subproc.wait(timeout=10.0) except BaseException as be: logger.exception(be) + atexit.register(kill_subproc) return subproc @@ -149,4 +157,5 @@ def cmd_list(command: List[str]) -> str: if __name__ == '__main__': import doctest + doctest.testmod()