Rich logging

logging.py

import random
import logging
from rich.logging import RichHandler
from rich.pretty import pretty_repr

class CustomRichHandler(RichHandler):
    def emit(self, record) -> None:
        if not isinstance(record.msg, str):
            record.msg = pretty_repr(record.msg)
        return super().emit(record)

def setup_logging(log_folder):
    import time

    log_folder.mkdir(parents=True, exist_ok=True)
    LOG_FILE_NAME = f"{log_folder}/log_{time.strftime('%Y%m%d%H%M%S', time.localtime())}_{random.randint(0, 10000)}.log"

    logger = logging.getLogger("logger")
    logger.propagate = False
    logger.setLevel(logging.DEBUG)

    # console_handler = logging.StreamHandler()
    console_handler = CustomRichHandler(markup=True, tracebacks_max_frames=4)
    console_handler.setLevel(logging.DEBUG)
    logger.addHandler(console_handler)

    file_handler = logging.FileHandler(LOG_FILE_NAME)
    file_handler.setLevel(logging.DEBUG)
    logger.addHandler(file_handler)

    # Set the log format
    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
    console_handler.setFormatter(formatter)
    file_handler.setFormatter(formatter)

    return LOG_FILE_NAME, logger

def highlighted(c, s):
    return f"[underline bold {c}]{s}[/underline bold {c}]"