Traceback logging decorator

traceback_decorator.py

import logging
import traceback
from typing import Callable


def report_failure_traceback(logger: logging.Logger, exc: Exception) -> None:
    logger.critical('*** Unexpected error ***')
    logger.critical('')
    lines = traceback.format_tb(exc.__traceback__)
    lines = ''.join(lines).split('\n')
    for line in lines:
        logger.critical(line)
    logger.critical(f'{exc.__class__.__name__}: {exc}')


def with_failure_traceback(logger: logging.Logger, propagate: bool = False) -> Callable:
    """
    If function raises an un-caught exception, this decorator sends into logger 
    a failure traceback using the exception's information. In case if 
    KeyboardInterrupt exception is received, only a warning about interruption 
    is printed.
    """
    def wrapper(f):
        def wrapped(*args, **kwargs):
            try:
                return f(*args, **kwargs)
            except KeyboardInterrupt:
                logger.warning('Keyboard interruption was sent. Terminating...')
            except Exception as e:
                report_failure_traceback(logger, e)
                if propagate:
                    raise RuntimeError(str(e))

        return wrapped

    return wrapper


logging.basicConfig()
logger = logging.getLogger()


@with_failure_traceback(logger)
def some_complex_logic(s: str) -> str:
    if s == 'error':
        raise ValueError(s)
    return s.upper()


def main():
    assert some_complex_logic('ok') == 'OK'
    assert some_complex_logic('error') is None


if __name__ == '__main__':
    main()