diff --git a/IPython/utils/io.py b/IPython/utils/io.py index 46296be..7f1467b 100644 --- a/IPython/utils/io.py +++ b/IPython/utils/io.py @@ -10,15 +10,17 @@ IO related utilities. # the file COPYING, distributed as part of this software. #----------------------------------------------------------------------------- from __future__ import print_function +from __future__ import absolute_import #----------------------------------------------------------------------------- # Imports #----------------------------------------------------------------------------- +import codecs import os import sys import tempfile from .capture import CapturedIO, capture_output -from .py3compat import string_types, input +from .py3compat import string_types, input, PY3 #----------------------------------------------------------------------------- # Code @@ -227,3 +229,26 @@ def raw_print_err(*args, **kw): # Short aliases for quick debugging, do NOT use these in production code. rprint = raw_print rprinte = raw_print_err + +def unicode_std_stream(stream='stdout'): + """Get a wrapper to write unicode to stdout/stderr as UTF-8. + + This ignores environment variables and default encodings, to reliably write + unicode to stdout or stderr. + + :: + + unicode_std_stream().write(u'ł@e¶ŧ←') + """ + assert stream in ('stdout', 'stderr') + stream = getattr(sys, stream) + if PY3: + try: + stream_b = stream.buffer + except AttributeError: + # sys.stdout has been replaced - use it directly + return stream + else: + stream_b = stream + + return codecs.getwriter('utf-8')(stream_b) diff --git a/IPython/utils/tests/test_io.py b/IPython/utils/tests/test_io.py index 4818c07..9f28ba2 100644 --- a/IPython/utils/tests/test_io.py +++ b/IPython/utils/tests/test_io.py @@ -12,7 +12,9 @@ # Imports #----------------------------------------------------------------------------- from __future__ import print_function +from __future__ import absolute_import +import io as stdlib_io import sys from subprocess import Popen, PIPE @@ -20,7 +22,8 @@ import unittest import nose.tools as nt -from IPython.utils.io import Tee, capture_output +from IPython.testing.decorators import skipif +from IPython.utils.io import Tee, capture_output, unicode_std_stream from IPython.utils.py3compat import doctest_refactor_print, PY3 if PY3: @@ -88,3 +91,34 @@ def test_capture_output(): nt.assert_equal(io.stdout, 'hi, stdout\n') nt.assert_equal(io.stderr, 'hi, stderr\n') + +def test_UnicodeStdStream(): + # Test wrapping a bytes-level stdout + if PY3: + stdoutb = stdlib_io.BytesIO() + stdout = stdlib_io.TextIOWrapper(stdoutb, encoding='ascii') + else: + stdout = stdoutb = stdlib_io.BytesIO() + + orig_stdout = sys.stdout + sys.stdout = stdout + try: + sample = u"@łe¶ŧ←" + unicode_std_stream().write(sample) + + output = stdoutb.getvalue().decode('utf-8') + nt.assert_equal(output, sample) + assert not stdout.closed + finally: + sys.stdout = orig_stdout + +@skipif(not PY3, "Not applicable on Python 2") +def test_UnicodeStdStream_nowrap(): + # If we replace stdout with a StringIO, it shouldn't get wrapped. + orig_stdout = sys.stdout + sys.stdout = StringIO() + try: + nt.assert_is(unicode_std_stream(), sys.stdout) + assert not sys.stdout.closed + finally: + sys.stdout = orig_stdout \ No newline at end of file