258 lines
8.7 KiB
Python
258 lines
8.7 KiB
Python
##############################################################################
|
|
# 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:
|
|
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
|