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