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