From 00351001863f22d61bda745e2be031e160b5f51f Mon Sep 17 00:00:00 2001 From: Todd Gamblin Date: Fri, 29 May 2015 17:20:08 -0700 Subject: [PATCH] SPACK-69: Add context manager to fork and log output in a with block. --- lib/spack/llnl/util/tty/log.py | 178 +++++++++++++++++++++++++++++++++ 1 file changed, 178 insertions(+) create mode 100644 lib/spack/llnl/util/tty/log.py diff --git a/lib/spack/llnl/util/tty/log.py b/lib/spack/llnl/util/tty/log.py new file mode 100644 index 0000000000..6ccd0e66d9 --- /dev/null +++ b/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