import logging
import colorama


class NoTracebackStreamHandler(logging.StreamHandler):
    def handle(self, record):
        info, cache = record.exc_info, record.exc_text
        record.exc_info, record.exc_text = None, None

        if record.levelname in ['WARNING', 'WARN']:
            color = colorama.Fore.LIGHTYELLOW_EX
        elif record.levelname == 'ERROR':
            color = colorama.Fore.LIGHTRED_EX
        elif record.levelname == 'INFO':
            color = colorama.Fore.LIGHTGREEN_EX
        elif record.levelname == 'CRITICAL':
            color = colorama.Fore.LIGHTCYAN_EX
        else:
            color = colorama.Fore.WHITE

        try:
            print(color, end='', flush=True)
            super().handle(record)
        finally:
            record.exc_info = info
            record.exc_text = cache
            print(colorama.Style.RESET_ALL, end='', flush=True)