Merge pull request #25 from MartijnBraam/master

Process stdout and stderr from subprocesses separately

pmb/helpers/logging.py

@@ -74,3 +74,5 @@ def init(args):
     log_handler._args = args
     handler.setFormatter(formatter)
     root_logger.addHandler(handler)
+
+    logging.debug('*' * 40)

pmb/helpers/run.py

@@ -16,8 +16,65 @@ GNU General Public License for more details.
 You should have received a copy of the GNU General Public License
 along with pmbootstrap.  If not, see <http://www.gnu.org/licenses/>.
 """
-import subprocess
 import logging
+import asyncio
+import locale
+
+
+@asyncio.coroutine
+def _execute(loop, args, cmd, log_message, log, return_stdout, check=True):
+    logging.debug(log_message)
+
+    class SubprocessProtocol(asyncio.SubprocessProtocol):
+        def __init__(self, future):
+            self.output = ""
+            self.error = ""
+            self.return_code = None
+            self._future = future
+
+        def pipe_data_received(self, fd, data):
+            nonlocal args
+            text = data.decode(locale.getpreferredencoding(False))
+
+            if fd == 1:
+                # stdout
+                if log:
+                    args.logfd.write(text)
+                else:
+                    print(text, end='')
+                self.output += text
+
+            elif fd == 2:
+                # stderr, possibly do something with color here
+                if log:
+                    args.logfd.write(text)
+                else:
+                    print(text, end='')
+                self.error += text
+
+            args.logfd.flush()
+
+        def process_exited(self):
+            self.return_code = 0
+            self._future.set_result(True)
+
+    exit_future = asyncio.Future(loop=loop)
+    create = loop.subprocess_exec(lambda: SubprocessProtocol(exit_future), *cmd)
+    transport, protocol = yield from create
+    yield from exit_future
+    transport.close()
+
+    return_code = transport.get_returncode()
+    if return_code != 0:
+        if check:
+            raise RuntimeError("Command failed: \n" + protocol.error)
+    args.logfd.write('Program exited with: {}\n'.format(transport.get_returncode()))
+    args.logfd.flush()
+
+    if return_stdout:
+        return protocol.output
+    else:
+        return return_code
 
 
 def core(args, cmd, log_message, log, return_stdout, check=True):
@@ -26,33 +83,14 @@ def core(args, cmd, log_message, log, return_stdout, check=True):
     Run the command and write the output to the log.
 
     :param check: raise an exception, when the command fails
+    :param log: send output to log instead of stdout
+    :param return_stdout: return the stdout from the called process
     """
-
-    try:
-        ret = None
-        if log:
-            if return_stdout:
-                ret = subprocess.check_output(cmd).decode("utf-8")
-                args.logfd.write(ret)
-            else:
-                subprocess.check_call(cmd, stdout=args.logfd,
-                                      stderr=args.logfd)
-            args.logfd.flush()
-        else:
-            logging.debug("*** output passed to pmbootstrap stdout, not" +
-                          " to this log ***")
-            subprocess.check_call(cmd)
-
-    except subprocess.CalledProcessError as exc:
-        if check:
-            if log:
-                logging.debug("^" * 70)
-                logging.info("NOTE: The failed command's output is above"
-                             " the ^^^ line in the logfile: " + args.log)
-            raise RuntimeError("Command failed: " + log_message) from exc
-        else:
-            pass
-    return ret
+    loop = asyncio.get_event_loop()
+    loop.set_debug(False)
+    task = _execute(loop, args, cmd, log_message, log, return_stdout, check)
+    result = loop.run_until_complete(task)
+    return result
 
 
 def user(args, cmd, log=True, working_dir=None, return_stdout=False,