"""Output capture utilities for function execution logging.
Provides OutputCapturer, a context manager that simultaneously captures and
displays stdout, stderr, and logging output. Uses a "tee" strategy where
output is duplicated: sent to both the original destination (for normal
display) and to an internal buffer (for storage in execution logs).
Design Rationale:
Simple redirection (e.g., sys.stdout = StringIO()) would suppress output
from the user's view. The tee approach preserves normal output behavior
while also capturing it for later analysis, making it suitable for both
interactive use and production logging.
"""
import sys
import io
import logging
import traceback
from contextlib import ExitStack
# TODO: see if we can use https://capturer.readthedocs.io/en/latest/index.html
# TODO: see if we can use similar functionality from pytest
[docs]
class OutputCapturer:
"""Context manager that captures stdout, stderr, and logging output.
Uses a dual-stream "tee" approach: output is sent to both the original
streams (preserving normal display) and to an internal buffer (enabling
capture). This ensures users see output in real-time while also storing
it for later retrieval via get_output().
Example:
>>> with OutputCapturer() as capturer:
... print("Hello") # Prints normally AND is captured
... logging.info("Test") # Logged normally AND is captured
>>> output = capturer.get_output()
>>> assert "Hello" in output
>>> assert "Test" in output
"""
class _TeeStream:
"""Internal stream that duplicates output to two destinations.
Duplicates all write() calls to both the original stream (for normal
display) and a StringIO buffer (for capture). This enables simultaneous
capture and display of stdout/stderr.
Args:
original: The original stream (stdout or stderr) to preserve.
buffer: The StringIO buffer to capture output.
"""
def __init__(self, *, original, buffer):
self.original = original
self.buffer = buffer
def write(self, data):
"""Write data to both the original stream and the capture buffer.
Args:
data: The data to be written.
"""
self.original.write(data)
self.buffer.write(data)
def flush(self):
"""Flush both streams to ensure all data is written."""
self.original.flush()
self.buffer.flush()
class _CaptureHandler(logging.Handler):
"""Internal logging handler that captures and forwards log records.
Captures logging output to a buffer while also forwarding records to
the original handlers, preserving normal logging behavior.
Args:
buffer: The StringIO buffer to capture logging output.
original_handlers: The original logging handlers to forward records to.
"""
def __init__(self, *, buffer, original_handlers):
super().__init__()
self.buffer = buffer
self.original_handlers = original_handlers
def emit(self, record):
"""Emit a log record to both the capture buffer and original handlers.
Args:
record: The log record to be captured and forwarded.
"""
msg = self.format(record)
self.buffer.write(msg + '\n')
for handler in self.original_handlers:
handler.emit(record)
[docs]
def __init__(self):
"""Initialize the OutputCapturer.
Creates the capture buffer. The actual stream redirection happens
in __enter__ to support proper nesting of multiple OutputCapturer
contexts.
"""
self.captured_buffer = io.StringIO()
self._stack = None
def __repr__(self) -> str:
"""Return a string representation of the OutputCapturer.
Returns:
A string showing the current size of the captured output buffer.
"""
captured_size = len(self.captured_buffer.getvalue())
return f"OutputCapturer(captured_chars={captured_size})"
def __enter__(self):
"""Start capturing stdout, stderr, and logging output.
Stream redirection happens here (not in __init__) to properly support
nested OutputCapturer contexts. Each capturer captures the current
streams at entry time, which may already be another capturer's tee.
Returns:
The OutputCapturer instance for use as a context variable.
"""
self._stack = ExitStack()
# Capture current state at ENTRY time (may be another capturer's tee)
original_stdout = sys.stdout
original_stderr = sys.stderr
original_handlers = logging.root.handlers[:]
# Create tees pointing to current streams
tee_stdout = self._TeeStream(original=original_stdout, buffer=self.captured_buffer)
tee_stderr = self._TeeStream(original=original_stderr, buffer=self.captured_buffer)
capture_handler = self._CaptureHandler(
buffer=self.captured_buffer, original_handlers=original_handlers)
# Register cleanup callbacks (LIFO order - restored in reverse)
self._stack.callback(setattr, sys, 'stdout', original_stdout)
self._stack.callback(setattr, sys, 'stderr', original_stderr)
self._stack.callback(setattr, logging.root, 'handlers', original_handlers)
# Install our tee streams
sys.stdout = tee_stdout
sys.stderr = tee_stderr
logging.root.handlers = [capture_handler]
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Stop capturing and restore stdout, stderr, and logging to original state.
If an exception occurred, prints the traceback to the captured output
before restoring streams. Delegates cleanup to ExitStack which restores
streams in LIFO order, ensuring proper unwinding of nested contexts.
Args:
exc_type: Exception class if an exception occurred, None otherwise.
exc_val: Exception instance if an exception occurred, None otherwise.
exc_tb: Traceback if an exception occurred, None otherwise.
Returns:
Whatever ExitStack.__exit__ returns (False by default, allowing
exceptions to propagate).
"""
if exc_type is not None:
traceback.print_exc()
return self._stack.__exit__(exc_type, exc_val, exc_tb)
[docs]
def get_output(self) -> str:
"""Retrieve all captured output as a single string.
Returns:
Combined stdout, stderr, and logging output captured during the
context manager's lifetime.
"""
return self.captured_buffer.getvalue()