IPython chained exception debugging

chain.py

""" %chain - IPython magic to inspect and move along exception chains.

JD Smith (2022)

Chained exceptions are associated with the messages:

 "During handling of the above exception, another exception occurred"
 "The above exception was the direct cause of the following exception"

and post-mortem debugging (%debug) by default can only examine the
most recent exception.  %chain allows you to inspect the chain, move
along it, and debug any exceptions found.

Note: Alters `sys.last_traceback`. 

Install: simply put this file into the IPython startup directory (e.g.
~/.ipython/profile_default/startup).

Usage: see %chain -h
"""

import sys
from IPython import get_ipython
from IPython.core.magic import register_line_magic
import argparse

def _print(top = 'Exception chain:'):
    col = get_ipython().SyntaxTB.color_scheme_table.active_colors
    depth = 0
    try:
        val = sys.last_value
    except AttributeError:
        print(f"{col.topline}No last exception.{col.Normal}")
        return
    vals = []
    if top: print(f"{col.topline}{top}{col.Normal}")
    while val:
        pref = "--->" if val.__traceback__ is sys.last_traceback else "    "
        vals.append(f"{col.em}{pref}{col.Normal} [{col.lineno}{depth:2d}{col.Normal}] "
                    f"{col.excName}{type(val).__name__}{col.Normal} {val}")
        depth += 1
        val = val.__context__
    print('\n'.join(reversed(vals)))

def _find_chain(n, relative = False):
    try:
        val = sys.last_value
    except AttributeError:
        return
    vals = []
    _end = isinstance(n, str) and n == 'end'
    active, depth = None, 0
    while val:
        if val.__traceback__ is sys.last_traceback: active = depth
        if relative:
            vals.append(val)
            if active is not None and (depth-active) >= n:
                if active + n < 0: return None # No such val
                return vals[active + n]
        elif _end:
            if not val.__context__: return val
        else:  # absolute depth
            if depth == n: return(val)
        val = val.__context__
        depth += 1

_desc="Examine and move along the context/cause chain of the last exception"
parser = argparse.ArgumentParser("%chain", exit_on_error=False, description = _desc)
parser.add_argument("-p", "--print", action="store_true",
                    help="Print the exception chain with current level")
parser.add_argument("-d", "--debug", action="store_true",
                    help="Enter post-mortem debug")
parser.add_argument("depth", nargs='?',
                    help=("Exception chain depth to move to (0 = most recent). "
                          "Prepend with '+' or '-' for a relative move. "
                          "With string `end', move to the end of the exception chain. "
                          "If omitted: move to end, print, and enter debug."))

@register_line_magic
def chain(line):
    try:
        args = parser.parse_args(line.split())
    except SystemExit:
        return

    args.relative = False

    if not line or line.isspace(): # Default: up 1, print, debug
        args.depth = 'end'
        args.print = True
        args.debug = True
    else: 
        if isinstance(args.depth, str) and args.depth != 'end':
            args.relative = args.depth[0] in set('-+')
            args.depth = int(args.depth)

    if args.depth is not None:
        new = _find_chain(args.depth, args.relative)

        if new is None:
            sgn = '+' if (args.relative and args.depth >= 0) else ''
            _print(top=f"Chain does not extend to {sgn}{args.depth}:")
            return
        else: 
            sys.last_traceback = new.__traceback__ # record place

    if args.print: _print()
    if args.debug: get_ipython().run_line_magic('debug','') # run pm-debug