Simplify handling of encoding

better_exceptions/__init__.py

@@ -16,37 +16,30 @@
 import logging
 import sys
 
-from .formatter import THEME, MAX_LENGTH, PIPE_CHAR, CAP_CHAR, ExceptionFormatter
-from .encoding import to_byte
-from .color import SUPPORTS_COLOR, SHOULD_ENCODE, STREAM
+from .formatter import THEME, MAX_LENGTH, ExceptionFormatter
+from .color import SUPPORTS_COLOR, STREAM
 from .log import BetExcLogger, patch as patch_logging
 from .repl import interact, get_repl
 
 
 __version__ = '0.2.2'
 
 
+ENCODING = STREAM.encoding
 THEME = THEME.copy()  # Users customizing the theme should not impact core
 
 
-def write_stream(data, stream=STREAM):
-    if SHOULD_ENCODE:
-        data = to_byte(data)
-        stream.buffer.write(data)
-    else:
-        stream.write(data)
-
-
 def format_exception(exc, value, tb):
     # Rebuild each time to take into account any changes made by the user to the global parameters
-    formatter = ExceptionFormatter(colored=SUPPORTS_COLOR, theme=THEME, max_length=MAX_LENGTH,
-                                   pipe_char=PIPE_CHAR, cap_char=CAP_CHAR)
+    formatter = ExceptionFormatter(
+        colored=SUPPORTS_COLOR, theme=THEME, max_length=MAX_LENGTH, encoding=ENCODING
+    )
     return list(formatter.format_exception(exc, value, tb))
 
 
 def excepthook(exc, value, tb):
     formatted = ''.join(format_exception(exc, value, tb))
-    write_stream(formatted, STREAM)
+    STREAM.write(formatted)
 
 
 def hook():

better_exceptions/color.py

@@ -11,12 +11,9 @@
 import struct
 import sys
 
-from .encoding import to_byte as _byte
 
-
-STREAM = sys.stderr
-SHOULD_ENCODE = True
 SUPPORTS_COLOR = False
+STREAM = sys.stderr
 
 
 def get_terminfo_file():
@@ -55,35 +52,13 @@ def get_terminfo_file():
     return f
 
 
-class ProxyBufferStreamWrapper(object):
-
-    def __init__(self, wrapped):
-        self.__wrapped = wrapped
-
-    def __getattr__(self, name):
-        return getattr(self.__wrapped, name)
-
-    def write(self, text):
-        data = _byte(text)
-        self.__wrapped.buffer.write(data)
-
-
 if os.name == 'nt':
     from colorama import init as init_colorama, AnsiToWin32
 
     init_colorama(wrap=False)
 
-    stream = sys.stderr
-
-    # Colorama cannot work with bytes-string
-    # The stream is wrapped so that encoding of the stream is done after
-    # (once Colorama found ANSI codes and converted them to win32 calls)
-    # See issue #23 for more information
-    stream = ProxyBufferStreamWrapper(stream)
-    SHOULD_ENCODE = False
-
-    STREAM = AnsiToWin32(stream).stream
     SUPPORTS_COLOR = True
+    STREAM = AnsiToWin32(STREAM).stream
 else:
     if os.getenv('FORCE_COLOR', None) == '1':
         SUPPORTS_COLOR = True

better_exceptions/formatter.py

@@ -9,20 +9,10 @@
 import sys
 import traceback
 
-from .color import STREAM, SUPPORTS_COLOR
-from .encoding import ENCODING, to_byte, to_unicode
+from .color import SUPPORTS_COLOR
 from .repl import get_repl
 
 
-PIPE_CHAR = '\u2502'
-CAP_CHAR = '\u2514'
-
-try:
-    PIPE_CHAR.encode(ENCODING)
-except UnicodeEncodeError:
-    PIPE_CHAR = '|'
-    CAP_CHAR = '->'
-
 THEME = {
     'comment': lambda s: '\x1b[2;37m{}\x1b[m'.format(s),
     'keyword': lambda s: '\x1b[33;1m{}\x1b[m'.format(s),
@@ -48,13 +38,21 @@ class ExceptionFormatter(object):
         'keywords': [getattr(ast, cls) for cls in dir(ast) if keyword.iskeyword(cls.lower()) and isast(getattr(ast, cls))],
     }
 
-    def __init__(self, colored=SUPPORTS_COLOR, theme=THEME, max_length=MAX_LENGTH,
-                       pipe_char=PIPE_CHAR, cap_char=CAP_CHAR):
+    def __init__(self, colored=SUPPORTS_COLOR, theme=THEME, max_length=MAX_LENGTH, encoding='ascii'):
         self._colored = colored
         self._theme = theme
         self._max_length = max_length
