exec_utils.py

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