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