SPACK-69: Add context manager to fork and log output in a with block.

lib/spack/llnl/util/tty/log.py

@@ -0,0 +1,178 @@
+##############################################################################
+# Copyright (c) 2013-2015, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+#
+# This file is part of Spack.
+# Written by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
+# LLNL-CODE-647188
+#
+# For details, see https://scalability-llnl.github.io/spack
+# Please also see the LICENSE file for our notice and the LGPL.
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License (as published by
+# the Free Software Foundation) version 2.1 dated February 1999.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
+# conditions of the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program; if not, write to the Free Software Foundation,
+# Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
+##############################################################################
+"""Utility classes for logging the output of blocks of code.
+"""
+import sys
+import os
+import re
+import select
+import inspect
+import llnl.util.tty as tty
+import llnl.util.tty.color as color
+
+# Use this to strip escape sequences
+_escape = re.compile(r'\x1b[^m]*m|\x1b\[?1034h')
+
+def _strip(line):
+    """Strip color and control characters from a line."""
+    return _escape.sub('', line)
+
+
+class _SkipWithBlock():
+    """Special exception class used to skip a with block."""
+    pass
+
+
+class log_output(object):
+    """Redirects output and error of enclosed block to a file.
+
+    Usage:
+        with log_output(open('logfile.txt', 'w')):
+           # do things ... output will be logged.
+
+    or:
+        with log_output(open('logfile.txt', 'w'), echo=True):
+           # do things ... output will be logged
+           # and also printed to stdout.
+
+    Closes the provided stream when done with the block.
+    If echo is True, also prints the output to stdout.
+    """
+    def __init__(self, stream, echo=False, force_color=False, debug=False):
+        self.stream = stream
+
+        # various output options
+        self.echo = echo
+        self.force_color = force_color
+        self.debug = debug
+
+    def trace(self, frame, event, arg):
+        """Jumps to __exit__ on the child process."""
+        raise _SkipWithBlock()
+
+
+    def __enter__(self):
+        """Redirect output from the with block to a file.
+
+        This forks the with block as a separate process, with stdout
+        and stderr redirected back to the parent via a pipe.  If
+        echo is set, also writes to standard out.
+
+        """
+        # remember these values for later.
+        self._force_color = color._force_color
+        self._debug = tty._debug
+
+        read, write = os.pipe()
+
+        self.pid = os.fork()
+        if self.pid:
+            # Parent: read from child, skip the with block.
+            os.close(write)
+
+            read_file = os.fdopen(read, 'r', 0)
+            with self.stream as log_file:
+                while True:
+                    rlist, w, x = select.select([read_file], [], [])
+                    if not rlist:
+                        break
+
+                    line = read_file.readline()
+                    if not line:
+                        break
+
+                    # Echo to stdout if requested.
+                    if self.echo:
+                        sys.stdout.write(line)
+
+                    # Stripped output to log file.
+                    log_file.write(_strip(line))
+
+            read_file.flush()
+            read_file.close()
+
+            # Set a trace function to skip the with block.
+            sys.settrace(lambda *args, **keys: None)
+            frame = inspect.currentframe(1)
+            frame.f_trace = self.trace
+
+        else:
+            # Child: redirect output, execute the with block.
+            os.close(read)
+
+            # Save old stdout and stderr
+            self._stdout = os.dup(sys.stdout.fileno())
+            self._stderr = os.dup(sys.stderr.fileno())
+
+            # redirect to the pipe.
+            os.dup2(write, sys.stdout.fileno())
+            os.dup2(write, sys.stderr.fileno())
+
+            if self.force_color:
+                color._force_color = True
+
+            if self.debug:
+                tty._debug = True
+
+
+    def __exit__(self, exc_type, exception, traceback):
+        """Exits on child, handles skipping the with block on parent."""
+        # Child should just exit here.
+        if self.pid == 0:
+            # Flush the log to disk.
+            sys.stdout.flush()
+            sys.stderr.flush()
+
+            if exception:
+                # Restore stdout on the child if there's an exception,
+                # and let it be raised normally.
+                #
+                # This assumes that even if the exception is caught,
+                # the child will exit with a nonzero return code.  If
+                # it doesn't, the child process will continue running.
+                #
+                # TODO: think about how this works outside install.
+                # TODO: ideally would propagate exception to parent...
+                os.dup2(self._stdout, sys.stdout.fileno())
+                os.dup2(self._stderr, sys.stderr.fileno())
+
+                return False
+
+            else:
+                # Die quietly if there was no exception.
+                os._exit(0)
+
+        else:
+            # If the child exited badly, parent also should exit.
+            pid, returncode = os.waitpid(self.pid, 0)
+            if returncode != 0:
+                os._exit(1)
+
+        # restore output options.
+        color._force_color = self._force_color
+        tty._debug = self._debug
+
+        # Suppresses exception if it's our own.
+        return exc_type is _SkipWithBlock