exception_handler.py

class _ExceptionHandlerBase:
    def __init__(self, handler_map):
        self._handler_map = handler_map

    def __call__(self, func):
        def wrapper(*args, **kwargs):
            with self:
                func(*args, **kwargs)

        return wrapper

    def __enter__(self):
        pass

    def __exit__(self, exc_class, exc_instance, traceback):
        if exc_class in self._handler_map:
            return self._handler_map[exc_class](exc_instance)
        return False


class ExceptionHandlerMap(_ExceptionHandlerBase):
    def add_handler(self, cls, callback):
        self._handler_map[cls] = callback

    def add_handlers(self, handler_map):
        self._handler_map.update(handler_map)

    def also_handle(self, handler_map):
        new_handler = type(self)(self._handler_map.copy())
        new_handler.add_handlers(handler_map)
        return new_handler


class ExceptionHandler(_ExceptionHandlerBase):
    def __init__(self, callback, *classes):
        super().__init__({cls: callback for cls in classes})
        self._callback = callback
        self._classes = classes

    def add_class(self, cls):
        self._handler_map[cls] = self._callback

    def instead_catch(self, *classes):
        return type(self)(self._callback, *classes)

    def also_catch(self, *classes):
        return self.instead_catch(*(*classes, *self._classes))

    @classmethod
    def catches(cls, *classes):
        def wrapper(callback):
            return cls(callback, *classes)

        return wrapper


# Example usage:
class FooException(Exception):
    pass


class BarException(Exception):
    pass


def handle_foo(e):
    print(f"handle_foo caught {type(e).__name__}: {e}")
    return True


def handle_bar(e):
    print(f"handle_bar caught {type(e).__name__}: {e}")
    return True


# we can use the "catches" class method as a decorator to create an exception handler
@ExceptionHandler.catches(FooException, BarException)
def foo_bar_handler(e):
    print(f"handle_foo_and_bar caught {type(e).__name__}: {e}")
    return True


# we can create an exception handler using the class directly
foo_handler = ExceptionHandler(handle_foo, FooException, BarException)
bar_handler = ExceptionHandler(handle_bar, FooException, BarException)


# we can use the exception handler as a decorator to handle exceptions within the function
@foo_handler
def raise_foo(text):
    raise FooException(text)


@foo_bar_handler
def raise_bar(text):
    raise BarException(text)


raise_foo("in raise_foo function")
raise_bar("in raise_bar function")


# we can also use the exception handler as a context manager to handle exceptions within it's block
with foo_bar_handler:
    raise FooException("inside context manager")

# we can replace which exceptions it catches for this invocation
with foo_handler.instead_catch(BarException):
    raise BarException("This WILL be handled")

# we can append which exceptions it catches for this invocation
with foo_handler.also_catch(BarException):
    raise BarException("This WILL be handled")

# we can create an object that has different callbacks for different exception classes
handler_map = ExceptionHandlerMap({FooException: handle_foo, BarException: handle_bar})

with handler_map:
    raise FooException("this is handled by the ExceptionHandlerMap")

with handler_map:
    raise BarException("this is also handled by the ExceptionHandlerMap")