diff --git a/Doc/library/test.rst b/Doc/library/test.rst --- a/Doc/library/test.rst +++ b/Doc/library/test.rst @@ -362,39 +362,49 @@ The :mod:`test.support` module defines t Here all warnings will be caught, and the test code tests the captured warnings directly. .. versionchanged:: 3.2 New optional arguments *filters* and *quiet*. .. function:: captured_stdin() - captured_stdout() - captured_stderr() + captured_stdout(encoding=None) + captured_stderr(encoding=None) A context managers that temporarily replaces the named stream with :class:`io.StringIO` object. Example use with output streams:: with captured_stdout() as stdout, captured_stderr() as stderr: print("hello") print("error", file=sys.stderr) assert stdout.getvalue() == "hello\n" assert stderr.getvalue() == "error\n" + You can use the *encoding* keyword-only parameter to simulate different + encodings for output streams:: + + with captured_stdout(encoding="ascii") as stdout: + print("hello") + assert stdout.getvalue() == b"hello\n" + Example use with input stream:: with captured_stdin() as stdin: stdin.write('hello\n') stdin.seek(0) # call test code that consumes from sys.stdin captured = input() self.assertEqual(captured, "hello") + .. versionchanged:: 3.5 + Added *encoding* keyword-only parameters to :func:`captured_stdout` and + :func:`captured_stderr`. .. function:: temp_dir(path=None, quiet=False) A context manager that creates a temporary directory at *path* and yields the directory. If *path* is None, the temporary directory is created using :func:`tempfile.mkdtemp`. If *quiet* is ``False``, the context manager diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py --- a/Lib/test/support/__init__.py +++ b/Lib/test/support/__init__.py @@ -1308,45 +1308,57 @@ def transient_internet(resource_name, *, filter_error(err) raise # XXX should we catch generic exceptions and look for their # __cause__ or __context__? finally: socket.setdefaulttimeout(old_timeout) +class StringIOWrapper: + + def __init__(self, stream, encoding): + self.stream = stream + self.encoding = encoding + + def getvalue(self): + if self.encoding is not None: + return self.stream.getvalue().encode(self.encoding) + return self.stream.getvalue() + + @contextlib.contextmanager -def captured_output(stream_name): +def captured_output(stream_name, *, encoding=None): """Return a context manager used by captured_stdout/stdin/stderr that temporarily replaces the sys stream *stream_name* with a StringIO.""" import io orig_stdout = getattr(sys, stream_name) setattr(sys, stream_name, io.StringIO()) try: - yield getattr(sys, stream_name) + yield StringIOWrapper(getattr(sys, stream_name), encoding) finally: setattr(sys, stream_name, orig_stdout) -def captured_stdout(): +def captured_stdout(*, encoding=None): """Capture the output of sys.stdout: with captured_stdout() as stdout: print("hello") self.assertEqual(stdout.getvalue(), "hello\n") """ - return captured_output("stdout") + return captured_output("stdout", encoding=encoding) -def captured_stderr(): +def captured_stderr(*, encoding=None): """Capture the output of sys.stderr: with captured_stderr() as stderr: print("hello", file=sys.stderr) self.assertEqual(stderr.getvalue(), "hello\n") """ - return captured_output("stderr") + return captured_output("stderr", encoding=encoding) def captured_stdin(): """Capture the input to sys.stdin: with captured_stdin() as stdin: stdin.write('hello\n') stdin.seek(0) # call test code that consumes from sys.stdin