Ran black code formatter on everything.

[python_utils.git] / exec_utils.py

diff --git a/exec_utils.py b/exec_utils.py
index 36d48fad2b4c3191c0db0c612a58db502e6be197..016310793152ddeb8872f5ec279c26ae1994655e 100644 (file)
--- a/exec_utils.py
+++ b/exec_utils.py
@@ -2,6 +2,7 @@
 
 import atexit
 import logging
+import selectors
 import shlex
 import subprocess
 import sys
@@ -11,7 +12,14 @@ from typing import List, Optional
 logger = logging.getLogger(__file__)
 
 
-def cmd_showing_output(command: str) -> None:
+def cmd_showing_output(
+    command: str,
+) -> int:
+    """Kick off a child process.  Capture and print all output that it
+    produces on stdout and stderr.  Wait for the subprocess to exit
+    and return the exit value as the return code of this function.
+
+    """
     line_enders = set([b'\n', b'\r'])
     p = subprocess.Popen(
         command,
@@ -21,16 +29,35 @@ def cmd_showing_output(command: str) -> None:
         stderr=subprocess.PIPE,
         universal_newlines=False,
     )
-    for char in iter(lambda: p.stdout.read(1), b''):
-        sys.stdout.buffer.write(char)
-        if char in line_enders:
-            sys.stdout.flush()
+    sel = selectors.DefaultSelector()
+    sel.register(p.stdout, selectors.EVENT_READ)
+    sel.register(p.stderr, selectors.EVENT_READ)
+    stream_ends = 0
+    while stream_ends < 2:
+        for key, _ in sel.select():
+            char = key.fileobj.read(1)
+            if not char:
+                stream_ends += 1
+                continue
+            if key.fileobj is p.stdout:
+                sys.stdout.buffer.write(char)
+                if char in line_enders:
+                    sys.stdout.flush()
+            else:
+                sys.stderr.buffer.write(char)
+                if char in line_enders:
+                    sys.stderr.flush()
     p.wait()
+    sys.stdout.flush()
+    sys.stderr.flush()
+    return p.returncode
 
 
 def cmd_with_timeout(command: str, timeout_seconds: Optional[float]) -> int:
-    """
-    Run a command but do not let it run for more than timeout seconds.
+    """Run a command but do not let it run for more than timeout seconds.
+    Doesn't capture or rebroadcast command output.  Function returns
+    the exit value of the command or raises a TimeoutExpired exception
+    if the deadline is exceeded.
 
     >>> cmd_with_timeout('/bin/echo foo', 10.0)
     0
@@ -47,8 +74,9 @@ def cmd_with_timeout(command: str, timeout_seconds: Optional[float]) -> int:
 
 
 def cmd(command: str, timeout_seconds: Optional[float] = None) -> str:
-    """Run a command with everything encased in a string and return
-    the output text as a string.  Raises subprocess.CalledProcessError.
+    """Run a command and capture its output to stdout (only) in a string.
+    Return that string as this function's output.  Raises
+    subprocess.CalledProcessError or TimeoutExpired on error.
 
     >>> cmd('/bin/echo foo')[:-1]
     'foo'
@@ -60,12 +88,16 @@ def cmd(command: str, timeout_seconds: Optional[float] = None) -> str:
 
     """
     ret = subprocess.run(
-        command, shell=True, capture_output=True, check=True, timeout=timeout_seconds,
+        command,
+        shell=True,
+        capture_output=True,
+        check=True,
+        timeout=timeout_seconds,
     ).stdout
     return ret.decode("utf-8")
 
 
-def run_silently(command: str) -> None:
+def run_silently(command: str, timeout_seconds: Optional[float] = None) -> None:
     """Run a command silently but raise subprocess.CalledProcessError if
     it fails.
 
@@ -78,31 +110,41 @@ def run_silently(command: str) -> None:
 
     """
     subprocess.run(
-        command, shell=True, stderr=subprocess.DEVNULL,
-        stdout=subprocess.DEVNULL, capture_output=False, check=True
+        command,
+        shell=True,
+        stderr=subprocess.DEVNULL,
+        stdout=subprocess.DEVNULL,
+        capture_output=False,
+        check=True,
+        timeout=timeout_seconds,
     )
 
 
 def cmd_in_background(
-        command: str, *, silent: bool = False
+    command: str, *, silent: bool = False
 ) -> subprocess.Popen:
     args = shlex.split(command)
     if silent:
-        subproc = subprocess.Popen(args,
-                                   stdin=subprocess.DEVNULL,
-                                   stdout=subprocess.DEVNULL,
-                                   stderr=subprocess.DEVNULL)
+        subproc = subprocess.Popen(
+            args,
+            stdin=subprocess.DEVNULL,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
     else:
         subproc = subprocess.Popen(args, stdin=subprocess.DEVNULL)
 
     def kill_subproc() -> None:
         try:
             if subproc.poll() is None:
-                logger.info("At exit handler: killing {}: {}".format(subproc, command))
+                logger.info(
+                    "At exit handler: killing {}: {}".format(subproc, command)
+                )
                 subproc.terminate()
                 subproc.wait(timeout=10.0)
         except BaseException as be:
             logger.exception(be)
+
     atexit.register(kill_subproc)
     return subproc
 
@@ -117,4 +159,5 @@ def cmd_list(command: List[str]) -> str:
 
 if __name__ == '__main__':
     import doctest
+
     doctest.testmod()