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

##############################################################################
# Copyright (c) 2013-2016, Lawrence Livermore National Security, LLC.
# Produced at the Lawrence Livermore National Laboratory.
#
# This file is part of Spack.
# Created by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
# LLNL-CODE-647188
#
# For details, see https://github.com/llnl/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 Lesser General Public License (as
# published by the Free Software Foundation) version 2.1, 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 Lesser 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 keyboard_input(object):
    """Disable canonical input and echo on a stream within a with block.

    Use this with sys.stdin for keyboard input, e.g.:

        with keyboard_input(sys.stdin):
            r, w, x = select.select([sys.stdin], [], [])
            # ... do something with keypresses ...

    When the with block completes, this will restore settings before
    canonical and echo were disabled.
    """
    def __init__(self, stream):
        self.stream = stream


    def __enter__(self):
        self.old_cfg = None

        # Ignore all this if the input stream is not a tty.
        if not self.stream.isatty():
            return

        try:
            # import and mark whether it worked.
            import termios

            # save old termios settings
            fd = self.stream.fileno()
            self.old_cfg = termios.tcgetattr(fd)

            # create new settings with canonical input and echo
            # disabled, so keypresses are immediate & don't echo.
            self.new_cfg = termios.tcgetattr(fd)
            self.new_cfg[3] &= ~termios.ICANON
            self.new_cfg[3] &= ~termios.ECHO

            # Apply new settings for terminal
            termios.tcsetattr(fd, termios.TCSADRAIN, self.new_cfg)

        except Exception, e:
            pass  # Some OS's do not support termios, so ignore.


    def __exit__(self, exc_type, exception, traceback):
        # If termios was avaialble, restore old settings after the
        # with block
        if self.old_cfg:
            import termios
            termios.tcsetattr(
                self.stream.fileno(), termios.TCSADRAIN, self.old_cfg)


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

        # Default is to try file-descriptor reassignment unless the system 
        # out/err streams do not have an associated file descriptor
        self.directAssignment = False

    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:
                with keyboard_input(sys.stdin):
                    while True:
                        rlist, w, x = select.select([read_file, sys.stdin], [], [])
                        if not rlist:
                            break

                        # Allow user to toggle echo with 'v' key.
                        # Currently ignores other chars.
                        if sys.stdin in rlist:
                            if sys.stdin.read(1) == 'v':
                                self.echo = not self.echo

                        # handle output from the with block process.
                        if read_file in rlist:
                            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)

            try:
                # 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())
            except AttributeError:
                self.directAssignment = True
                self._stdout = sys.stdout
                self._stderr = sys.stderr
                output_redirect = os.fdopen(write, 'w')
                sys.stdout = output_redirect
                sys.stderr = output_redirect

            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...
                if self.directAssignment:
                    sys.stdout = self._stdout
                    sys.stderr = self._stderr
                else:
                    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