Merge pull request #336 from pv/unicode-fix

Fix encoding-fallback printing on Py3

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
 
 

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