SageRescue

sagerescue.py

import sys, collections
try:
    import rich.traceback, rich.console
    console = rich.console.Console(file=sys.stderr)
    def nice_traceback(_app, t, v, tb):
        console.print(rich.traceback.Traceback.from_exception(t, v, tb))
except ImportError:
    def nice_traceback(app, t, v, tb):
        app.shell.excepthook(t, v, tb)

class SageRescueScope(collections.ChainMap):
    def __init__(self, app, typ, val, tb):
        self.__app = app
        self.__typ = typ
        self.__val = val
        self.__tb = tb
        super().__init__()

    def add_child(self, m):
        self.maps.insert(0, m)

    def __getattr__(self, name):
        '''Fallback method to allow quicker access to the variables being exposed'''
        return self[name]

    def tb(self):
        nice_traceback(self.__app, self.__typ, self.__val, self.__tb)


def install():
    def hook(old):
        def f(t, v, tb):
            ## Restore sanity, avoid recursion
            sys.excepthook = old

            ## Create the repl
            from sage.repl.interpreter import SageTerminalApp
            app = SageTerminalApp.instance()

            ## Walk up the stack
            scope = SageRescueScope(app, t, v, tb)
            walktb = tb
            while walktb is not None:
                scope.add_child(walktb.tb_frame.f_globals)
                scope.add_child(walktb.tb_frame.f_locals)
                walktb = walktb.tb_next

            ## Setup the app and add our variable
            app.initialize()
            app.shell.user_module.__dict__.update({"scope": scope})
            ## Show the traceback
            scope.tb()
            app.start()
        return f

    sys.excepthook = hook(sys.excepthook)