Skip to content

Commit

Permalink
Replace buffers inside of stdout and stderr
Browse files Browse the repository at this point in the history
The current implementation replaces the streams in `sys`,
which is not effective if those streams have already been
stored in other modules, and can be problematic if the
temporary replacement streams are copied into other modules
as they are not able to be cleaned up by the context manager.

Related to #75
  • Loading branch information
jayvdb committed Sep 6, 2019
1 parent 38b4b9b commit 6f8fc9f
Show file tree
Hide file tree
Showing 5 changed files with 48 additions and 11 deletions.
2 changes: 1 addition & 1 deletion .travis.yml
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,6 @@ jobs:
script:
- python --version
- pip list
- pytest --cov=src -p no:warnings
- pytest --cov=src -p no:warnings -s
- if (( $PYTHON_MAJOR == 3 && $PYTHON_MINOR == 7 )); then tox -e flake8; else echo "No flake8."; fi
- if (( $PYTHON_MAJOR == 3 && $PYTHON_MINOR == 7 )); then codecov; else echo "No codecov."; fi
4 changes: 2 additions & 2 deletions src/stdio_mgr/stdio_mgr.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@
TextIOWrapper,
)

from .types import _MultiCloseContextManager, ReplaceSysIoContextManager
from .types import _OutStreamsCloseContextManager, InjectSysIoContextManager


class _PersistedBytesIO(BytesIO):
Expand Down Expand Up @@ -266,7 +266,7 @@ class SafeCloseTeeStdin(_SafeCloseIOBase, TeeStdin):
"""


class StdioManager(ReplaceSysIoContextManager, _MultiCloseContextManager):
class StdioManager(InjectSysIoContextManager, _OutStreamsCloseContextManager):
r"""Substitute temporary text buffers for `stdio` in a managed context.
Context manager.
Expand Down
41 changes: 41 additions & 0 deletions src/stdio_mgr/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -177,3 +177,44 @@ def __exit__(self, exc_type, exc_value, traceback):
(sys.stdin, sys.stdout, sys.stderr) = self._prior_streams

return super().__exit__(exc_type, exc_value, traceback)


class _OutStreamsCloseContextManager(_MultiCloseContextManager):
"""Close stdout and stderr only."""

def __enter__(self):
"""Enter context of all members."""
with ExitStack() as stack:
with suppress(AttributeError):
stack.enter_context(self.stdout)
with suppress(AttributeError):
stack.enter_context(self.stderr)

self._close_files = stack.pop_all().close

return super().__enter__()


class InjectSysIoContextManager(StdioTuple):
"""Replace sys stdio with members of the tuple."""

def __enter__(self):
"""Enter context, replacing sys stdio objects with capturing streams."""
self._prior_stdin = sys.stdin
self._prior_raw_out = (sys.stdout.buffer.raw, sys.stderr.buffer.raw)

super().__enter__()

sys.stdin = self.stdin
sys.stdout.buffer.__init__(self.stdout.buffer.raw)
sys.stderr.buffer.__init__(self.stderr.buffer.raw)

return self

def __exit__(self, exc_type, exc_value, traceback):
"""Exit context, restoring state of sys module."""
sys.stdin = self._prior_stdin
sys.stdout.buffer.__init__(self._prior_raw_out[0])
sys.stderr.buffer.__init__(self._prior_raw_out[1])

return super().__exit__(exc_type, exc_value, traceback)
7 changes: 1 addition & 6 deletions tests/test_stdiomgr_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -481,11 +481,6 @@ def test_stdout_detached(convert_newlines):

assert convert_newlines("test str\n") == o.getvalue()

with pytest.raises(ValueError) as err:
print("anything")

assert str(err.value) == "underlying buffer has been detached"

f.write(convert_newlines("second test str\n").encode("utf8"))
f.flush()

Expand Down Expand Up @@ -529,7 +524,7 @@ def test_stdout_access_buffer_after_close(convert_newlines):
with pytest.raises(ValueError) as err:
print("anything")

assert str(err.value) == "I/O operation on closed file."
assert str(err.value) == "write to closed file"

assert convert_newlines("test str\nsecond test str\n") == o.getvalue()

Expand Down
5 changes: 3 additions & 2 deletions tox.ini
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ envlist=

[testenv]
commands=
pytest -p no:warnings
pytest -p no:warnings -s
deps=
pytest

Expand Down Expand Up @@ -43,7 +43,8 @@ commands=
[pytest]
addopts = -v -rsxX
xfail_strict = True

faulthandler_timeout = 5
timeout = 5

[flake8]
# W503: black formats binary operators to start of line
Expand Down

0 comments on commit 6f8fc9f

Please sign in to comment.