X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=exec_utils.py;h=b52f52f0dc785033fc437a638c734d79cf8d5aa8;hb=287360114f0a9d61d5dc3c3f168344df856ffbd5;hp=1b587405fb1a706a1d7b1c128fde2ad878401137;hpb=97fbe845e5dfdbda22521117c1783e1fd8515952;p=python_utils.git diff --git a/exec_utils.py b/exec_utils.py index 1b58740..b52f52f 100644 --- a/exec_utils.py +++ b/exec_utils.py @@ -1,29 +1,78 @@ #!/usr/bin/env python3 +import atexit +import logging import shlex import subprocess from typing import List, Optional +logger = logging.getLogger(__file__) + + +def cmd_showing_output(command: str) -> None: + p = subprocess.Popen( + command, + shell=True, + bufsize=0, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + for line in iter(p.stdout.readline, b''): + print(line.decode('utf-8'), end='') + p.stdout.close() + p.wait() + + 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. + + >>> cmd_with_timeout('/bin/echo foo', 10.0) + 0 + + >>> cmd_with_timeout('/bin/sleep 2', 0.1) + Traceback (most recent call last): + ... + 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 ) -def cmd(command: str) -> str: +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. + + >>> cmd('/bin/echo foo')[:-1] + 'foo' + + >>> cmd('/bin/sleep 2', 0.1) + Traceback (most recent call last): + ... + subprocess.TimeoutExpired: Command '/bin/sleep 2' timed out after 0.1 seconds + """ ret = subprocess.run( - command, shell=True, capture_output=True, check=True + command, shell=True, capture_output=True, check=True, timeout=timeout_seconds, ).stdout return ret.decode("utf-8") def run_silently(command: str) -> None: """Run a command silently but raise subprocess.CalledProcessError if - it fails.""" + it fails. + + >>> run_silently("/usr/bin/true") + + >>> run_silently("/usr/bin/false") + Traceback (most recent call last): + ... + subprocess.CalledProcessError: Command '/usr/bin/false' returned non-zero exit status 1. + + """ subprocess.run( command, shell=True, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL, capture_output=False, check=True @@ -35,13 +84,23 @@ def cmd_in_background( ) -> subprocess.Popen: args = shlex.split(command) if silent: - return 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: - return subprocess.Popen(args, - stdin=subprocess.DEVNULL) + 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)) + subproc.terminate() + subproc.wait(timeout=10.0) + except BaseException as be: + logger.exception(be) + atexit.register(kill_subproc) + return subproc def cmd_list(command: List[str]) -> str: @@ -50,3 +109,8 @@ def cmd_list(command: List[str]) -> str: """ ret = subprocess.run(command, capture_output=True, check=True).stdout return ret.decode("utf-8") + + +if __name__ == '__main__': + import doctest + doctest.testmod()