exec_utils.py

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