exec_utils.py

   1 #!/usr/bin/env python3
   2
   3 # © Copyright 2021-2022, Scott Gasch
   4
   5 """Helper methods concerned with executing subprocesses."""
   6
   7 import atexit
   8 import logging
   9 import os
  10 import selectors
  11 import shlex
  12 import subprocess
  13 import sys
  14 from typing import List, Optional
  15
  16 logger = logging.getLogger(__file__)
  17
  18
  19 def cmd_showing_output(
  20     command: str,
  21     *,
  22     timeout_seconds: Optional[float] = None,
  23 ) -> int:
  24     """Kick off a child process.  Capture and emit all output that it
  25     produces on stdout and stderr in a character by character manner
  26     so that we don't have to wait on newlines.  This was done to
  27     capture the output of a subprocess that created dots to show
  28     incremental progress on a task and render it correctly.
  29
  30     Args:
  31         command: the command to execute
  32         timeout_seconds: terminate the subprocess if it takes longer
  33             than N seconds; None means to wait as long as it takes.
  34
  35     Returns:
  36         the exit status of the subprocess once the subprocess has
  37         exited.  Raises TimeoutExpired after killing the subprocess
  38         if the timeout expires.
  39
  40     Side effects:
  41         prints all output of the child process (stdout or stderr)
  42     """
  43
  44     def timer_expired(p):
  45         p.kill()
  46         raise subprocess.TimeoutExpired(command, timeout_seconds)
  47
  48     line_enders = set([b'\n', b'\r'])
  49     sel = selectors.DefaultSelector()
  50     with subprocess.Popen(
  51         command,
  52         shell=True,
  53         bufsize=0,
  54         stdout=subprocess.PIPE,
  55         stderr=subprocess.PIPE,
  56         universal_newlines=False,
  57     ) as p:
  58         timer = None
  59         if timeout_seconds:
  60             import threading
  61
  62             timer = threading.Timer(timeout_seconds, timer_expired(p))
  63             timer.start()
  64         try:
  65             sel.register(p.stdout, selectors.EVENT_READ)  # type: ignore
  66             sel.register(p.stderr, selectors.EVENT_READ)  # type: ignore
  67             done = False
  68             while not done:
  69                 for key, _ in sel.select():
  70                     char = key.fileobj.read(1)  # type: ignore
  71                     if not char:
  72                         sel.unregister(key.fileobj)
  73                         if len(sel.get_map()) == 0:
  74                             sys.stdout.flush()
  75                             sys.stderr.flush()
  76                             sel.close()
  77                             done = True
  78                     if key.fileobj is p.stdout:
  79                         os.write(sys.stdout.fileno(), char)
  80                         if char in line_enders:
  81                             sys.stdout.flush()
  82                     else:
  83                         os.write(sys.stderr.fileno(), char)
  84                         if char in line_enders:
  85                             sys.stderr.flush()
  86             p.wait()
  87         finally:
  88             if timer:
  89                 timer.cancel()
  90         return p.returncode
  91
  92
  93 def cmd_exitcode(command: str, timeout_seconds: Optional[float] = None) -> int:
  94     """Run a command silently and return its exit code once it has
  95     finished.  If timeout_seconds is provided and the command runs too
  96     long it will raise a TimeoutExpired exception.
  97
  98     Args:
  99         command: the command to run
 100         timeout_seconds: the max number of seconds to allow the subprocess
 101             to execute or None to indicate no timeout
 102
 103     Returns:
 104         the exit status of the subprocess once the subprocess has
 105         exited
 106
 107     >>> cmd_exitcode('/bin/echo foo', 10.0)
 108     0
 109
 110     >>> cmd_exitcode('/bin/sleep 2', 0.01)
 111     Traceback (most recent call last):
 112     ...
 113     subprocess.TimeoutExpired: Command '['/bin/bash', '-c', '/bin/sleep 2']' timed out after 0.01 seconds
 114
 115     """
 116     return subprocess.check_call(["/bin/bash", "-c", command], timeout=timeout_seconds)
 117
 118
 119 def cmd(command: str, timeout_seconds: Optional[float] = None) -> str:
 120     """Run a command and capture its output to stdout and stderr into a
 121     string buffer.  Return that string as this function's output.
 122     Raises subprocess.CalledProcessError or TimeoutExpired on error.
 123
 124     Args:
 125         command: the command to run
 126         timeout_seconds: the max number of seconds to allow the subprocess
 127             to execute or None to indicate no timeout
 128
 129     Returns:
 130         The captured output of the subprocess' stdout as a string buffer
 131
 132     >>> cmd('/bin/echo foo')[:-1]
 133     'foo'
 134
 135     >>> cmd('/bin/sleep 2', 0.01)
 136     Traceback (most recent call last):
 137     ...
 138     subprocess.TimeoutExpired: Command '/bin/sleep 2' timed out after 0.01 seconds
 139
 140     """
 141     ret = subprocess.run(
 142         command,
 143         shell=True,
 144         stdout=subprocess.PIPE,
 145         stderr=subprocess.STDOUT,
 146         check=True,
 147         timeout=timeout_seconds,
 148     ).stdout
 149     return ret.decode("utf-8")
 150
 151
 152 def run_silently(command: str, timeout_seconds: Optional[float] = None) -> None:
 153     """Run a command silently but raise subprocess.CalledProcessError if
 154     it fails and raise a TimeoutExpired if it runs too long.
 155
 156     Args:
 157         command: the command to run
 158         timeout_seconds: the max number of seconds to allow the subprocess
 159             to execute or None to indicate no timeout
 160
 161     Returns:
 162         No return value; error conditions (including non-zero child process
 163         exits) produce exceptions.
 164
 165     >>> run_silently("/usr/bin/true")
 166
 167     >>> run_silently("/usr/bin/false")
 168     Traceback (most recent call last):
 169     ...
 170     subprocess.CalledProcessError: Command '/usr/bin/false' returned non-zero exit status 1.
 171
 172     """
 173     subprocess.run(
 174         command,
 175         shell=True,
 176         stderr=subprocess.DEVNULL,
 177         stdout=subprocess.DEVNULL,
 178         capture_output=False,
 179         check=True,
 180         timeout=timeout_seconds,
 181     )
 182
 183
 184 def cmd_in_background(command: str, *, silent: bool = False) -> subprocess.Popen:
 185     """Spawns a child process in the background and registers an exit
 186     handler to make sure we kill it if the parent process (us) is
 187     terminated.
 188
 189     Args:
 190         command: the command to run
 191         silent: do not allow any output from the child process to be displayed
 192             in the parent process' window
 193
 194     Returns:
 195         the :class:`Popen` object that can be used to communicate
 196             with the background process.
 197     """
 198     args = shlex.split(command)
 199     if silent:
 200         subproc = subprocess.Popen(
 201             args,
 202             stdin=subprocess.DEVNULL,
 203             stdout=subprocess.DEVNULL,
 204             stderr=subprocess.DEVNULL,
 205         )
 206     else:
 207         subproc = subprocess.Popen(args, stdin=subprocess.DEVNULL)
 208
 209     def kill_subproc() -> None:
 210         try:
 211             if subproc.poll() is None:
 212                 logger.info('At exit handler: killing %s (%s)', subproc, command)
 213                 subproc.terminate()
 214                 subproc.wait(timeout=10.0)
 215         except BaseException as be:
 216             logger.exception(be)
 217
 218     atexit.register(kill_subproc)
 219     return subproc
 220
 221
 222 def cmd_list(command: List[str]) -> str:
 223     """Run a command with args encapsulated in a list and return the
 224     output text as a string.  Raises subprocess.CalledProcessError.
 225     """
 226     ret = subprocess.run(command, capture_output=True, check=True).stdout
 227     return ret.decode("utf-8")
 228
 229
 230 if __name__ == '__main__':
 231     import doctest
 232
 233     doctest.testmod()