-
Notifications
You must be signed in to change notification settings - Fork 15
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add stderr context manager. related: sktime/sktime#6891, sktime/sktime#6653
- Loading branch information
Showing
3 changed files
with
98 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,64 @@ | ||
# -*- coding: utf-8 -*- | ||
"""Context manager to suppress stderr.""" | ||
|
||
__author__ = ["XinyuWu"] | ||
|
||
import io | ||
import sys | ||
|
||
|
||
class StderrMute: | ||
"""A context manager to suppress stderr. | ||
Exception handling on exit can be customized by overriding | ||
the ``_handle_exit_exceptions`` method. | ||
Parameters | ||
---------- | ||
active : bool, default=True | ||
Whether to suppress stderr or not. | ||
If True, stderr is suppressed. | ||
If False, stderr is not suppressed, and the context manager does nothing | ||
except catch and suppress ModuleNotFoundError. | ||
""" | ||
|
||
def __init__(self, active=True): | ||
self.active = active | ||
|
||
def __enter__(self): | ||
"""Context manager entry point.""" | ||
# capture stderr if active | ||
# store the original stderr so it can be restored in __exit__ | ||
if self.active: | ||
self._stderr = sys.stderr | ||
sys.stderr = io.StringIO() | ||
|
||
def __exit__(self, type, value, traceback): # noqa: A002 | ||
"""Context manager exit point.""" | ||
# restore stderr if active | ||
# if not active, nothing needs to be done, since stderr was not replaced | ||
if self.active: | ||
sys.stderr = self._stderr | ||
|
||
if type is not None: | ||
return self._handle_exit_exceptions(type, value, traceback) | ||
|
||
# if no exception was raised, return True to indicate successful exit | ||
# return statement not needed as type was None, but included for clarity | ||
return True | ||
|
||
def _handle_exit_exceptions(self, type, value, traceback): # noqa: A002 | ||
"""Handle exceptions raised during __exit__. | ||
Parameters | ||
---------- | ||
type : type | ||
The type of the exception raised. | ||
Known to be not-None and Exception subtype when this method is called. | ||
value : Exception | ||
The exception instance raised. | ||
traceback : traceback | ||
The traceback object associated with the exception. | ||
""" | ||
# by default, all exceptions are raised | ||
return False |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,31 @@ | ||
# -*- coding: utf-8 -*- | ||
# copyright: skbase developers, BSD-3-Clause License (see LICENSE file) | ||
"""Tests of stdout_mute and stderr_mute.""" | ||
import io | ||
import sys | ||
from contextlib import redirect_stderr, redirect_stdout | ||
|
||
import pytest | ||
|
||
from skbase.utils.stderr_mute import StderrMute | ||
from skbase.utils.stdout_mute import StdoutMute | ||
|
||
__author__ = ["XinyuWu"] | ||
|
||
|
||
@pytest.mark.parametrize( | ||
"mute, expected", [(True, ["", ""]), (False, ["test stdout", "test sterr"])] | ||
) | ||
def test_std_mute(mute, expected): | ||
"""Test StderrMute.""" | ||
stderr_io = io.StringIO() | ||
stdout_io = io.StringIO() | ||
|
||
try: | ||
with redirect_stderr(stderr_io), redirect_stdout(stdout_io): | ||
with StderrMute(mute), StdoutMute(mute): | ||
sys.stdout.write("test stdout") | ||
sys.stderr.write("test sterr") | ||
1 / 0 | ||
except ZeroDivisionError: | ||
assert expected == [stdout_io.getvalue(), stderr_io.getvalue()] |