X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=exec_utils.py;h=49484c61e40e4bcf9332e213c8afce2a00795e90;hb=e46158e49121b8a955bb07b73f5bcf9928b79c90;hp=0bcfc2771741040e79ae69d7bb93b02d85abe251;hpb=cdced3d08fad56ae3f4311d7f21d70e393942a0b;p=python_utils.git

diff --git a/exec_utils.py b/exec_utils.py
index 0bcfc27..49484c6 100644
--- a/exec_utils.py
+++ b/exec_utils.py
@@ -18,6 +18,8 @@ logger = logging.getLogger(__file__)
 
 def cmd_showing_output(
     command: str,
+    *,
+    timeout_seconds: Optional[float] = None,
 ) -> int:
     """Kick off a child process.  Capture and emit all output that it
     produces on stdout and stderr in a character by character manner
@@ -27,15 +29,22 @@ def cmd_showing_output(
 
     Args:
         command: the command to execute
+        timeout_seconds: terminate the subprocess if it takes longer
+            than N seconds; None means to wait as long as it takes.
 
     Returns:
         the exit status of the subprocess once the subprocess has
-        exited
+        exited.  Raises TimeoutExpired after killing the subprocess
+        if the timeout expires.
 
     Side effects:
         prints all output of the child process (stdout or stderr)
     """
 
+    def timer_expired(p):
+        p.kill()
+        raise subprocess.TimeoutExpired(command, timeout_seconds)
+
     line_enders = set([b'\n', b'\r'])
     sel = selectors.DefaultSelector()
     with subprocess.Popen(
@@ -46,36 +55,45 @@ def cmd_showing_output(
         stderr=subprocess.PIPE,
         universal_newlines=False,
     ) as p:
-        sel.register(p.stdout, selectors.EVENT_READ)  # type: ignore
-        sel.register(p.stderr, selectors.EVENT_READ)  # type: ignore
-        done = False
-        while not done:
-            for key, _ in sel.select():
-                char = key.fileobj.read(1)  # type: ignore
-                if not char:
-                    sel.unregister(key.fileobj)
-                    if len(sel.get_map()) == 0:
-                        sys.stdout.flush()
-                        sys.stderr.flush()
-                        sel.close()
-                        done = True
-                if key.fileobj is p.stdout:
-                    os.write(sys.stdout.fileno(), char)
-                    if char in line_enders:
-                        sys.stdout.flush()
-                else:
-                    os.write(sys.stderr.fileno(), char)
-                    if char in line_enders:
-                        sys.stderr.flush()
-        p.wait()
+        timer = None
+        if timeout_seconds:
+            import threading
+
+            timer = threading.Timer(timeout_seconds, timer_expired(p))
+            timer.start()
+        try:
+            sel.register(p.stdout, selectors.EVENT_READ)  # type: ignore
+            sel.register(p.stderr, selectors.EVENT_READ)  # type: ignore
+            done = False
+            while not done:
+                for key, _ in sel.select():
+                    char = key.fileobj.read(1)  # type: ignore
+                    if not char:
+                        sel.unregister(key.fileobj)
+                        if len(sel.get_map()) == 0:
+                            sys.stdout.flush()
+                            sys.stderr.flush()
+                            sel.close()
+                            done = True
+                    if key.fileobj is p.stdout:
+                        os.write(sys.stdout.fileno(), char)
+                        if char in line_enders:
+                            sys.stdout.flush()
+                    else:
+                        os.write(sys.stderr.fileno(), char)
+                        if char in line_enders:
+                            sys.stderr.flush()
+            p.wait()
+        finally:
+            if timer:
+                timer.cancel()
         return p.returncode
 
 
 def cmd_exitcode(command: str, timeout_seconds: Optional[float] = None) -> int:
-    """Run a command but do not let it run for more than timeout_seconds.
-    This code doesn't capture or rebroadcast the command's output.  It
-    returns the exit value of the command or raises a TimeoutExpired
-    exception if the deadline is exceeded.
+    """Run a command silently and return its exit code once it has
+    finished.  If timeout_seconds is provided and the command runs too
+    long it will raise a TimeoutExpired exception.
 
     Args:
         command: the command to run
@@ -99,9 +117,9 @@ def cmd_exitcode(command: str, timeout_seconds: Optional[float] = None) -> int:
 
 
 def cmd(command: str, timeout_seconds: Optional[float] = None) -> str:
-    """Run a command and capture its output to stdout (only) into a string
-    buffer.  Return that string as this function's output.  Raises
-    subprocess.CalledProcessError or TimeoutExpired on error.
+    """Run a command and capture its output to stdout and stderr into a
+    string buffer.  Return that string as this function's output.
+    Raises subprocess.CalledProcessError or TimeoutExpired on error.
 
     Args:
         command: the command to run
@@ -133,7 +151,7 @@ def cmd(command: str, timeout_seconds: Optional[float] = None) -> str:
 
 def run_silently(command: str, timeout_seconds: Optional[float] = None) -> None:
     """Run a command silently but raise subprocess.CalledProcessError if
-    it fails.
+    it fails and raise a TimeoutExpired if it runs too long.
 
     Args:
         command: the command to run