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