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