Pop IPython embed shell when any exception is raised (before ei)

ipydbg.py

import sys
import inspect
import traceback
import importlib
import readline
import contextlib


__all__ = [
    'probe',
    'DebugHook',
    'patch',
    'unpatch',
    'capture',
    'embed'
]


assert importlib.util.find_spec('IPython') is not None, (
        'IPython is required for this module')


_original_excepthook = None


def probe():
    from IPython import embed
    import IPython.core.ultratb
    
    print('IPython debug hook is available!')
    
    for name, value in IPython.core.ultratb.__dict__.items():
        if not isinstance(value, type):
            continue
        if value is IPython.core.ultratb.TBTools:
            continue
        if not issubclass(value, IPython.core.ultratb.TBTools):
            continue
        
        tb = value()
        modes = getattr(tb, 'valid_modes', None)
        colors = list(getattr(tb, 'color_scheme_table').keys()) \
                if hasattr(tb, 'color_scheme_table') else None
                
        print('  class: {}'.format(name))
        print('    modes: {}'.format(modes))
        print('    colors: {}'.format(colors))


class DebugHook:
    def __init__(self, select=False, tb_class='AutoFormattedTB', tb_mode='Context', tb_color='Neutral'):
        self.select = select
        self.tb_class = tb_class
        self.tb_mode = tb_mode
        self.tb_color = tb_color
    
    def __call__(self, exc_type, exc, tb):
        from IPython import embed
        import IPython.core.ultratb
        
        TB = getattr(IPython.core.ultratb, self.tb_class)
        tb_handler = TB(mode=self.tb_mode, color_scheme=self.tb_color)
        tb_handler(exc_type, exc, tb)
        
        frames = [t[0] for t in traceback.walk_tb(tb)][::-1]
        
        while True:
            print()
            
            if self.select:
                print('Select a stack frame to embed IPython shell (default: 0):')
                for i, frame in enumerate(frames):
                    print('{}. {}'.format(i, frame))
                try:
                    s = input('> ').strip()
                    n = int(s) if s else 0
                    frame = frames[n]
                except (KeyboardInterrupt, EOFError):
                    break
                except:
                    continue
            else:
                frame = frames[0]

            print('Embedded into', frame)
            
            user_module = inspect.getmodule(frame)
            user_ns = frame.f_locals
            
            user_ns.setdefault('etype', exc_type)
            user_ns.setdefault('evalue', exc)
            user_ns.setdefault('etb', tb)
        
            embed(banner1='', user_module=user_module, user_ns=user_ns, colors='Neutral')
            
            if not self.select:
                break


def patch(**kwargs):
    global _original_excepthook
    _original_excepthook = sys.excepthook
    sys.excepthook = DebugHook(**kwargs)


def unpatch():
    global _original_excepthook
    assert _original_excepthook is not None, 'not patched'
    sys.excepthook = _original_excepthook
    _original_excepthook = None


@contextlib.contextmanager
def capture(**kwargs):
    patch(**kwargs)
    yield
    unpatch()


def embed(message=None, exit=False):
    from IPython import embed
    
    e = sys.exc_info()[1]
    if e is not None and message is None:
        message = traceback.format_exc()
        
    frame = inspect.stack()[1].frame
    user_ns = frame.f_locals
    user_module = inspect.getmodule(frame)
    
    embed(header=message, user_ns=user_ns, user_module=user_module)
    
    if exit:
        sys.exit()

Install

wget https://gist.githubusercontent.com/djosix/545af4d7eb4ec0417c3a1dccecfc19b6/raw/15724a6a584c0b85d8f6d66152f5b058f8e648f4/ipydbg.py

Example 1

def test():
    a = 0
    return 1 / a

import ipydbg
ipydbg.patch(select=True):

test()

Example 2

def test():
    a = 0
    return 1 / a

import ipydbg
with ipydbg.capture(select=True):
    test()

Example 3

import ipydbg

def main():
    try:
        a = 1
        b = 0
        c = a / b
    except:
        ipydbg.embed(exit=True)
        
    print(a, b, c)

main()