Fix raw codes appearing with Python 3 on Windows cmd

Colorama does not wrap the 'buffer' attributes of streams, so writing
to it would not convert ANSI codes on Windows.
The workaround is to use '.write()' without encoding the string,
and rather wrap the stream used by Colorama to encode substrings
which are sent once win32 call are made.

better_exceptions/__init__.py

@@ -17,15 +17,15 @@ import ast
 import inspect
 import keyword
 import linecache
-import locale
 import logging
 import os
 import re
 import sys
 import traceback
-import codecs
 
-from .color import STREAM, SUPPORTS_COLOR
+from .color import STREAM, SUPPORTS_COLOR, SHOULD_ENCODE
+from .context import PY3
+from .encoding import ENCODING, to_byte as _byte, to_unicode as _unicode
 from .log import BetExcLogger, patch as patch_logging
 from .repl import interact, get_repl
 
@@ -37,8 +37,6 @@ def isast(v):
     return inspect.isclass(v) and issubclass(v, ast.AST)
 
 
-ENCODING = locale.getpreferredencoding()
-
 PIPE_CHAR = u'\u2502'
 CAP_CHAR = u'\u2514'
 
@@ -66,32 +64,6 @@ THEME = {
 
 MAX_LENGTH = 128
 
-PY3 = sys.version_info[0] >= 3
-
-
-def _byte(val):
-    unicode_type = str if PY3 else unicode
-    if isinstance(val, unicode_type):
-        try:
-            return val.encode(ENCODING)
-        except UnicodeEncodeError:
-            if PY3:
-                return codecs.escape_decode(val)[0]
-            else:
-                return val.encode("unicode-escape").decode("string-escape")
-
-    return val
-
-
-def _unicode(val):
-    if isinstance(val, bytes):
-        try:
-            return val.decode(ENCODING)
-        except UnicodeDecodeError:
-            return val.decode("unicode-escape")
-
-    return val
-
 
 def colorize_comment(source):
     match = COMMENT_REGXP.match(source)
@@ -320,10 +292,13 @@ def format_traceback(tb=None):
 
 
 def write_stream(data):
-    data = _byte(data)
+    if SHOULD_ENCODE:
+        data = _byte(data)
 
-    if PY3:
-        STREAM.buffer.write(data)
+        if PY3:
+            STREAM.buffer.write(data)
+        else:
+            STREAM.write(data)
     else:
         STREAM.write(data)
 

better_exceptions/color.py

@@ -9,8 +9,12 @@ import os
 import struct
 import sys
 
+from .context import PY3
+from .encoding import to_byte as _byte
+
 
 STREAM = sys.stderr
+SHOULD_ENCODE = True
 SUPPORTS_COLOR = False
 
 
@@ -50,11 +54,35 @@ 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 = AnsiToWin32(sys.stderr).stream
+
+    stream = sys.stderr
+
+    if PY3:
+        # 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
 else:
     if os.getenv('FORCE_COLOR', None) == '1':

better_exceptions/context.py

@@ -0,0 +1,3 @@
+import sys
+
+PY3 = sys.version_info[0] >= 3

better_exceptions/encoding.py

@@ -0,0 +1,32 @@
+import codecs
+import locale
+import sys
+
+from .context import PY3
+
+
+ENCODING = locale.getpreferredencoding()
+
+
+def to_byte(val):
+    unicode_type = str if PY3 else unicode
+    if isinstance(val, unicode_type):
+        try:
+            return val.encode(ENCODING)
+        except UnicodeEncodeError:
+            if PY3:
+                return codecs.escape_decode(val)[0]
+            else:
+                return val.encode("unicode-escape").decode("string-escape")
+
+    return val
+
+
+def to_unicode(val):
+    if isinstance(val, bytes):
+        try:
+            return val.decode(ENCODING)
+        except UnicodeDecodeError:
+            return val.decode("unicode-escape")
+
+    return val