exec_utils.py

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