log_ouptut can take either a filename or a file object

This commit is contained in:
Todd Gamblin 2017-08-22 13:33:31 -07:00
parent 139d5bfa6b
commit 4f444c5f58

View File

@ -31,6 +31,8 @@
import sys
import traceback
from contextlib import contextmanager
from six import string_types
from six import StringIO
import llnl.util.tty as tty
@ -144,10 +146,10 @@ def __getattr__(self, attr):
return getattr(self.stream, attr)
def _file_descriptors_work():
"""Whether we can get file descriptors for stdout and stderr.
def _file_descriptors_work(*streams):
"""Whether we can get file descriptors for the streams specified.
This tries to call ``fileno()`` on ``sys.stdout`` and ``sys.stderr``
This tries to call ``fileno()`` on all streams in the argument list,
and returns ``False`` if anything goes wrong.
This can happen, when, e.g., the test framework replaces stdout with
@ -161,8 +163,8 @@ def _file_descriptors_work():
"""
# test whether we can get fds for out and error
try:
sys.stdout.fileno()
sys.stderr.fileno()
for stream in streams:
stream.fileno()
return True
except:
return False
@ -204,16 +206,22 @@ class log_output(object):
work within test frameworks like nose and pytest.
"""
def __init__(self, filename=None, echo=False, debug=False, buffer=False):
def __init__(self, file_like=None, echo=False, debug=False, buffer=False):
"""Create a new output log context manager.
Args:
filename (str): name of file where output should be logged
file_like (str or stream): open file object or name of file where
output should be logged
echo (bool): whether to echo output in addition to logging it
debug (bool): whether to enable tty debug mode during logging
buffer (bool): pass buffer=True to skip unbuffering output; note
this doesn't set up any *new* buffering
log_output can take either a file object or a filename. If a
filename is passed, the file will be opened and closed entirely
within ``__enter__`` and ``__exit__``. If a file object is passed,
this assumes the caller owns it and will close it.
By default, we unbuffer sys.stdout and sys.stderr because the
logger will include output from executed programs and from python
calls. If stdout and stderr are buffered, their output won't be
@ -222,16 +230,19 @@ def __init__(self, filename=None, echo=False, debug=False, buffer=False):
Logger daemon is not started until ``__enter__()``.
"""
self.filename = filename
self.file_like = file_like
self.echo = echo
self.debug = debug
self.buffer = buffer
self._active = False # used to prevent re-entry
def __call__(self, filename=None, echo=None, debug=None, buffer=None):
def __call__(self, file_like=None, echo=None, debug=None, buffer=None):
"""Thie behaves the same as init. It allows a logger to be reused.
Arguments are the same as for ``__init__()``. Args here take
precedence over those passed to ``__init__()``.
With the ``__call__`` function, you can save state between uses
of a single logger. This is useful if you want to remember,
e.g., the echo settings for a prior ``with log_output()``::
@ -245,8 +256,8 @@ def __call__(self, filename=None, echo=None, debug=None, buffer=None):
# log things; logger remembers prior echo settings.
"""
if filename is not None:
self.filename = filename
if file_like is not None:
self.file_like = file_like
if echo is not None:
self.echo = echo
if debug is not None:
@ -259,9 +270,23 @@ def __enter__(self):
if self._active:
raise RuntimeError("Can't re-enter the same log_output!")
if self.filename is None:
if self.file_like is None:
raise RuntimeError(
"filename must be set by either __init__ or __call__")
"file argument must be set by either __init__ or __call__")
# set up a stream for the daemon to write to
self.close_log_in_parent = True
self.write_log_in_parent = False
if isinstance(self.file_like, string_types):
self.log_file = open(self.file_like, 'w')
elif _file_descriptors_work(self.file_like):
self.log_file = self.file_like
self.close_log_in_parent = False
else:
self.log_file = StringIO()
self.write_log_in_parent = True
# record parent color settings before redirecting. We do this
# because color output depends on whether the *original* stdout
@ -304,7 +329,7 @@ def __enter__(self):
sys.stderr.flush()
# Now do the actual output rediction.
self.use_fds = _file_descriptors_work()
self.use_fds = _file_descriptors_work(sys.stdout, sys.stderr)
if self.use_fds:
# We try first to use OS-level file descriptors, as this
# redirects output for subprocesses and system calls.
@ -366,6 +391,14 @@ def __exit__(self, exc_type, exc_val, exc_tb):
sys.stdout = self._saved_stdout
sys.stderr = self._saved_stderr
# print log contents in parent if needed.
if self.write_log_in_parent:
string = self.parent.recv()
self.file_like.write(string)
if self.close_log_in_parent:
self.log_file.close()
# recover and store echo settings from the child before it dies
self.echo = self.parent.recv()
@ -409,51 +442,56 @@ def _writer_daemon(self, stdin):
# list of streams to select from
istreams = [in_pipe, stdin] if stdin else [in_pipe]
log_file = self.log_file
try:
with open(self.filename, 'w') as log_file:
with keyboard_input(stdin):
while True:
# Without the last parameter (timeout) select will
# wait until at least one of the two streams are
# ready. This may cause the function to hang.
rlist, _, xlist = select.select(istreams, [], [], 0)
with keyboard_input(stdin):
while True:
# Without the last parameter (timeout) select will
# wait until at least one of the two streams are
# ready. This may cause the function to hang.
rlist, _, xlist = select.select(istreams, [], [], 0)
# Allow user to toggle echo with 'v' key.
# Currently ignores other chars.
if stdin in rlist:
if stdin.read(1) == 'v':
echo = not echo
# Allow user to toggle echo with 'v' key.
# Currently ignores other chars.
if stdin in rlist:
if stdin.read(1) == 'v':
echo = not echo
# Handle output from the with block process.
if in_pipe in rlist:
# If we arrive here it means that in_pipe was
# ready for reading : it should never happen that
# line is false-ish
line = in_pipe.readline()
if not line:
break # EOF
# Handle output from the with block process.
if in_pipe in rlist:
# If we arrive here it means that in_pipe was
# ready for reading : it should never happen that
# line is false-ish
line = in_pipe.readline()
if not line:
break # EOF
# find control characters and strip them.
controls = control.findall(line)
line = re.sub(control, '', line)
# find control characters and strip them.
controls = control.findall(line)
line = re.sub(control, '', line)
# Echo to stdout if requested or forced
if echo or force_echo:
sys.stdout.write(line)
sys.stdout.flush()
# Echo to stdout if requested or forced
if echo or force_echo:
sys.stdout.write(line)
sys.stdout.flush()
# Stripped output to log file.
log_file.write(_strip(line))
log_file.flush()
if xon in controls:
force_echo = True
if xoff in controls:
force_echo = False
# Stripped output to log file.
log_file.write(_strip(line))
log_file.flush()
if xon in controls:
force_echo = True
if xoff in controls:
force_echo = False
except:
tty.error("Exception occurred in writer daemon!")
traceback.print_exc()
finally:
# send written data back to parent if we used a StringIO
if self.write_log_in_parent:
self.child.send(log_file.getvalue())
log_file.close()
# send echo value back to the parent so it can be preserved.
self.child.send(echo)