logging with colors and tqdm progress bar + pickle + ctrl-c handler

logutils.py

import json
import sys
import os
import logging
import argparse
import time
import requests
import re

from urllib.parse import urlparse
from dataclasses import dataclass
from typing import List, Dict
from enum import Enum
import tqdm

import datetime

import urllib3
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
import pickle
import signal
saved_state = None

def handler(signum, frame):

    global saved_state

    msg = "Ctrl-c was pressed. Do you really want to exit? y/n "
    print(msg, end="", flush=True)
    res = input()
    if res.lower() == 'y':
        with open("saved_state.pickle", "wb") as f:
            pickle.dump(saved_state, f)
        logger.ok("State successfully saved @ saved_state.pickle")
        sys.exit(0)

    else:
        print("", end="\r", flush=True)
        print(" " * len(msg), end="", flush=True)  # clear the printed line
        print("    ", end="\r", flush=True)


signal.signal(signal.SIGINT, handler)


def ok(msg, *args, **kwargs):
    logger.log(15, args[0])

logging.addLevelName(15, "OK")
logging.ok = ok
logging.Logger.ok = ok



class CustomFormatter(logging.Formatter):

    grey = '\x1b[38;21m'
    blue = '\x1b[38;5;39m'
    yellow = '\x1b[38;5;226m'
    red = '\x1b[38;5;196m'
    bold_red = '\x1b[31;1m'
    reset = '\x1b[0m'
    okgreen = '\033[92m'

    def __init__(self, fmt):
        super().__init__()
        self.fmt = fmt
        self.FORMATS = {
            logging.DEBUG: self.grey + self.fmt + self.reset,
            logging.INFO: self.blue + self.fmt + self.reset,
            logging.WARNING: self.yellow + self.fmt + self.reset,
            logging.ERROR: self.red + self.fmt + self.reset,
            logging.CRITICAL: self.bold_red + self.fmt + self.reset,
            15: self.okgreen + self.fmt + self.reset
        }

    def format(self, record):
        log_fmt = self.FORMATS.get(record.levelno)
        formatter = logging.Formatter(log_fmt)
        return formatter.format(record)

class TqdmLoggingHandler(logging.Handler):
    def __init__(self, level=logging.NOTSET):
        super().__init__(level)

    def emit(self, record):
        try:
            msg = self.format(record)
            tqdm.tqdm.write(msg)
            self.flush()
        except Exception:
            self.handleError(record)
# Create custom logger logging all five levels
logger = logging.getLogger(__name__)
logger.setLevel(5)

# Define format for logs
fmt = '%(asctime)s | %(levelname)8s | [%(filename)s:%(lineno)3d] %(funcName)s() | %(message)s'

# Create stdout handler for logging to the console (logs all five levels)
stdout_handler = logging.StreamHandler()
stdout_handler.setLevel(logging.DEBUG)
stdout_handler.setFormatter(CustomFormatter(fmt))

# Create file handler for logging to a file (logs all five levels)
today = datetime.date.today()
file_handler = logging.FileHandler('test_{}.log'.format(today.strftime('%Y_%m_%d')))
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(logging.Formatter(fmt))

# Add both handlers to the logger
#logger.addHandler(stdout_handler)
logger.addHandler(file_handler)
tqdm_fmtter = TqdmLoggingHandler()
tqdm_fmtter.setFormatter(CustomFormatter(fmt))
tqdm_fmtter.setLevel(logging.DEBUG)
logger.addHandler(tqdm_fmtter)