bitbot-3.11-fork/src/Logging.py

import datetime, logging, logging.handlers, os, queue, sys, time, typing
from src import utils

LEVELS = {
    "trace": logging.DEBUG-1,
    "debug": logging.DEBUG,
    "info": logging.INFO,
    "warn": logging.WARN,
    "error": logging.ERROR,
    "critical": logging.CRITICAL
}

class BitBotFormatter(logging.Formatter):
    def formatTime(self, record, datefmt=None):
        datetime_obj = datetime.datetime.fromtimestamp(record.created)
        return utils.iso8601_format(datetime_obj, milliseconds=True)

class HookedHandler(logging.StreamHandler):
    def __init__(self, func: typing.Callable[[int, str], None]):
        logging.StreamHandler.__init__(self)
        self._func = func

    def emit(self, record):
        self._func(record.levelno, self.format(record))

class Log(object):
    def __init__(self, to_file: bool, level: str, location: str):
        self._hooks = []

        logging.addLevelName(LEVELS["trace"], "TRACE")
        self.logger = logging.getLogger(__name__)

        if not level.lower() in LEVELS:
            raise ValueError("Unknown log level '%s'" % level)
        stdout_level = LEVELS[level.lower()]

        self.logger.setLevel(LEVELS["trace"])

        formatter = BitBotFormatter("%(asctime)s [%(levelname)s] %(message)s")

        stdout_handler = logging.StreamHandler(sys.stdout)
        stdout_handler.setLevel(stdout_level)
        stdout_handler.setFormatter(formatter)
        self.logger.addHandler(stdout_handler)

        test_handler = HookedHandler(self._on_log)
        test_handler.setLevel(LEVELS["debug"])
        test_handler.setFormatter(formatter)
        self.logger.addHandler(test_handler)

        if to_file:
            trace_path = os.path.join(location, "trace.log")
            trace_handler = logging.handlers.TimedRotatingFileHandler(
                trace_path, when="midnight", backupCount=5)
            trace_handler.setLevel(LEVELS["trace"])
            trace_handler.setFormatter(formatter)
            self.logger.addHandler(trace_handler)

            info_path = os.path.join(location, "info.log")
            info_handler = logging.handlers.TimedRotatingFileHandler(
                info_path, when="midnight", backupCount=0)
            info_handler.setLevel(LEVELS["info"])
            info_handler.setFormatter(formatter)
            self.logger.addHandler(info_handler)

            warn_path = os.path.join(location, "warn.log")
            warn_handler = logging.FileHandler(warn_path)
            warn_handler.setLevel(LEVELS["warn"])
            warn_handler.setFormatter(formatter)
            self.logger.addHandler(warn_handler)

    def hook(self, func: typing.Callable[[int, str], None]):
        self._hooks.append(func)
    def _on_log(self, levelno, line):
        for func in self._hooks:
            func(levelno, line)

    def trace(self, message: str, params: typing.List=None, **kwargs):
        self._log(message, params, LEVELS["trace"], kwargs)
    def debug(self, message: str, params: typing.List=None, **kwargs):
        self._log(message, params, logging.DEBUG, kwargs)
    def info(self, message: str, params: typing.List=None, **kwargs):
        self._log(message, params, logging.INFO, kwargs)
    def warn(self, message: str, params: typing.List=None, **kwargs):
        self._log(message, params, logging.WARN, kwargs)
    def error(self, message: str, params: typing.List=None, **kwargs):
        self._log(message, params, logging.ERROR, kwargs)
    def critical(self, message: str, params: typing.List=None, **kwargs):
        self._log(message, params, logging.CRITICAL, kwargs)
    def _log(self, message: str, params: typing.Optional[typing.List],
            level: int, kwargs: dict):
        self.logger.log(level, message, *(params or []), **kwargs)