From 35eaebff0b01aaff792a7337fb916cdd3d7f024c Mon Sep 17 00:00:00 2001
From: Delgan <adrien.gabillaud@gmail.com>
Date: Mon, 4 Sep 2017 20:32:20 +0200
Subject: [PATCH] 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 | 43 ++++++++---------------------------
 better_exceptions/color.py    | 30 +++++++++++++++++++++++-
 better_exceptions/context.py  |  3 +++
 better_exceptions/encoding.py | 32 ++++++++++++++++++++++++++
 4 files changed, 73 insertions(+), 35 deletions(-)
 create mode 100644 better_exceptions/context.py
 create mode 100644 better_exceptions/encoding.py

diff --git a/better_exceptions/__init__.py b/better_exceptions/__init__.py
index 447a253..00a854d 100644
--- a/better_exceptions/__init__.py
+++ b/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)
 
diff --git a/better_exceptions/color.py b/better_exceptions/color.py
index bb3315b..b4e8f9f 100644
--- a/better_exceptions/color.py
+++ b/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':
diff --git a/better_exceptions/context.py b/better_exceptions/context.py
new file mode 100644
index 0000000..bbd91c3
--- /dev/null
+++ b/better_exceptions/context.py
@@ -0,0 +1,3 @@
+import sys
+
+PY3 = sys.version_info[0] >= 3
diff --git a/better_exceptions/encoding.py b/better_exceptions/encoding.py
new file mode 100644
index 0000000..7b209ac
--- /dev/null
+++ b/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