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