faceswap/lib/logger.py

#!/usr/bin/python
""" Logging Setup """
import collections
import logging
from logging.handlers import RotatingFileHandler
import os
import sys
import traceback

from datetime import datetime
from tqdm import tqdm

from numpy import ndarray


class FaceswapLogger(logging.Logger):
    """ Create custom logger  with custom levels """
    def __init__(self, name):
        for new_level in (("VERBOSE", 15), ("TRACE", 5)):
            level_name, level_num = new_level
            if hasattr(logging, level_name):
                continue
            logging.addLevelName(level_num, level_name)
            setattr(logging, level_name, level_num)
        super().__init__(name)

    def verbose(self, msg, *args, **kwargs):
        """
        Log 'msg % args' with severity 'VERBOSE'.
        """
        if self.isEnabledFor(15):
            self._log(15, msg, args, **kwargs)

    def trace(self, msg, *args, **kwargs):
        """
        Log 'msg % args' with severity 'VERBOSE'.
        """
        if self.isEnabledFor(5):
            self._log(5, msg, args, **kwargs)


class FaceswapFormatter(logging.Formatter):
    """ Override formatter to strip newlines from logger arguments """
    def format(self, record):
        if isinstance(record.msg, str) and ("\n" in record.msg or "\r" in record.msg):
            record.msg = record.msg.replace("\n", "\\n").replace("\r", "\\r")
        if any(self.reformat_check(arg) for arg in record.args):
            record.args = self.reformat_args(record.args)
        return super().format(record)

    @staticmethod
    def reformat_check(arg):
        """ Check if the argument should be reformatted
            The argument is a string with a line break
            The argument is a numpy array
        """
        return ((isinstance(arg, str) and ("\n" in arg or "\r" in arg))
                or isinstance(arg, ndarray))

    @staticmethod
    def reformat_args(args):
        """ Reformat args that require new lines removing """
        new_args = []
        for arg in args:
            if isinstance(arg, ndarray):
                # Convert numpy arrays to string for reformatting
                arg = str(ndarray)
            if isinstance(arg, str) and ("\n" in arg or "\r" in arg):
                arg = arg.replace("\n", "\\n").replace("\r", "\\r")
            new_args.append(arg)
        return tuple(new_args)


class RollingBuffer(collections.deque):
    """File-like that keeps a certain number of lines of text in memory."""
    def write(self, buffer):
        """ Write line to buffer """
        for line in buffer.rstrip().splitlines():
            self.append(line + "\n")


class TqdmHandler(logging.StreamHandler):
    """ Use TQDM Write for outputting to console """
    def emit(self, record):
        msg = self.format(record)
        tqdm.write(msg)


def set_root_logger(loglevel=logging.INFO):
    """ Setup the root logger. """
    rootlogger = logging.getLogger()
    rootlogger.setLevel(loglevel)
    return rootlogger


def log_setup(loglevel, logfile, command, is_gui=False):
    """ initial log set up. """
    numeric_loglevel = get_loglevel(loglevel)
    root_loglevel = min(logging.DEBUG, numeric_loglevel)
    rootlogger = set_root_logger(loglevel=root_loglevel)
    log_format = FaceswapFormatter("%(asctime)s %(processName)-15s %(threadName)-15s "
                                   "%(module)-15s %(funcName)-25s %(levelname)-8s %(message)s",
                                   datefmt="%m/%d/%Y %H:%M:%S")
    f_handler = file_handler(numeric_loglevel, logfile, log_format, command)
    s_handler = stream_handler(numeric_loglevel, is_gui)
    c_handler = crash_handler(log_format)
    rootlogger.addHandler(f_handler)
    rootlogger.addHandler(s_handler)
    rootlogger.addHandler(c_handler)
    logging.info("Log level set to: %s", loglevel.upper())


def file_handler(loglevel, logfile, log_format, command):
    """ Add a logging rotating file handler """
    if logfile is not None:
        filename = logfile
    else:
        filename = os.path.join(os.path.dirname(os.path.realpath(sys.argv[0])), "faceswap")
        # Windows has issues sharing the log file with subprocesses, so log GUI separately
        filename += "_gui.log" if command == "gui" else ".log"

    should_rotate = os.path.isfile(filename)
    log_file = RotatingFileHandler(filename, backupCount=1)
    if should_rotate:
        log_file.doRollover()
    log_file.setFormatter(log_format)
    log_file.setLevel(loglevel)
    return log_file


def stream_handler(loglevel, is_gui):
    """ Add a logging cli handler """
    # Don't set stdout to lower than verbose
    loglevel = max(loglevel, 15)
    log_format = FaceswapFormatter("%(asctime)s %(levelname)-8s %(message)s",
                                   datefmt="%m/%d/%Y %H:%M:%S")

    if is_gui:
        # tqdm.write inserts extra lines in the GUI, so use standard output as
        # it is not needed there.
        log_console = logging.StreamHandler(sys.stdout)
    else:
        log_console = TqdmHandler(sys.stdout)
    log_console.setFormatter(log_format)
    log_console.setLevel(loglevel)
    return log_console


def crash_handler(log_format):
    """ Add a handler that sores the last 100 debug lines to 'debug_buffer'
        for use in crash reports """
    log_crash = logging.StreamHandler(debug_buffer)
    log_crash.setFormatter(log_format)
    log_crash.setLevel(logging.DEBUG)
    return log_crash


def get_loglevel(loglevel):
    """ Check valid log level supplied and return numeric log level """
    numeric_level = getattr(logging, loglevel.upper(), None)
    if not isinstance(numeric_level, int):
        raise ValueError("Invalid log level: %s" % loglevel)

    return numeric_level


def crash_log():
    """ Write debug_buffer to a crash log on crash """
    from lib.sysinfo import sysinfo
    path = os.getcwd()
    filename = os.path.join(path, datetime.now().strftime("crash_report.%Y.%m.%d.%H%M%S%f.log"))

    freeze_log = list(debug_buffer)
    with open(filename, "w") as outfile:
        outfile.writelines(freeze_log)
        traceback.print_exc(file=outfile)
        outfile.write(sysinfo)
    return filename


old_factory = logging.getLogRecordFactory()  # pylint: disable=invalid-name


def faceswap_logrecord(*args, **kwargs):
    """ Add a flag to logging.LogRecord to not strip formatting from particular records """
    record = old_factory(*args, **kwargs)
    record.strip_spaces = True
    return record


logging.setLogRecordFactory(faceswap_logrecord)

# Set logger class to custom logger
logging.setLoggerClass(FaceswapLogger)

# Stores the last 100 debug messages
debug_buffer = RollingBuffer(maxlen=100)  # pylint: disable=invalid-name