From 87413759cc30204b7c7a1fa556929312c026900a Mon Sep 17 00:00:00 2001
From: Josh <Qix-@users.noreply.github.com>
Date: Wed, 22 Mar 2017 18:41:02 -0700
Subject: [PATCH] use better color detection methods and add color support on
 windows

Makes color detection a little more intuitive and adds proper Windows support via colorama (only installed on Windows).
---
 better_exceptions/__init__.py | 16 ++++----
 better_exceptions/color.py    | 70 +++++++++++++++++++++++++++++++++++
 setup.py                      |  3 ++
 test_color.py                 |  3 ++
 4 files changed, 84 insertions(+), 8 deletions(-)
 create mode 100644 better_exceptions/color.py
 create mode 100644 test_color.py

diff --git a/better_exceptions/__init__.py b/better_exceptions/__init__.py
index 6e9592e..31ba5e0 100644
--- a/better_exceptions/__init__.py
+++ b/better_exceptions/__init__.py
@@ -22,13 +22,13 @@ import re
 import sys
 import traceback
 
+from better_exceptions.color import STREAM, SUPPORTS_COLOR
+
 
 def isast(v):
     return inspect.isclass(v) and issubclass(v, ast.AST)
 
 
-NOCOLOR = not os.isatty(2) or os.name == 'nt' or os.getenv('TERM', '')[:5] != 'xterm'
-
 ENCODING = locale.getpreferredencoding()
 
 PIPE_CHAR = u'\u2502'
@@ -52,7 +52,7 @@ THEME = {
     'keyword': lambda s: '\x1b[33;1m{}\x1b[m'.format(s),
     'builtin': lambda s: '\x1b[35;1m{}\x1b[m'.format(s),
     'literal': lambda s: '\x1b[31m{}\x1b[m'.format(s),
-    'inspect': lambda s: s if NOCOLOR else u'\x1b[36m{}\x1b[m'.format(s),
+    'inspect': lambda s: s if not SUPPORTS_COLOR else u'\x1b[36m{}\x1b[m'.format(s),
 }
 
 MAX_LENGTH = 128
@@ -66,7 +66,7 @@ def colorize_comment(source):
 
 
 def colorize_tree(tree, source):
-    if NOCOLOR:
+    if not SUPPORTS_COLOR:
         # quick fail
         return source
 
@@ -193,13 +193,13 @@ def format_traceback(tb=None):
     return ''.join(lines), final_source
 
 
-def write_stderr(data):
+def write_stream(data):
     data = data.encode(ENCODING)
 
     if sys.version_info[0] < 3:
-        sys.stderr.write(data)
+        STREAM.write(data)
     else:
-        sys.stderr.buffer.write(data)
+        STREAM.buffer.write(data)
 
 
 def excepthook(exc, value, tb):
@@ -210,7 +210,7 @@ def excepthook(exc, value, tb):
     title = traceback.format_exception_only(exc, value)
 
     full_trace = u'Traceback (most recent call last):\n{}{}\n'.format(formatted, title[0].strip())
-    write_stderr(full_trace)
+    write_stream(full_trace)
 
 
 sys.excepthook = excepthook
diff --git a/better_exceptions/color.py b/better_exceptions/color.py
new file mode 100644
index 0000000..f217cfa
--- /dev/null
+++ b/better_exceptions/color.py
@@ -0,0 +1,70 @@
+"""Checks if the current terminal supports colors.
+
+Also specifies the stream to write to. On Windows, this is a wrapped
+stream.
+"""
+
+import errno
+import os
+import struct
+import sys
+
+
+STREAM = sys.stderr
+SUPPORTS_COLOR = False
+
+
+def get_terminfo_file():
+    term = os.getenv('TERM', None)
+
+    if term is None:
+        return None
+
+    terminfo_path = os.path.join('/usr/share/terminfo', ('%0.2X' % ord(term[0])), term)
+    try:
+        f = open(terminfo_path, 'rb')
+    except IOError as e:
+        if e.errno != errno.ENOENT:
+            raise
+
+        terminfo_path = os.path.join('/usr/share/terminfo', term[0], term)
+
+        try:
+            f = open(terminfo_path, 'rb')
+        except IOError as e:
+            if e.errno != errno.ENOENT:
+                raise
+            return None
+
+    return f
+
+
+if os.name == 'nt':
+    from colorama import init as init_colorama, AnsiToWin32
+
+    init_colorama(wrap=False)
+    STREAM = AnsiToWin32(sys.stderr).stream
+    SUPPORTS_COLOR = True
+else:
+    if os.getenv('FORCE_COLOR', None) == '1':
+        SUPPORTS_COLOR = True
+    else:
+        f = get_terminfo_file()
+        if f is not None:
+            with f:
+                 # f is a valid terminfo; seek and read!
+                 magic_number = struct.unpack('<h', f.read(2))[0]
+
+                 if magic_number == 0x11A:
+                     # the opened terminfo file is valid.
+                     offset = 2 + 10  # magic number + size section (the next thing we read from)
+                     offset += struct.unpack('<h', f.read(2))[0]  # skip over names section
+                     offset += struct.unpack('<h', f.read(2))[0]  # skip over bool section
+                     offset += offset % 2  # align to short boundary
+                     offset += 13 * 2  # maxColors is the 13th numeric value
+
+                     f.seek(offset)
+                     max_colors = struct.unpack('<h', f.read(2))[0]
+
+                     if max_colors >= 8:
+                         SUPPORTS_COLOR = True
diff --git a/setup.py b/setup.py
index f62cf8e..b5ccc8b 100644
--- a/setup.py
+++ b/setup.py
@@ -13,4 +13,7 @@ setup(
     download_url = 'https://github.com/qix-/better-exceptions/archive/{}.tar.gz'.format(VERSION),
     keywords = ['pretty', 'better', 'exceptions', 'exception', 'error', 'local', 'debug', 'debugging', 'locals'],
     classifiers = [],
+    extras_require = {
+        ':sys_platform=="win32"': ['colorama']
+    }
 )
diff --git a/test_color.py b/test_color.py
new file mode 100644
index 0000000..83b9f85
--- /dev/null
+++ b/test_color.py
@@ -0,0 +1,3 @@
+from __future__ import print_function
+from better_exceptions import color
+print(color.SUPPORTS_COLOR)