From 7b551bd1d1c5b3d4ef2dcebd56ed42cdcd9aec76 Mon Sep 17 00:00:00 2001
From: "Brian C. Lane" <bcl@redhat.com>
Date: Wed, 29 Jul 2015 15:41:55 -0700
Subject: [PATCH] Add execReadlines to executils.

Returns output in realtime instead of buffering it.
---
 src/pylorax/executils.py | 96 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 92 insertions(+), 4 deletions(-)

diff --git a/src/pylorax/executils.py b/src/pylorax/executils.py
index 313cf777c..ac87a9dd7 100644
--- a/src/pylorax/executils.py
+++ b/src/pylorax/executils.py
@@ -20,8 +20,8 @@
 
 import os
 import subprocess
+from subprocess import TimeoutExpired
 import signal
-from time import sleep
 
 import logging
 log = logging.getLogger("pylorax")
@@ -151,9 +151,13 @@ def _run_program(argv, root='/', stdin=None, stdout=None, env_prune=None, log_ou
 
         if callback:
             while callback(proc) and proc.poll() is None:
-                sleep(1)
-
-        (output_string, err_string) = proc.communicate()
+                try:
+                    (output_string, err_string) = proc.communicate(timeout=1)
+                    break
+                except TimeoutExpired:
+                    pass
+        else:
+            (output_string, err_string) = proc.communicate()
         if output_string:
             if binary_output:
                 output_lines = [output_string]
@@ -224,6 +228,7 @@ def execWithCapture(command, argv, stdin=None, root='/', log_output=True, filter
         :param log_output: Whether to log the output of command
         :param filter_stderr: Whether stderr should be excluded from the returned output
         :param raise_err: whether to raise a CalledProcessError if the returncode is non-zero
+        :param raise_err: whether to raise a CalledProcessError if the returncode is non-zero
         :return: The output of the command
     """
     argv = [command] + list(argv)
@@ -231,6 +236,89 @@ def execWithCapture(command, argv, stdin=None, root='/', log_output=True, filter
                         raise_err=raise_err, callback=callback, env_add=env_add,
                         reset_handlers=reset_handlers, reset_lang=reset_lang)[1]
 
+def execReadlines(command, argv, stdin=None, root='/', env_prune=None, filter_stderr=False,
+                  raise_err=False, callback=lambda x: True, env_add=None, reset_handlers=True, reset_lang=True):
+    """ Execute an external command and return the line output of the command
+        in real-time.
+
+        This method assumes that there is a reasonably low delay between the
+        end of output and the process exiting. If the child process closes
+        stdout and then keeps on truckin' there will be problems.
+
+        NOTE/WARNING: UnicodeDecodeError will be raised if the output of the
+                      external command can't be decoded as UTF-8.
+
+        :param command: The command to run
+        :param argv: The argument list
+        :param stdin: The file object to read stdin from.
+        :param stdout: Optional file object to redirect stdout and stderr to.
+        :param root: The directory to chroot to before running command.
+        :param env_prune: environment variable to remove before execution
+        :param filter_stderr: Whether stderr should be excluded from the returned output
+        :param raise_err: whether to raise a CalledProcessError if the returncode is non-zero
+        :param raise_err: whether to raise a CalledProcessError if the returncode is non-zero
+
+        Output from the file is not logged to program.log
+        This returns an iterator with the lines from the command until it has finished
+    """
+
+    class ExecLineReader(object):
+        """Iterator class for returning lines from a process and cleaning
+           up the process when the output is no longer needed.
+        """
+
+        def __init__(self, proc, argv, callback):
+            self._proc = proc
+            self._argv = argv
+            self._callback = callback
+
+        def __iter__(self):
+            return self
+
+        def __del__(self):
+            # See if the process is still running
+            if self._proc.poll() is None:
+                # Stop the process and ignore any problems that might arise
+                try:
+                    self._proc.terminate()
+                except OSError:
+                    pass
+
+        def __next__(self):
+            # Read the next line, blocking if a line is not yet available
+            line = self._proc.stdout.readline().decode("utf-8")
+            if line == '' or not self._callback(self._proc):
+                # Output finished, wait for the process to end
+                self._proc.communicate()
+
+                # Check for successful exit
+                if self._proc.returncode < 0:
+                    raise OSError("process '%s' was killed by signal %s" %
+                            (self._argv, -self._proc.returncode))
+                elif self._proc.returncode > 0:
+                    raise OSError("process '%s' exited with status %s" %
+                            (self._argv, self._proc.returncode))
+                raise StopIteration
+
+            return line.strip()
+
+    argv = [command] + argv
+
+    if filter_stderr:
+        stderr = subprocess.DEVNULL
+    else:
+        stderr = subprocess.STDOUT
+
+    try:
+        proc = startProgram(argv, root=root, stdin=stdin, stdout=subprocess.PIPE, stderr=stderr, bufsize=1,
+                            env_prune=env_prune, env_add=env_add, reset_handlers=reset_handlers, reset_lang=reset_lang)
+    except OSError as e:
+        with program_log_lock:
+            program_log.error("Error running %s: %s", argv[0], e.strerror)
+        raise
+
+    return ExecLineReader(proc, argv, callback)
+
 def runcmd(cmd, **kwargs):
     """ run execWithRedirect with raise_err=True
     """