Generically wrapping exceptions

decorator.py

import contextlib
import functools
from typing import Awaitable, Callable, Coroutine, Dict, Type


ExceptionMapping = Dict[Type[Exception], Type[Exception]]
CoroutineFunction = Callable[..., Awaitable]


@contextlib.contextmanager
def wrap_exceptions(exception_mapping: ExceptionMapping):
    """
    Decorator and context manager.
    Change exception types using provided mapping (keys become values)
    Original exception is available through exc.__cause__.
    New exception is instanciated without arguments.
    
    >>> @wrap_exceptions({ValueError: KeyError})
    ... def myfunc():
    ...     raise ValueError
    
    >>> myfunc()
    KeyError
    
    """
    try:
        yield
    except tuple(exception_mapping) as exc:
        raise exception_mapping[type(exc)] from exc
        
def wrap_exceptions_coro(
    exception_mapping: ExceptionMapping,
) -> Callable[[CoroutineFunction], CoroutineFunction]:
    """
    Decorator, same as wrap_exceptions but for coroutines (wrapping when awaited, not
    when called). Not a context manager.
    
    >>> @wrap_exceptions_coro({ValueError: KeyError})
    ... async def myfunc():
    ...     raise ValueError
    
    >>> myfunc()  # does nothing
    
    >>> await myfunc()
    KeyError:
    """
    def decorator(coro: CoroutineFunction) -> CoroutineFunction:
        @functools.wraps(coro)
        async def wrapper(*args, **kwargs):
            with wrap_exceptions(exception_mapping=exception_mapping):
                return await coro(*args, **kwargs)
        return wrapper
    return decorator