exec_utils.py

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