diff --git a/README.md b/README.md index f719e8e..9641b76 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,21 @@ with sys_pipes(): call_some_c_function() ``` +Forward C-level output to Python Logger objects (new in 3.1). +Each line of output will be a log message. + +```python +from wurlitzer import pipes, STDOUT +import logging + +logger = logging.getLogger("my.log") +logger.setLevel(logging.INFO) +logger.addHandler(logging.FileHandler("mycode.log")) + +with pipes(logger, stderr=STDOUT): + call_some_c_function() +``` + Or even simpler, enable it as an IPython extension: ``` diff --git a/test.py b/test.py index 24813c6..fcf5a1b 100644 --- a/test.py +++ b/test.py @@ -2,6 +2,7 @@ from __future__ import print_function import io +import logging import os import platform import sys @@ -177,3 +178,31 @@ def test_bufsize(): with wurlitzer.pipes(bufsize=bufsize) as (stdout, stderr): assert fcntl(sys.__stdout__, wurlitzer.F_GETPIPE_SZ) == bufsize assert fcntl(sys.__stderr__, wurlitzer.F_GETPIPE_SZ) == bufsize + + +def test_log_pipes(caplog): + with caplog.at_level(logging.INFO), wurlitzer.pipes( + logging.getLogger("wurlitzer.stdout"), logging.getLogger("wurlitzer.stderr") + ): + sys.__stdout__.write("some stdout\n") + sys.__stderr__.write("some stderr\n") + + stdout_logs = [] + stderr_logs = [] + for t in caplog.record_tuples: + if "stdout" in t[0]: + stdout_logs.append(t) + else: + stderr_logs.append(t) + + assert stdout_logs == [ + ("wurlitzer.stdout", logging.INFO, "some stdout"), + ] + assert stderr_logs == [ + ("wurlitzer.stderr", logging.ERROR, "some stderr"), + ] + + for record in caplog.records: + # check 'stream' extra + assert record.stream + assert record.name == f"wurlitzer.{record.stream}" diff --git a/wurlitzer.py b/wurlitzer.py index 1e65c92..5626f98 100644 --- a/wurlitzer.py +++ b/wurlitzer.py @@ -20,6 +20,7 @@ import ctypes import errno import io +import logging import os import platform import selectors @@ -183,6 +184,12 @@ def __init__( default: use /proc/sys/fs/pipe-max-size up to a max of 1MB if 0, will do nothing. """ + # accept logger objects + if stdout and isinstance(stdout, logging.Logger): + stdout = _LogPipe(stdout, stream_name="stdout", level=logging.INFO) + if stderr and isinstance(stderr, logging.Logger): + stderr = _LogPipe(stderr, stream_name="stderr", level=logging.ERROR) + self._stdout = stdout if stderr == STDOUT: self._stderr = self._stdout @@ -380,6 +387,18 @@ def pipes(stdout=PIPE, stderr=PIPE, encoding=_default_encoding, bufsize=None): The return value for the context manager is (stdout, stderr). + Args: + + stdout (optional, default: PIPE): None or PIPE or Writable or Logger + stderr (optional, default: PIPE): None or PIPE or STDOUT or Writable or Logger + encoding (optional): probably 'utf-8' + bufsize (optional): set explicit buffer size if the default doesn't work + + .. versionadded:: 3.1 + Accept Logger objects for stdout/stderr. + If a Logger is specified, each line will produce a log message. + stdout messages will be at INFO level, stderr messages at ERROR level. + .. versionchanged:: 3.0 when using `PIPE` (default), the type of captured output @@ -400,6 +419,13 @@ def pipes(stdout=PIPE, stderr=PIPE, encoding=_default_encoding, bufsize=None): PipeIO = io.StringIO else: PipeIO = io.BytesIO + + # accept logger objects + if stdout and isinstance(stdout, logging.Logger): + stdout = _LogPipe(stdout, stream_name="stdout", level=logging.INFO) + if stderr and isinstance(stderr, logging.Logger): + stderr = _LogPipe(stderr, stream_name="stderr", level=logging.ERROR) + # setup stdout if stdout == PIPE: stdout_r = stdout_w = PipeIO() @@ -420,6 +446,10 @@ def pipes(stdout=PIPE, stderr=PIPE, encoding=_default_encoding, bufsize=None): with w: yield stdout_r, stderr_r finally: + if stdout and isinstance(stdout, _LogPipe): + stdout.flush() + if stderr and isinstance(stderr, _LogPipe): + stderr.flush() # close pipes if stdout_pipe: # seek to 0 so that it can be read after exit @@ -429,6 +459,52 @@ def pipes(stdout=PIPE, stderr=PIPE, encoding=_default_encoding, bufsize=None): stderr_r.seek(0) +class _LogPipe(io.BufferedWriter): + """Writeable that writes lines to a Logger object as they arrive from captured pipes""" + + def __init__(self, logger, stream_name, level=logging.INFO): + self.logger = logger + self.stream_name = stream_name + self._buf = "" + self.level = level + + def _log(self, line): + """Log one line""" + self.logger.log(self.level, line.rstrip(), extra={"stream": self.stream_name}) + + def write(self, chunk): + """Given chunk, split into lines + + Log each line as a discrete message + + If it ends with a partial line, save it until the next one + """ + lines = chunk.splitlines(True) + if self._buf: + lines[0] = self._buf + lines[0] + if lines[-1].endswith("\n"): + self._buf = "" + else: + # last line is incomplete + self._buf = lines[-1] + lines = lines[:-1] + + for line in lines: + self._log(line) + + def flush(self): + """Write buffer as a last message if there is one""" + if self._buf: + self._log(self._buf) + self._buf = "" + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + self.flush() + + def sys_pipes(encoding=_default_encoding, bufsize=None): """Redirect C-level stdout/stderr to sys.stdout/stderr