logger.py

"""
Provides a logger class for logging messages to both console and file.
"""

from __future__ import annotations

import os, sys
import datetime
import uuid
import time
import logging
import argparse
import shutil
import typing

import common_util.decorator as deco
import common_util.env_variable as env
import common_util.colored_print as cprint
import common_util.python_util as putil

_logger_argparser = argparse.ArgumentParser()

_logger_argparser.add_argument(
    "--log_dir",
    type=str,
    nargs="?",
    default=argparse.SUPPRESS,
    help="Path to dir that stores log files",
)

_logger_argparser.add_argument(
    "--create_log_dir",
    action="store_true",
    default=False,
    help="Create log directory if it does not exist",
)

_logger_args, _ = _logger_argparser.parse_known_args()
_logger_default_log_dir = os.path.abspath(os.path.join(os.getcwd(), "log"))


class LoggingCustomStreamFormatter(logging.Formatter):
    def __init__(
        self,
        fmt=None,
        datefmt=None,
        style: logging._FormatStyle = '%',
        validate=True,
        *,
        defaults=None,
    ):
        super().__init__(fmt=fmt, datefmt=datefmt, style=style, validate=validate)
        self.__formats = {
            level: f"{cprint.ColoredPrintSetting.MSG_COLOR_DICT[level]}{fmt}{cprint.ANSIColors.ENDC}"
            for level in cprint.ColoredPrintSetting.MSG_COLOR_DICT.keys()
        }

    def format(self, record: logging.LogRecord) -> str:
        if record.levelno in self.__formats:
            self._style._fmt = self.__formats[record.levelno]
        return super().format(record)


