ubelt.util_stream module

Functions for capturing and redirecting IO streams with optional tee-functionality.

The CaptureStdout captures all text sent to stdout and optionally prevents it from actually reaching stdout.

The TeeStringIO does the same thing but for arbitrary streams. It is how the former is implemented.

class ubelt.util_stream.TeeStringIO(redirect: IOBase | None = None)

Bases: StringIO

An IO object that writes to itself and another IO stream.

Variables:

redirect (io.IOBase | None) – The other stream to write to.

Example

>>> import ubelt as ub
>>> import io
>>> redirect = io.StringIO()
>>> self = ub.TeeStringIO(redirect)
>>> self.write('spam')
>>> assert self.getvalue() == 'spam'
>>> assert redirect.getvalue() == 'spam'

Parameters:

redirect (io.IOBase) – The other stream to write to.

isatty() → bool

Returns true of the redirect is a terminal.

Note

Needed for IPython.embed to work properly when this class is used to override stdout / stderr.

SeeAlso:

io.IOBase.isatty()

Returns:

bool

fileno() → int

Returns underlying file descriptor of the redirected IOBase object if one exists.

Returns:

the integer corresponding to the file descriptor

Return type:

int

SeeAlso:

io.IOBase.fileno()

Example

>>> import ubelt as ub
>>> dpath = ub.Path.appdir('ubelt/tests/util_stream').ensuredir()
>>> fpath = dpath / 'fileno-test.txt'
>>> with open(fpath, 'w') as file:
>>>     self = ub.TeeStringIO(file)
>>>     descriptor = self.fileno()
>>>     print(f'descriptor={descriptor}')
>>>     assert isinstance(descriptor, int)

Example

>>> # Test errors
>>> # Not sure the best way to test, this func is important for
>>> # capturing stdout when ipython embedding
>>> import io
>>> import pytest
>>> import ubelt as ub
>>> with pytest.raises(io.UnsupportedOperation):
>>>     ub.TeeStringIO(redirect=io.StringIO()).fileno()
>>> with pytest.raises(io.UnsupportedOperation):
>>>     ub.TeeStringIO(None).fileno()

property encoding: Any

Gets the encoding of the redirect IO object

FIXME:

Mypy complains that this violates the Liskov substitution principle because the return type can be str or None, whereas the parent class always returns a None. In the future we may raise an exception instead of returning None.

SeeAlso:

io.TextIOBase.encoding

Example

>>> import ubelt as ub
>>> redirect = io.StringIO()
>>> assert ub.TeeStringIO(redirect).encoding is None
>>> assert ub.TeeStringIO(None).encoding is None
>>> assert ub.TeeStringIO(sys.stdout).encoding is sys.stdout.encoding
>>> redirect = io.TextIOWrapper(io.StringIO())
>>> assert ub.TeeStringIO(redirect).encoding is redirect.encoding

write(msg: str) → int

Write to this and the redirected stream

Parameters:

msg (str) – the data to write

SeeAlso:

io.TextIOBase.write()

Example

>>> import ubelt as ub
>>> dpath = ub.Path.appdir('ubelt/tests/util_stream').ensuredir()
>>> fpath = dpath / 'write-test.txt'
>>> with open(fpath, 'w') as file:
>>>     self = ub.TeeStringIO(file)
>>>     n = self.write('hello world')
>>>     assert n == 11
>>> assert self.getvalue() == 'hello world'
>>> assert fpath.read_text() == 'hello world'

flush() → None

Flush to this and the redirected stream

SeeAlso:

io.IOBase.flush()

class ubelt.util_stream.CaptureStdout(suppress: bool = True, enabled: bool = True)

Bases: CaptureStream

Context manager that captures stdout and stores it in an internal stream.

Depending on the value of suppress, the user can control if stdout is printed (i.e. if stdout is tee-ed or suppressed) while it is being captured.

SeeAlso:

contextlib.redirect_stdout() - similar, but does not have the

ability to print stdout while it is being captured.

Variables:

• text (str | None) – internal storage for the most recent part

• parts (List[str]) – internal storage for all parts

• cap_stdout (None | TeeStringIO) – internal stream proxy

• orig_stdout (io.TextIOBase) – internal pointer to the original stdout stream

Example

>>> import ubelt as ub
>>> self = ub.CaptureStdout(suppress=True)
>>> print('dont capture the table flip (╯°□°）╯︵ ┻━┻')
>>> with self:
...     text = 'capture the heart ♥'
...     print(text)
>>> print('dont capture look of disapproval ಠ_ಠ')
>>> assert isinstance(self.text, str)
>>> assert self.text == text + '\n', 'failed capture text'

Example

>>> import ubelt as ub
>>> self = ub.CaptureStdout(suppress=False)
>>> with self:
...     print('I am captured and printed in stdout')
>>> assert self.text.strip() == 'I am captured and printed in stdout'

Example

>>> import ubelt as ub
>>> self = ub.CaptureStdout(suppress=True, enabled=False)
>>> with self:
...     print('dont capture')
>>> assert self.text is None

_get_stream() → TextIO

_set_stream(value: TextIO) → None

property cap_stdout: TeeStringIO | None

Backward-compatibility alias for cap_stream.

property orig_stdout: TextIO | None

Backward-compatibility alias for orig_stream.

class ubelt.util_stream.CaptureStderr(suppress: bool = True, enabled: bool = True)

Bases: CaptureStream

Context manager that captures stderr and stores it in an internal stream.

Behavior mirrors CaptureStdout, but for sys.stderr.

Example

>>> import sys
>>> self = CaptureStderr(suppress=True)
>>> with self:
...     print('to stdout (not captured)')
...     print('to stderr (captured)', file=sys.stderr)
>>> assert 'to stderr (captured)' in (self.text or '')

_get_stream() → TextIO

_set_stream(value: TextIO) → None

class ubelt.util_stream.CaptureStream(suppress: bool = True, enabled: bool = True)

Bases: object

Generic context manager for capturing a global text stream (stdout/stderr), with optional tee/suppress behavior and incremental reads.

Subclasses must override _get_stream() and _set_stream(value) to read/write the process-global stream they manage.

Variables:

• text (str | None) – most recent captured chunk from log_part().

• parts (list[str]) – all captured chunks appended by log_part().

• cap_stream (None | TeeStringIO) – proxy stream used while capturing.

• orig_stream (TextIO | None) – original global stream restored on stop.

• suppress (bool) – if True, do not tee to the original stream while capturing.

• enabled (bool) – if False, acts as a no-op context manager.

• started (bool) – True while the capture is active.

_get_stream() → TextIO

_set_stream(value: TextIO) → None

_make_proxy() → TeeStringIO

Create a fresh TeeStringIO proxy with appropriate redirect target depending on suppress. Called at start of each capture.

log_part() → None

Log what has been captured since the last call to log_part().

start() → None

Begin capturing. Swaps the global stream to our TeeStringIO.

stop() → None

Stop capturing. Restores the original global stream.

close() → None

Close and drop the proxy buffer to release memory.