############################################################################## # 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