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