exec_utils.py

   1 #!/usr/bin/env python3
   2
   3 import atexit
   4 import logging
   5 import selectors
   6 import shlex
   7 import subprocess
   8 import sys
   9 from typing import List, Optional
  10
  11
  12 logger = logging.getLogger(__file__)
  13
  14
  15 def cmd_showing_output(command: str, ) -> int:
  16     """Kick off a child process.  Capture and print all output that it
  17     produces on stdout and stderr.  Wait for the subprocess to exit
  18     and return the exit value as the return code of this function.
  19
  20     """
  21     line_enders = set([b'\n', b'\r'])
  22     p = subprocess.Popen(
  23         command,
  24         shell=True,
  25         bufsize=0,
  26         stdout=subprocess.PIPE,
  27         stderr=subprocess.PIPE,
  28         universal_newlines=False,
  29     )
  30     sel = selectors.DefaultSelector()
  31     sel.register(p.stdout, selectors.EVENT_READ)
  32     sel.register(p.stderr, selectors.EVENT_READ)
  33     stream_ends = 0
  34     while stream_ends < 2:
  35         for key, _ in sel.select():
  36             char = key.fileobj.read(1)
  37             if not char:
  38                 stream_ends += 1
  39                 continue
  40             if key.fileobj is p.stdout:
  41                 sys.stdout.buffer.write(char)
  42                 if char in line_enders:
  43                     sys.stdout.flush()
  44             else:
  45                 sys.stderr.buffer.write(char)
  46                 if char in line_enders:
  47                     sys.stderr.flush()
  48     p.wait()
  49     return p.returncode
  50
  51
  52 def cmd_with_timeout(command: str, timeout_seconds: Optional[float]) -> int:
  53     """Run a command but do not let it run for more than timeout seconds.
  54     Doesn't capture or rebroadcast command output.  Function returns
  55     the exit value of the command or raises a TimeoutExpired exception
  56     if the deadline is exceeded.
  57
  58     >>> cmd_with_timeout('/bin/echo foo', 10.0)
  59     0
  60
  61     >>> cmd_with_timeout('/bin/sleep 2', 0.1)
  62     Traceback (most recent call last):
  63     ...
  64     subprocess.TimeoutExpired: Command '['/bin/bash', '-c', '/bin/sleep 2']' timed out after 0.1 seconds
  65
  66     """
  67     return subprocess.check_call(
  68         ["/bin/bash", "-c", command], timeout=timeout_seconds
  69     )
  70
  71
  72 def cmd(command: str, timeout_seconds: Optional[float] = None) -> str:
  73     """Run a command and capture its output to stdout (only) in a string.
  74     Return that string as this function's output.  Raises
  75     subprocess.CalledProcessError or TimeoutExpired on error.
  76
  77     >>> cmd('/bin/echo foo')[:-1]
  78     'foo'
  79
  80     >>> cmd('/bin/sleep 2', 0.1)
  81     Traceback (most recent call last):
  82     ...
  83     subprocess.TimeoutExpired: Command '/bin/sleep 2' timed out after 0.1 seconds
  84
  85     """
  86     ret = subprocess.run(
  87         command,
  88         shell=True,
  89         capture_output=True,
  90         check=True,
  91         timeout=timeout_seconds,
  92     ).stdout
  93     return ret.decode("utf-8")
  94
  95
  96 def run_silently(command: str, timeout_seconds: Optional[float] = None) -> None:
  97     """Run a command silently but raise subprocess.CalledProcessError if
  98     it fails.
  99
 100     >>> run_silently("/usr/bin/true")
 101
 102     >>> run_silently("/usr/bin/false")
 103     Traceback (most recent call last):
 104     ...
 105     subprocess.CalledProcessError: Command '/usr/bin/false' returned non-zero exit status 1.
 106
 107     """
 108     subprocess.run(
 109         command,
 110         shell=True,
 111         stderr=subprocess.DEVNULL,
 112         stdout=subprocess.DEVNULL,
 113         capture_output=False,
 114         check=True,
 115         timeout=timeout_seconds,
 116     )
 117
 118
 119 def cmd_in_background(
 120         command: str, *, silent: bool = False
 121 ) -> subprocess.Popen:
 122     args = shlex.split(command)
 123     if silent:
 124         subproc = subprocess.Popen(args,
 125                                    stdin=subprocess.DEVNULL,
 126                                    stdout=subprocess.DEVNULL,
 127                                    stderr=subprocess.DEVNULL)
 128     else:
 129         subproc = subprocess.Popen(args, stdin=subprocess.DEVNULL)
 130
 131     def kill_subproc() -> None:
 132         try:
 133             if subproc.poll() is None:
 134                 logger.info("At exit handler: killing {}: {}".format(subproc, command))
 135                 subproc.terminate()
 136                 subproc.wait(timeout=10.0)
 137         except BaseException as be:
 138             logger.exception(be)
 139     atexit.register(kill_subproc)
 140     return subproc
 141
 142
 143 def cmd_list(command: List[str]) -> str:
 144     """Run a command with args encapsulated in a list and return the
 145     output text as a string.  Raises subprocess.CalledProcessError.
 146     """
 147     ret = subprocess.run(command, capture_output=True, check=True).stdout
 148     return ret.decode("utf-8")
 149
 150
 151 if __name__ == '__main__':
 152     import doctest
 153     doctest.testmod()