@deco.singleton
class Logger:
    """
    Wrapper of a two-level hierarchical logging.Logger
    """

    def __init__(self) -> None:
        # default log entry and log file time format
        self.__log_time_format = r"%Y-%m-%d %H:%M:%S %z"
        self.__dir_time_format = r"%Y-%m-%dT%H:%M:%S%z"

        # default log entry format for file logger and stderr logger
        self.__log_file_format = "[%(asctime)s](%(filename)s:%(lineno)d) %(levelname)s %(message)s"
        self.__log_stderr_format = "%(levelname)s %(message)s"

        # initialize the root logger for all modules using this utils
        self.__default_logger = logging.getLogger("Utils Root Logger")
        if hasattr(_logger_args, "log_dir"):
            logging_folder = putil.safeval(_logger_args.log_dir, _logger_default_log_dir)

            if not os.path.isdir(logging_folder):
                if not _logger_args.create_log_dir:
                    cprint.eprintf(
                        f"{self.__default_logger.name} - "
                        f"Logger directory creation disabled with "
                        f"target directory {logging_folder} missing, abort",
                        file=sys.stderr,
                    )
                    exit(1)
                os.makedirs(logging_folder)

            dirname = datetime.datetime.now().astimezone().strftime(self.__dir_time_format)

            disambiguated_dirname = dirname
            disambiguated_abspath = os.path.join(logging_folder, dirname)
            while os.path.isdir(disambiguated_abspath):
                disambiguated_dirname = f"{dirname}.{uuid.uuid4()}"
                disambiguated_abspath = os.path.join(logging_folder, disambiguated_dirname)
            self.__log_folder = disambiguated_dirname
            self.__log_name = "python_rt.log"
            self.__log_dirpath = os.path.realpath(disambiguated_abspath)
            self.__log_path = os.path.realpath(os.path.join(disambiguated_abspath, self.__log_name))
            os.makedirs(os.path.join(logging_folder, self.__log_folder))

            # set logging file handler and format
            cur_root_handlers = logging.root.handlers
            assert len(cur_root_handlers) == 1
            cur_root_handlers[0].setFormatter(
                LoggingCustomStreamFormatter(
                    fmt=self.__log_stderr_format,
                    datefmt=self.__log_time_format,
                )
            )

            # add a file handler on top of the default stream handler
            handler = logging.FileHandler(
                filename=self.__log_path,
                mode="w",
                delay=True,
            )
            handler.setFormatter(
                logging.Formatter(
                    fmt=self.__log_file_format,
                    datefmt=self.__log_time_format,
                )
            )
            logging.root.addHandler(handler)
        else:
            self.__log_folder = ""
            self.__log_name = ""
            self.__log_dirpath = ""
            self.__log_path = ""

            stderr_handler = logging.StreamHandler(sys.stderr)
            stderr_handler.setFormatter(
                LoggingCustomStreamFormatter(
                    fmt=self.__log_stderr_format,
                    datefmt=self.__log_time_format,
                )
            )
            self.__default_logger.addHandler(stderr_handler)

        # register this component and a default logger
        self.__default_logger.setLevel(logging.DEBUG if env.is_debug() else logging.WARNING)
        self.__registered_logger_names = set()

    @property
    def default_logging_level(self) -> int:
        return self.__default_logger.level

    @property
    def default_logging_level_name(self) -> str:
        return logging.getLevelName(self.__default_logger.level)

    @property
    def log_folder(self) -> str:
        return self.__log_folder

    @property
    def log_filename(self) -> str:
        return self.__log_name

    @property
    def log_dirpath(self) -> str:
        return self.__log_dirpath

    @property
    def log_path(self) -> str:
        return self.__log_path

    @property
    def log_time_format(self) -> str:
        return self.__log_time_format

    @property
    def dir_time_format(self) -> str:
        return self.__dir_time_format

    def __get_comp_logger(self, comp_name: str) -> logging.Logger | None:
        return (
            self.__default_logger.getChild(comp_name)
            if comp_name in self.__registered_logger_names
            else None
        )

    def __get_comp_logger_or_default(self, comp_name: str | None) -> logging.Logger:
        logger = None
        if comp_name is not None:
            logger = self.__get_comp_logger(comp_name)
        return self.__default_logger if logger is None else logger

    def __register_comp_logger(self, comp_name: str, level: str | int | None) -> None:
        if comp_name in self.__registered_logger_names:
            return
        self.__registered_logger_names.add(comp_name)
        self.__default_logger.getChild(comp_name).setLevel(
            level if level is not None else logging.NOTSET
        )

    def set_default_logging_level(self, level: str | int | None) -> int:
        self.__default_logger.setLevel(level if level is not None else logging.NOTSET)
        return self.__default_logger.level

    def set_component_logging_level(self, comp_name: str, level: str | int | None) -> int:
        logger = self.__get_comp_logger(comp_name)
        assert logger is not None, comp_logger.log(
            logging.ERROR, f"Component {comp_name} not registered"
        )
        logger.setLevel(level if level is not None else logging.NOTSET)
        return logger.getEffectiveLevel()

    def __convert_name_to_readable(self, name: str) -> str:
        name = name.replace("_", " ").replace("-", " ").strip()
        if " " in name:
            # for names with snake_case or kebab-case
            return " ".join([f"{word[0].upper()}{word[1:]}" for word in name.split()])
        else:
            # for names with camelCase or PascalCase
            split_idxs = [
                0,
                *[
                    i
                    for i, c in enumerate(name)
                    if i != 0
                    and c.isupper()
                    and (
                        (i > 0 and name[i - 1].islower())
                        or (i < len(name) and name[i + 1].islower())
                    )
                ],
                len(name),
            ]
            return " ".join(
                [
                    f"{word[0].upper()}{word[1:]}"
                    for word in [
                        name[split_idxs[i] : split_idxs[i + 1]] for i in range(len(split_idxs) - 1)
                    ]
                ]
            )

    def __get_readable_name(self, comp_name: str, name_level: int) -> str:
        """
        Get more human-readable name, interpreted from input component name (which is likely to be
        __file__ by design).

        Input Args:
            `comp_name`: input name, likely to be __file__ of corresponding component
            `name_level`: level of name to be returned, indicating number of directory levels
                included in front of the component name.

        Returns:
            more human-readable component name
        """
        comp_name = os.path.abspath(comp_name)

        dir_names = []
        comp_dir = os.path.dirname(comp_name)
        for _ in range(name_level):
            dir_basename = os.path.basename(comp_dir)
            dir_names.append(self.__convert_name_to_readable(dir_basename))
            comp_dir = os.path.dirname(comp_dir)
        dir_name = " / ".join(dir_names[::-1])

        comp_name = self.__convert_name_to_readable(
            " ".join(os.path.basename(comp_name).split(".")[:-1])
        )

        if len(dir_name) != 0:
            comp_name = f"{dir_name} / {comp_name}"

        return comp_name

    def register_component(
        self,
        comp_name: str,
        level: str | int | None = None,
        auto_readable: bool = True,
        name_level: int = 0,
    ) -> CompLogger:
        """
        Register a component with a name and logging level.

        If `auto_readable` is True, the component name will be converted to a more human-readable
        form, which is interpreted from the input component name (which is likely to be __file__ by
        design). The human-readable name will be formatted as:
        "[dir1] / [dir2] / ... / [component_name]", where [dir1], [dir2], ... are the directory
        names of the component file, and [component_name] is the name of the component file without
        extension, with underscores replaced by spaces and each word capitalized. The number of
        directory levels included in the name is determined by `name_level`. If a custom name is
        desired, set `auto_readable` to False and `comp_name` to the desired name.

        Args:
            `comp_name`: name of the component, likely to be __file__ of corresponding component
            `level`: logging level for this component, default to None (which means NOTSET)
            `auto_readable`: whether to convert the component name to a more human-readable form
                (default True)
            `name_level`: level of name to be returned, indicating number of directory levels
                included in front of the component name (default 0)

        Returns:
            A CompLogger instance for the component.

        Raises:
            AssertionError: if the component name is already registered.
        """
        if auto_readable:
            comp_name = self.__get_readable_name(comp_name, name_level)
        assert self.__get_comp_logger(comp_name) is None, comp_logger.log(
            logging.ERROR, f"Component name {comp_name} is registered twice"
        )
        self.__register_comp_logger(comp_name, level)
        return CompLogger(comp_name)

    def get_component_logging_header(self) -> str:
        return f"<%s> "

    def component_should_log(self, comp_name: str | None, level: int) -> bool:
        logger = None
        if comp_name is not None:
            logger = self.__get_comp_logger(comp_name)
        logger = self.__default_logger if logger is None else logger
        return logger.isEnabledFor(level)

    def log(
        self, comp_name: str | None, level: int, msg: str, *args, stacklevel=3, **kwargs
    ) -> None:
        """
        Log a message with the specified component name and logging level.

        Args:
            comp_name: The name of the component.
            level: The logging level.
            msg: The message to log.
            *args: Additional arguments to pass to the logger.
            stacklevel: The stack level to use for the logger, default to be 3 assuming calling from
                component logger.
            **kwargs: Additional keyword arguments to pass to the logger.
        """
        header: str = (
            self.get_component_logging_header() % comp_name if comp_name is not None else ""
        )
        logger: logging.Logger = self.__get_comp_logger_or_default(comp_name)
        logger.log(level, header + msg, *args, stacklevel=stacklevel, **kwargs)


