traceback.py

import sys
from pygments.style import Style
from pygments.token import *
from pygments.formatter import Formatter
from curtsies.bpythonparse import parse

default_colors = {
        'keyword': 'y',
        'name': 'c',
        'comment': 'b',
        'string': 'm',
        'error': 'r',
        'number': 'G',
        'operator': 'Y',
        'punctuation': 'y',
        'token': 'C',
        'background': 'd',
        'output': 'w',
        'main': 'c',
        'paren': 'R',
        'prompt': 'c',
        'prompt_more': 'g',
    }

Parenthesis = Token.Punctuation.Parenthesis
theme_map = {
    Keyword: 'keyword',
    Name: 'name',
    Comment: 'comment',
    String: 'string',
    Literal: 'string',
    Error: 'error',
    Number: 'number',
    Token.Literal.Number.Float: 'number',
    Operator: 'operator',
    Punctuation: 'punctuation',
    Token: 'token',
    Whitespace: 'background',
    Parenthesis: 'paren',
    Parenthesis.UnderCursor: 'operator'}


class BPythonFormatter(Formatter):
    """This is the custom formatter for bpython.
    Its format() method receives the tokensource
    and outfile params passed to it from the
    Pygments highlight() method and slops
    them into the appropriate format string
    as defined above, then writes to the outfile
    object the final formatted string.

    See the Pygments source for more info; it's pretty
    straightforward."""

    def __init__(self, color_scheme, **options):
        self.f_strings = {}
        for k, v in theme_map.iteritems():
            self.f_strings[k] = '\x01%s' % (color_scheme[v],)
            if k is Parenthesis:
                # FIXME: Find a way to make this the inverse of the current
                # background colour
                self.f_strings[k] += 'I'
        Formatter.__init__(self, **options)

    def format(self, tokensource, outfile):
        o = ''
        for token, text in tokensource:
            while token not in self.f_strings:
                token = token.parent
            o += "%s\x03%s\x04" % (self.f_strings[token], text)
        return o


def myexcepthook(type, value, tb):
    import traceback
    from pygments import highlight
    from pygments.lexers import get_lexer_by_name
    from pygments.formatters import Terminal256Formatter
    from pygments.styles import get_style_by_name

    tbtext = ''.join(traceback.format_exception(type, value, tb))
    lexer = get_lexer_by_name("pytb", stripall=True)
    formatter = BPythonFormatter(default_colors)
    #print tbtext
    tokens = list(lexer.get_tokens(tbtext))
    format_string = formatter.format(tokens, sys.stderr)
    print repr(parse(format_string))
    print parse(format_string)


def g(): f()

def f(): 1/0

sys.excepthook = myexcepthook
g()

#colorscheme={Name.Builtin:("red","blue")}

#['monokai', 'manni', 'rrt', 'perldoc', 'borland', 'colorful', 'default', 'murphy', 'vs', 'trac', 'tango', 'fruity', 'autumn', 'bw', 'emacs', 'vim', 'pastie', 'friendly', 'native']