Fixing Mako tracebacks

mako_tb.py

import linecache
import traceback
import sys
import mako.template

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.
    """
    tb = sys.exc_info()[-1]
    while tb is not None:
        f = tb.tb_frame
        co = f.f_code
        filename = co.co_filename
        if filename.startswith('memory:'):
            info = mako.template._get_module_info(filename)
            linecache.cache[filename] = (None, None,
                info.module_source.splitlines(True), filename)
        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()

output.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 40, 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:0x7fd7c0ed32d0", line 28, in render_body
ZeroDivisionError: integer division or modulo by zero

= Improved traceback =

Traceback (most recent call last):
  File "mako_tb.py", line 50, 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:0x276c6d0", line 28, in render_body
    __M_writer(unicode(a / b))
ZeroDivisionError: integer division or modulo by zero

There's an updated version at https://gist.github.com/4269249