Coloring exceptions from Python on a terminal

Question

Is there an easy way to get the message of the exception to be colored on the command line? For example

def g():    f()
def f():    1/0
g()

Gives the error

Traceback (most recent call last):
  File "test.py", line 3, in <module>
    g()
  File "test.py", line 1, in g
    def g():    f()
  File "test.py", line 2, in f
    def f():    1/0
ZeroDivisionError: integer division or modulo by zero

I would like "integer division or modulo by zero" to be colored or highlighted on the terminal so that I can quickly pick it out of a long traceback (Linux only). Ideally, I wouldn't want to write a custom class for each Exception, but somehow catch and format all kinds.

EDIT: The question linked in the comments gives examples on how to solve the problem with external software, but I'm interested in an internal Python solution.

Answer 1

You can assign a custom function to the sys.excepthook handler. The function is called whenever there is a unhandled exception (so one that exits the interpreter).

def set_highlighted_excepthook():
    import sys, traceback
    from pygments import highlight
    from pygments.lexers import get_lexer_by_name
    from pygments.formatters import TerminalFormatter

    lexer = get_lexer_by_name("pytb" if sys.version_info.major < 3 else "py3tb")
    formatter = TerminalFormatter()

    def myexcepthook(type, value, tb):
        tbtext = ''.join(traceback.format_exception(type, value, tb))
        sys.stderr.write(highlight(tbtext, lexer, formatter))

    sys.excepthook = myexcepthook

set_highlighted_excepthook()

This version uses the pygments library to convert the traceback text into one formatted with ANSI coloring, before writing it to stderr.

Someone turned this into a project that detects terminal support and lets you set the pygments style, see colored-traceback.py.

Answer 2

Found another way to do this using the IPython module which is likely a dependency that everyone already has installed:

from IPython.core.ultratb import ColorTB
c = ColorTB()
exc = sys.exc_info()
print(''.join(c.structured_traceback(*exc)))