# bettertb.by written by Nicholas FitzRoy-Dale, Jul 10, 2010
# Public domain code. Do what you want with it.
# Email me for details (nicholas, lardcave, net)
# Features: 

# 1) Only print code lines for last three TB elements
# 2) Colourise
# 3) Columnised output: function name, line number, file name

# Syntax highlight if we can
try:
	import pygments
	import pygments.lexers
	import pygments.formatters
except ImportError:
	pygments = None

import traceback
import sys
import os

FILTER_OUT_CWD = True
BRIEF_CUTOFF = 5
VERBOSE_CUTOFF = 3

cwd = None

class Colours(object):
	PREFIX  = '\x1b['

	# Red: 1, Green: 2, Yellow: 3, Blue: 4, Magenta: 5, Cyan: 6, White: 7
	banner    = PREFIX + '31m'
	briefwarn = PREFIX + '31m'
	lineno    = PREFIX + '30m'
	fn        = PREFIX + '34m'
	func      = PREFIX + '30m'
	reset     = PREFIX + '0m'
	exception = PREFIX + '30m'

def syntax_highlight(text):
	if pygments:
		text = pygments.highlight(
				text,
				pygments.lexers.PythonLexer(),
				pygments.formatters.TerminalFormatter())
		text = text.strip()

	return text

def filter_out_cwd(fn):
	global cwd

	if FILTER_OUT_CWD and fn.startswith(cwd):
		fn = fn[len(cwd):]
		if fn.startswith(os.path.sep):
			fn = fn[1:]

	return fn

def clean_fn(fn):
	# Remove extension
	fn = fn.rsplit('.', 1)[0]

	return fn

def print_tb(exc, tb):
	global cwd

	w = sys.stdout.write
	c = Colours
	cwd = os.getcwd()

	w(c.banner + 'Traceback' + c.reset)

	tb_lines = traceback.extract_tb(tb)
	if VERBOSE_CUTOFF:
		verbose_start = len(tb_lines) - VERBOSE_CUTOFF
	else:
		verbose_start = 0
	
	if BRIEF_CUTOFF:
		brief_start = len(tb_lines) - BRIEF_CUTOFF
		w(c.briefwarn)
		w(' (%d frames omitted)' % (brief_start))
	else:
		brief_start = 0

	w(c.banner + ':\n' + c.reset)
	count = 0
	for fn, lineno, func, text in traceback.extract_tb(tb):
		if(count >= brief_start):
			fn = clean_fn(filter_out_cwd(fn))

			w(c.lineno)
			w('% 4d ' % (lineno))
			w(c.fn)
			w('%-30s ' % (fn))
			w(c.func)
			w('%s\n' % (func))

			if(count >= verbose_start):
				w('     %s\n\n' % (syntax_highlight(text)))

		count += 1
	
	w(c.exception)
	exc_str = str(exc)
	if exc_str:
		w("%s: %s\n" % (exc.__class__.__name__, exc))
	else:
		w("%s\n" % (exc.__class__.__name__))

	w(c.reset)

def run(func, *args, **kwargs):
	try:
		result = func(*args, **kwargs)
	except:
		info = sys.exc_info()
		print_tb(info[1], info[2])
		sys.exit(1)
	else:
		return result