exec_utils.py

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