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_with_timeout(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_with_timeout('/bin/echo foo', 10.0)
  90     0
  91
  92     >>> cmd_with_timeout('/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         capture_output=True,
 127         check=True,
 128         timeout=timeout_seconds,
 129     ).stdout
 130     return ret.decode("utf-8")
 131
 132
 133 def run_silently(command: str, timeout_seconds: Optional[float] = None) -> None:
 134     """Run a command silently but raise subprocess.CalledProcessError if
 135     it fails.
 136
 137     Args:
 138         command: the command to run
 139         timeout_seconds: the max number of seconds to allow the subprocess
 140             to execute or None to indicate no timeout
 141
 142     Returns:
 143         No return value; error conditions (including non-zero child process
 144         exits) produce exceptions.
 145
 146     >>> run_silently("/usr/bin/true")
 147
 148     >>> run_silently("/usr/bin/false")
 149     Traceback (most recent call last):
 150     ...
 151     subprocess.CalledProcessError: Command '/usr/bin/false' returned non-zero exit status 1.
 152
 153     """
 154     subprocess.run(
 155         command,
 156         shell=True,
 157         stderr=subprocess.DEVNULL,
 158         stdout=subprocess.DEVNULL,
 159         capture_output=False,
 160         check=True,
 161         timeout=timeout_seconds,
 162     )
 163
 164
 165 def cmd_in_background(command: str, *, silent: bool = False) -> subprocess.Popen:
 166     """Spawns a child process in the background and registers an exit
 167     handler to make sure we kill it if the parent process (us) is
 168     terminated.
 169
 170     Args:
 171         command: the command to run
 172         silent: do not allow any output from the child process to be displayed
 173             in the parent process' window
 174
 175     Returns:
 176         the :class:`Popen` object that can be used to communicate
 177             with the background process.
 178     """
 179     args = shlex.split(command)
 180     if silent:
 181         subproc = subprocess.Popen(
 182             args,
 183             stdin=subprocess.DEVNULL,
 184             stdout=subprocess.DEVNULL,
 185             stderr=subprocess.DEVNULL,
 186         )
 187     else:
 188         subproc = subprocess.Popen(args, stdin=subprocess.DEVNULL)
 189
 190     def kill_subproc() -> None:
 191         try:
 192             if subproc.poll() is None:
 193                 logger.info('At exit handler: killing %s (%s)', subproc, command)
 194                 subproc.terminate()
 195                 subproc.wait(timeout=10.0)
 196         except BaseException as be:
 197             logger.exception(be)
 198
 199     atexit.register(kill_subproc)
 200     return subproc
 201
 202
 203 def cmd_list(command: List[str]) -> str:
 204     """Run a command with args encapsulated in a list and return the
 205     output text as a string.  Raises subprocess.CalledProcessError.
 206     """
 207     ret = subprocess.run(command, capture_output=True, check=True).stdout
 208     return ret.decode("utf-8")
 209
 210
 211 if __name__ == '__main__':
 212     import doctest
 213
 214     doctest.testmod()