class CompLogger:
    def __init__(self, comp_name: str):
        self.__comp_name = comp_name
        self.__logger = Logger()

    def log(self, level: int, msg: str, *args, **kwargs) -> None:
        self.__logger.log(self.__comp_name, level, msg, *args, **kwargs)

    @property
    def comp_name(self) -> str:
        return self.__comp_name

    def get_augmented_message(self, msg: str) -> str:
        header: str = (
            Logger().get_component_logging_header() if self.__comp_name is not None else ""
        )
        return header % self.__comp_name + msg


comp_logger = Logger().register_component(__file__)


def addLoggingLevel(levelName, levelNum, methodName=None):
    """
    Comprehensively adds a new logging level to the `logging` module and the
    currently configured logging class.

    `levelName` becomes an attribute of the `logging` module with the value
    `levelNum`. `methodName` becomes a convenience method for both `logging`
    itself and the class returned by `logging.getLoggerClass()` (usually just
    `logging.Logger`). If `methodName` is not specified, `levelName.lower()` is
    used.

    To avoid accidental clobberings of existing attributes, this method will
    raise an `AttributeError` if the level name is already an attribute of the
    `logging` module or if the method name is already present

    Example
    -------
    >>> addLoggingLevel('TRACE', logging.DEBUG - 5)
    >>> logging.getLogger(__name__).setLevel("TRACE")
    >>> logging.getLogger(__name__).trace('that worked')
    >>> logging.trace('so did this')
    >>> logging.TRACE
    5

    Reference: https://stackoverflow.com/questions/2183233/how-to-add-a-custom-loglevel-to-pythons-logging-facility/35804945#35804945
    """
    if not methodName:
        methodName = levelName.lower()

    if hasattr(logging, levelName):
        raise AttributeError('{} already defined in logging module'.format(levelName))
    if hasattr(logging, methodName):
        raise AttributeError('{} already defined in logging module'.format(methodName))
    if hasattr(logging.getLoggerClass(), methodName):
        raise AttributeError('{} already defined in logger class'.format(methodName))

    # This method was inspired by the answers to Stack Overflow post
    # http://stackoverflow.com/q/2183233/2988730, especially
    # http://stackoverflow.com/a/13638084/2988730
    def logForLevel(self, message, *args, **kwargs):
        if self.isEnabledFor(levelNum):
            self._log(levelNum, message, args, **kwargs)

    def logToRoot(message, *args, **kwargs):
        logging.log(levelNum, message, *args, **kwargs)

    logging.addLevelName(levelNum, levelName)
    setattr(logging, levelName, levelNum)
    setattr(logging.getLoggerClass(), methodName, logForLevel)
    setattr(logging, methodName, logToRoot)


