Fixing Mako tracebacks (version 2)

demo.txt

$ virtualenv /tmp/sandbox
...

$ /tmp/sandbox/bin/pip install mako
...

$ /tmp/sandbox/bin/python mako_tb.py 

= Default traceback =

Traceback (most recent call last):
  File "mako_tb.py", line 53, in <module>
    print t.render()
  File "/tmp/sandbox/local/lib/python2.7/site-packages/mako/template.py", line 412, in render
    return runtime._render(self, self.callable_, args, data)
  File "/tmp/sandbox/local/lib/python2.7/site-packages/mako/runtime.py", line 766, in _render
    **_kwargs_for_callable(callable_, data))
  File "/tmp/sandbox/local/lib/python2.7/site-packages/mako/runtime.py", line 798, in _render_context
    _exec_template(inherit, lclcontext, args=args, kwargs=kwargs)
  File "/tmp/sandbox/local/lib/python2.7/site-packages/mako/runtime.py", line 824, in _exec_template
    callable_(context, *args, **kwargs)
  File "memory:0x7f6998e352d0", line 28, in render_body
ZeroDivisionError: integer division or modulo by zero

= Improved traceback =

Traceback (most recent call last):
  File "mako_tb.py", line 63, in <module>
    print t.render()
  File "/tmp/sandbox/local/lib/python2.7/site-packages/mako/template.py", line 412, in render
    return runtime._render(self, self.callable_, args, data)
  File "/tmp/sandbox/local/lib/python2.7/site-packages/mako/runtime.py", line 766, in _render
    **_kwargs_for_callable(callable_, data))
  File "/tmp/sandbox/local/lib/python2.7/site-packages/mako/runtime.py", line 798, in _render_context
    _exec_template(inherit, lclcontext, args=args, kwargs=kwargs)
  File "/tmp/sandbox/local/lib/python2.7/site-packages/mako/runtime.py", line 819, in _exec_template
    _render_error(template, context, e)
  File "/tmp/sandbox/local/lib/python2.7/site-packages/mako/runtime.py", line 828, in _render_error
    result = template.error_handler(context, error)
  File "/tmp/sandbox/local/lib/python2.7/site-packages/mako/runtime.py", line 817, in _exec_template
    callable_(context, *args, **kwargs)
  File "memory:0x1962750", line 28, in render_body
    __M_writer(unicode(a / b))
    # memory:0x1962750 line 4 in render_body:
    # a / b = ${a / b}
ZeroDivisionError: integer division or modulo by zero

mako_tb.py

import linecache
import traceback
import sys
import mako.template
import mako.exceptions

def mako_error_handler(context, error):
    """Decorate tracebacks when Mako errors happen.

    Evil hack: walk the traceback frames, find compiled Mako templates,
    stuff their (transformed) source into linecache.cache.
    """
    rich_tb = mako.exceptions.RichTraceback(error)
    rich_iter = iter(rich_tb.traceback)
    tb = sys.exc_info()[-1]
    source = {}
    annotated = set()
    while tb is not None:
        cur_rich = next(rich_iter)
        f = tb.tb_frame
        co = f.f_code
        filename = co.co_filename
        lineno = tb.tb_lineno
        if filename.startswith('memory:'):
            lines = source.get(filename)
            if lines is None:
                info = mako.template._get_module_info(filename)
                lines = source[filename] = info.module_source.splitlines(True)
                linecache.cache[filename] = (None, None, lines, filename)
            if (filename, lineno) not in annotated:
                annotated.add((filename, lineno))
                extra = '    # {} line {} in {}:\n    # {}'.format(*cur_rich)
                lines[lineno-1] += extra
        tb = tb.tb_next
    # Don't return False -- that will lose the actual Mako frame.  Instead
    # re-raise.
    raise

# demo

template_source = """
<% a = 10; b = 0 %>

a / b = ${a / b}
"""

print
print "= Default traceback ="
print

t = mako.template.Template(template_source)
try:
    print t.render()
except:
    traceback.print_exc()

print
print "= Improved traceback ="
print

t = mako.template.Template(template_source, error_handler=mako_error_handler)
try:
    print t.render()
except:
    # NB: while traceback.print_exc() picks up the module source from
    # linecache.cache correctly, if you just let this exception propagate to
    # the top level, it'll end up in some Python internal traceback printer
    # which apparently ignores linecache
    traceback.print_exc()