Nathan
83 lines
2.0 KiB
Python
83 lines
2.0 KiB
Python
# SPDX-FileCopyrightText: 2026 Blender Authors
|
|
#
|
|
# SPDX-License-Identifier: GPL-3.0-or-later
|
|
|
|
"""
|
|
Context manager to capture STDOUT/STDERR while also forwarding to the real output.
|
|
|
|
Useful so the LLM may use print(..) style debugging and receive the results as part of the response.
|
|
"""
|
|
|
|
__all__ = (
|
|
"CaptureOutput",
|
|
)
|
|
|
|
import io
|
|
import sys
|
|
from typing import IO, Self
|
|
|
|
|
|
class _Tee(io.TextIOBase):
|
|
"""Write to both a :class:`io.StringIO` buffer and the original stream."""
|
|
__slots__ = (
|
|
"_buffer",
|
|
"_original",
|
|
)
|
|
|
|
def __init__(self, original: IO[str]) -> None:
|
|
self._buffer = io.StringIO()
|
|
self._original = original
|
|
|
|
def write(self, s: str) -> int:
|
|
self._original.write(s)
|
|
return self._buffer.write(s)
|
|
|
|
def flush(self) -> None:
|
|
self._original.flush()
|
|
self._buffer.flush()
|
|
|
|
def getvalue(self) -> str:
|
|
return self._buffer.getvalue()
|
|
|
|
|
|
class CaptureOutput:
|
|
"""
|
|
Context manager that captures STDOUT & STDERR.
|
|
|
|
Output is forwarded to the original streams in real time
|
|
and also stored for retrieval via :meth:`stdout` and :meth:`stderr`.
|
|
"""
|
|
__slots__ = (
|
|
"_tee_out",
|
|
"_tee_err",
|
|
"_original_out",
|
|
"_original_err",
|
|
)
|
|
|
|
def __enter__(self) -> Self:
|
|
self._original_out = sys.stdout
|
|
self._original_err = sys.stderr
|
|
self._tee_out = _Tee(self._original_out)
|
|
self._tee_err = _Tee(self._original_err)
|
|
sys.stdout = self._tee_out
|
|
sys.stderr = self._tee_err
|
|
return self
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_val: BaseException | None,
|
|
exc_tb: object,
|
|
) -> None:
|
|
del exc_type, exc_val, exc_tb
|
|
sys.stdout = self._original_out
|
|
sys.stderr = self._original_err
|
|
|
|
@property
|
|
def stdout(self) -> str:
|
|
return self._tee_out.getvalue()
|
|
|
|
@property
|
|
def stderr(self) -> str:
|
|
return self._tee_err.getvalue()
|