addLoggingLevel("VERBOSE1", logging.INFO - 1, "verbose1")
addLoggingLevel("VERBOSE2", logging.INFO - 2, "verbose2")

import re
import sys
import atexit
import signal
import traceback

# saving the default exception handler
default_excepthook = sys.excepthook


def exc_handler(exctype, value, tb):
    """
    Replaced exception handler, added the functionality to log the exception and raise a SIGABRT.
    """
    # # remove default stream handler so exceptions do not print to stderr
    # for handler in logging.root.handlers:
    #     if isinstance(handler, logging.StreamHandler):
    #         logging.root.removeHandler(handler)

    # invoke the default exception hook
    default_excepthook(exctype, value, tb)

    log_filename = Logger().log_path
    # only log when there is a log file
    if len(log_filename) != 0:
        # format message and log it
        msg = re.subn(r"%", r"%%", "".join(traceback.format_exception(exctype, value, tb)))[0]
        comp_logger.log(logging.FATAL, msg)
        # comp_logger must exist at this point
        # NOTE a ":0" is appended so that the colon in filename does not confuse some smart file path
        # resolvers (e.g. in VSCode)
        cprint.wprintf(
            comp_logger.get_augmented_message(f"Log saved to file {log_filename}:0"),
            file=sys.stderr,
        )
    # raise sigterm to signal program termination
    signal.raise_signal(signal.SIGABRT)


# register the custom exception handler with system
sys.excepthook = exc_handler


def save_config_to_log_dir(config_path: str):
    """Copy the given config file into log_dirpath/config/"""
    log_dir = Logger().log_dirpath
    config_dir = os.path.join(log_dir, "config")
    os.makedirs(config_dir, exist_ok=True)

    # Copy config file
    dest_path = os.path.join(config_dir, os.path.basename(config_path))
    shutil.copy2(config_path, dest_path)
    print(f"[INFO] Config file saved to {dest_path}")