'''
annex - ANNotated EXception tracebacks.

This module provides an excepthook which annotates each source line with
the relevant reference values.  This can reduce time spent understanding
exceptions.

Example:

import sys
import annex
sys.excepthook = annex.annotated_excepthook

# ... code which may generate exceptions.

'''

import parser, token, keyword, traceback, linecache


def annotated_excepthook(etype, evalue, etb):
    '''Use this as sys.excepthook to get annotated exceptions.'''
    # Lot's of this code cut'n'pasted from traceback.py
    print 'Traceback (most recent call last):'
    while etb is not None:
        f = etb.tb_frame
        lineno = etb.tb_lineno
        co = f.f_code
        filename = co.co_filename
        name = co.co_name
        print '  File "%s", line %d, in %s' % (filename, lineno, name)
        linecache.checkcache(filename)
        line = linecache.getline(filename, lineno)
        if line:
            print '    ' + line.strip()
            scopes = Scopes(f.f_locals, f.f_globals, f.f_builtins)
            varnames = match_references(line, co.co_names)
            if varnames:
                print '    # With bindings:'
                for varname in varnames:
                    #if varname.startswith('_['):
                    #    continue
                    print '    # %s = %r' % (varname, scopes[varname])
        
        etb = etb.tb_next

    lines = traceback.format_exception_only(etype, evalue)
    for line in lines:
        print line
        
def match_references(expr, names):
    refs = find_all_names(expr)
    refs = refs.intersection(names)
    return sorted(refs)

def find_all_names(expr, _names=None):
    ast = idempotent_parse(expr)
    if _names is None:
        _names = set()

    nodetype = ast[0]
    if nodetype == token.NAME:
        name = ast[1]
        if not keyword.iskeyword(name):
            _names.add(name)
    elif token.ISNONTERMINAL(nodetype):
        [find_all_names(c, _names) for c in ast[1:]]
    return _names

def idempotent_parse(source):
    if type(source) is str:
        source = parser.suite(source.lstrip())

    if type(source) is parser.ASTType:
        source = source.totuple(True)

    return source


class Scopes (object):
    def __init__(self, *scopes):
        self.scopes = scopes

    def __getitem__(self, key):
        for scope in self.scopes:
            try:
                return scope[key]
            except KeyError:
                pass
        raise KeyError(key)