diff --git a/asv/console.py b/asv/console.py index 07a81b412..29076bec7 100644 --- a/asv/console.py +++ b/asv/console.py @@ -8,6 +8,7 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) +import io import codecs import contextlib import locale @@ -120,7 +121,8 @@ def _write_with_fallback(s, write, fileobj): Write the supplied string with the given write function like ``write(s)``, but use a writer for the locale's preferred encoding in case of a UnicodeEncodeError. Failing that attempt to write - with 'utf-8' or 'latin-1'. + with 'utf-8' or 'latin-1'. *fileobj* can be text or byte stream, + *s* can be unicode or bytes. """ try: write(s) @@ -135,6 +137,14 @@ def _write_with_fallback(s, write, fileobj): except LookupError: Writer = codecs.getwriter('utf-8') + if isinstance(fileobj, io.TextIOBase): + # Get the byte stream + fileobj = fileobj.buffer + + if six.PY3 and isinstance(s, bytes): + # Writers expect unicode input + s = _decode_preferred_encoding(s) + f = Writer(fileobj) write = f.write @@ -157,7 +167,7 @@ def _write_with_fallback(s, write, fileobj): write(s) return write except UnicodeEncodeError: - write(s.encode('ascii', 'replace')) + write(s.encode('ascii', 'replace').decode('ascii')) return write diff --git a/test/test_console.py b/test/test_console.py new file mode 100644 index 000000000..e4d661038 --- /dev/null +++ b/test/test_console.py @@ -0,0 +1,111 @@ +# -*- coding: utf-8 -*- +# Licensed under a 3-clause BSD style license - see LICENSE.rst + +from __future__ import (absolute_import, division, print_function, + unicode_literals) + +import six +import io +import sys +import locale +import itertools +import pytest +from contextlib import contextmanager + +from asv.console import _write_with_fallback, color_print + + +def test_write_with_fallback(capfd): + + def check_write(value, expected, stream_encoding, preferred_encoding): + old_getpreferredencoding = locale.getpreferredencoding + try: + locale.getpreferredencoding = lambda: preferred_encoding + + # Check writing to io.StringIO + stream = io.StringIO() + _write_with_fallback(value, stream.write, stream) + assert stream.getvalue() == value + + # Check writing to a text stream + buf = io.BytesIO() + stream = io.TextIOWrapper(buf, encoding=stream_encoding) + _write_with_fallback(value, stream.write, stream) + stream.flush() + got = buf.getvalue() + assert got == expected + + # Check writing to a byte stream (no stream encoding, so + # it should write in locale encoding) + if stream_encoding == preferred_encoding: + buf = io.BytesIO() + _write_with_fallback(value, buf.write, buf) + got = buf.getvalue() + assert got == expected + + # Check writing to a file + with io.open('tmp.txt', 'w', encoding=stream_encoding) as stream: + _write_with_fallback(value, stream.write, stream) + with open('tmp.txt', 'rb') as stream: + got = stream.read() + assert got == expected + + # Check writing to Py2 files + if not six.PY3: + if stream_encoding == preferred_encoding: + # No stream encoding: write in locale encoding + for mode in ['w', 'wb']: + with open('tmp.txt', mode) as stream: + _write_with_fallback(value, stream.write, stream) + with open('tmp.txt', 'rb') as stream: + got = stream.read() + assert got == expected + finally: + locale.getpreferredencoding = old_getpreferredencoding + + # What is printed should follow the following rules: + # + # - Try printing in stream encoding. + # - Try printing in locale preferred encoding. + # - Otherwise, map characters produced by asv to ascii equivalents, and + # - Try to print in latin1 + # - Try to print in ascii, replacing all non-ascii characters + encodings = ['utf-8', 'latin1', 'ascii', 'euc-jp'] + strings = ["helloμ", "hello·", "hello難", "helloä"] + repmap = {"helloμ": "hellou", "hello·": "hello-"} + + for pref_enc, stream_enc, s in itertools.product(encodings, encodings, strings): + expected = None + for enc in [stream_enc, pref_enc]: + try: + expected = s.encode(enc) + break + except UnicodeError: + pass + else: + s2 = repmap.get(s, s) + try: + expected = s2.encode('latin1') + except UnicodeError: + expected = s2.encode('ascii', 'replace') + + check_write(s, expected, stream_enc, pref_enc) + + # Should not bail out on bytes input + _write_with_fallback("a".encode('ascii'), sys.stdout.write, sys.stdout) + out, err = capfd.readouterr() + assert out == "a" + + +def test_color_print_nofail(capfd): + # Try out color print + + color_print("hello", "red") + color_print("indeed難", "blue") + color_print(b"really\xfe", "green", "not really") + + out, err = capfd.readouterr() + assert 'hello' in out + assert 'indeed' in out + assert 'really' in out + assert 'not really' in out