samgis_lisa_on_cuda/utilities/session_logger.py

import contextvars
import logging
from functools import wraps
from typing import Callable, Tuple

logging_uuid = contextvars.ContextVar("uuid")
default_formatter = '%(asctime)s | %(uuid)s [%(pathname)s:%(module)s %(lineno)d] %(levelname)s | %(message)s'


loggingType = logging.CRITICAL | logging.ERROR | logging.WARNING | logging.INFO | logging.DEBUG


def setup_logging(
        debug: bool = False, formatter: str = default_formatter, name: str = "logger"
) -> Tuple[logging, contextvars.ContextVar]:
    """
    Create a logging instance with log string formatter.

    Args:
        debug: logging debug argument
        formatter: log string formatter
        name: logger name

    Returns:
        Logger

    """

    old_factory = logging.getLogRecordFactory()

    def record_factory(*args, **kwargs):
        record = old_factory(*args, **kwargs)
        record.uuid = logging_uuid.get("uuid")
        if isinstance(record.msg, str):
            record.msg = record.msg.replace("\\", "\\\\").replace("\n", "\\n")
        return record

    logging.setLogRecordFactory(record_factory)
    logging.basicConfig(level=logging.DEBUG, format=default_formatter, force=True)

    logger = logging.getLogger(name=name)

    # create a console handler
    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)

    # create formatter and add to the console
    formatter = logging.Formatter(formatter)
    ch.setFormatter(formatter)

    # add the console handler to logger
    logger.addHandler(ch)
    return logger, logging_uuid


def set_uuid_logging(func: Callable) -> Callable:
    @wraps(func)
    def wrapper(*args, **kwargs):
        import uuid

        current_uuid = f"{uuid.uuid4()}"
        logging_uuid.set(current_uuid)
        return func(*args, **kwargs)

    return wrapper