-        self._pipe_char = pipe_char
-        self._cap_char = cap_char
+        self._encoding = encoding
+        self._pipe_char = self._get_char('\u2502', '|')
+        self._cap_char = self._get_char('\u2514', '->')
+
+    def _get_char(self, char, default):
+        try:
+            char.encode(self._encoding)
+        except UnicodeEncodeError:
+            return default
+        else:
+            return char
 
     def colorize_comment(self, source):
         match = self.COMMENT_REGXP.match(source)
@@ -240,7 +238,7 @@ def format_traceback_frame(self, tb):
 
             line += '{}{} {}'.format((' ' * (col - index)), self._cap_char, val)
             lines.append(self._theme['inspect'](line) if self._colored else line)
-        formatted = '\n    '.join([to_unicode(x) for x in lines])
+        formatted = '\n    '.join(lines)
 
         return (filename, lineno, function, formatted), color_source
 

test/output/python-dumb-ascii-color.out

@@ -34,11 +34,11 @@ Traceback (most recent call last):
     div()
     -> <function div at 0xDEADBEEF>
   File "test/test_encoding.py", line 11, in div
-    return _deep('天')
+    return _deep('\u5929')
            -> <function _deep at 0xDEADBEEF>
   File "test/test_encoding.py", line 8, in _deep
     return 1 / val
-               -> '天'
+               -> '\u5929'
 TypeError: unsupported operand type(s) for /: 'int' and 'str'
 
 

test/output/python-dumb-ascii-nocolor.out

@@ -34,11 +34,11 @@ Traceback (most recent call last):
     div()
     -> <function div at 0xDEADBEEF>
   File "test/test_encoding.py", line 11, in div
-    return _deep("天")
+    return _deep("\u5929")
            -> <function _deep at 0xDEADBEEF>
   File "test/test_encoding.py", line 8, in _deep
     return 1 / val
-               -> '天'
+               -> '\u5929'
 TypeError: unsupported operand type(s) for /: 'int' and 'str'
 
 

test/output/python-vt100-ascii-color.out

@@ -34,11 +34,11 @@ Traceback (most recent call last):
     div()
     -> <function div at 0xDEADBEEF>
   File "test/test_encoding.py", line 11, in div
-    return _deep('天')
+    return _deep('\u5929')
            -> <function _deep at 0xDEADBEEF>
   File "test/test_encoding.py", line 8, in _deep
     return 1 / val
-               -> '天'
+               -> '\u5929'
 TypeError: unsupported operand type(s) for /: 'int' and 'str'
 
 

test/output/python-vt100-ascii-nocolor.out

@@ -34,11 +34,11 @@ Traceback (most recent call last):
     div()
     -> <function div at 0xDEADBEEF>
   File "test/test_encoding.py", line 11, in div
-    return _deep("天")
+    return _deep("\u5929")
            -> <function _deep at 0xDEADBEEF>
   File "test/test_encoding.py", line 8, in _deep
     return 1 / val
-               -> '天'
+               -> '\u5929'
 TypeError: unsupported operand type(s) for /: 'int' and 'str'
 
 

test/output/python-xterm-ascii-color.out

@@ -34,11 +34,11 @@ Traceback (most recent call last):
     div()
     -> <function div at 0xDEADBEEF>
   File "test/test_encoding.py", line 11, in div
-    return _deep('天')
+    return _deep('\u5929')
            -> <function _deep at 0xDEADBEEF>
   File "test/test_encoding.py", line 8, in _deep
     return 1 / val
-               -> '天'
+               -> '\u5929'
 TypeError: unsupported operand type(s) for /: 'int' and 'str'
 
 

test/output/python-xterm-ascii-nocolor.out

@@ -34,11 +34,11 @@ Traceback (most recent call last):
     div()
     -> <function div at 0xDEADBEEF>
   File "test/test_encoding.py", line 11, in div
-    return _deep("天")
+    return _deep("\u5929")
            -> <function _deep at 0xDEADBEEF>
   File "test/test_encoding.py", line 8, in _deep
     return 1 / val
-               -> '天'
+               -> '\u5929'
 TypeError: unsupported operand type(s) for /: 'int' and 'str'
 
 

test_all.sh

@@ -52,10 +52,7 @@ for encoding in ascii "UTF-8"; do
 			[[ $color == "1" ]] && color_filename="color" || color_filename="nocolor"
 			filename="test/output/python${pv}-${term}-${encoding}-${color_filename}.out"
 
-			export LANG="en_US.${encoding}"
-			export LC_ALL="${LANG}"
-			export PYTHONCOERCECLOCALE=0
-			export PYTHONUTF8=0
+			export PYTHONIOENCODING="${encoding}"
 			export TERM="${term}"
 			export FORCE_COLOR="${color}"
 			export BETEXC_PYTHON="python"