utils/logger.py

import datetime
import functools
import inspect
import logging
import os
import shutil
import subprocess
from termcolor import colored


def add_fillers(text, filler="=", fill_side="both"):
    terminal_width = shutil.get_terminal_size().columns
    text = text.strip()
    text_width = len(text)
    if text_width >= terminal_width:
        return text

    if fill_side[0].lower() == "b":
        leading_fill_str = filler * ((terminal_width - text_width) // 2 - 1) + " "
        trailing_fill_str = " " + filler * (
            terminal_width - text_width - len(leading_fill_str) - 1
        )
    elif fill_side[0].lower() == "l":
        leading_fill_str = filler * (terminal_width - text_width - 1) + " "
        trailing_fill_str = ""
    elif fill_side[0].lower() == "r":
        leading_fill_str = ""
        trailing_fill_str = " " + filler * (terminal_width - text_width - 1)
    else:
        raise ValueError("Invalid fill_side")

    filled_str = f"{leading_fill_str}{text}{trailing_fill_str}"
    return filled_str


class OSLogger(logging.Logger):
    LOG_METHODS = {
        "err": ("error", "red"),
        "warn": ("warning", "light_red"),
        "note": ("info", "light_magenta"),
        "mesg": ("info", "light_cyan"),
        "file": ("info", "light_blue"),
        "line": ("info", "white"),
        "success": ("info", "light_green"),
        "fail": ("info", "light_red"),
        "back": ("debug", "light_cyan"),
    }
    INDENT_METHODS = [
        "indent",
        "set_indent",
        "reset_indent",
        "store_indent",
        "restore_indent",
        "log_indent",
    ]
    LEVEL_METHODS = [
        "set_level",
        "store_level",
        "restore_level",
        "quiet",
        "enter_quiet",
        "exit_quiet",
    ]
    LEVEL_NAMES = {
        "critical": logging.CRITICAL,
        "error": logging.ERROR,
        "warning": logging.WARNING,
        "info": logging.INFO,
        "debug": logging.DEBUG,
    }

    def __init__(self, name=None, prefix=False):
        if not name:
            frame = inspect.stack()[1]
            module = inspect.getmodule(frame[0])
            name = module.__name__

        super().__init__(name)
        self.setLevel(logging.INFO)

        if prefix:
            formatter_prefix = "[%(asctime)s] - [%(name)s] - [%(levelname)s]\n"
        else:
            formatter_prefix = ""

        self.formatter = logging.Formatter(formatter_prefix + "%(message)s")

        stream_handler = logging.StreamHandler()
        stream_handler.setLevel(logging.INFO)
        stream_handler.setFormatter(self.formatter)
        self.addHandler(stream_handler)

        self.log_indent = 0
        self.log_indents = []

        self.log_level = "info"
        self.log_levels = []

    def indent(self, indent=2):
        self.log_indent += indent

    def set_indent(self, indent=2):
        self.log_indent = indent

    def reset_indent(self):
        self.log_indent = 0

    def store_indent(self):
        self.log_indents.append(self.log_indent)

    def restore_indent(self):
        self.log_indent = self.log_indents.pop(-1)

    def set_level(self, level):
        self.log_level = level
        self.setLevel(self.LEVEL_NAMES[level])

    def store_level(self):
        self.log_levels.append(self.log_level)

    def restore_level(self):
        self.log_level = self.log_levels.pop(-1)
        self.set_level(self.log_level)

    def quiet(self):
        self.set_level("critical")

    def enter_quiet(self, quiet=False):
        if quiet:
            self.store_level()
            self.quiet()

    def exit_quiet(self, quiet=False):
        if quiet:
            self.restore_level()

    def log(
        self,
        level,
        color,
        msg,
        indent=0,
        fill=False,
        fill_side="both",
        end="\n",
        *args,
        **kwargs,
    ):
        if type(msg) == str:
            msg_str = msg
        else:
            msg_str = repr(msg)
            quotes = ["'", '"']
            if msg_str[0] in quotes and msg_str[-1] in quotes:
                msg_str = msg_str[1:-1]

        indent_str = " " * (self.log_indent + indent)
        indented_msg = "\n".join([indent_str + line for line in msg_str.split("\n")])

        if fill:
            indented_msg = add_fillers(indented_msg, fill_side=fill_side)

        handler = self.handlers[0]
        handler.terminator = end

        getattr(self, level)(colored(indented_msg, color), *args, **kwargs)

    def route_log(self, method, msg, *args, **kwargs):
        level, method = method
        functools.partial(self.log, level, method, msg)(*args, **kwargs)

    def err(self, msg: str = "", *args, **kwargs):
        self.route_log(("error", "red"), msg, *args, **kwargs)

    def warn(self, msg: str = "", *args, **kwargs):
        self.route_log(("warning", "light_red"), msg, *args, **kwargs)

    def note(self, msg: str = "", *args, **kwargs):
        self.route_log(("info", "light_magenta"), msg, *args, **kwargs)

    def mesg(self, msg: str = "", *args, **kwargs):
        self.route_log(("info", "light_cyan"), msg, *args, **kwargs)

    def file(self, msg: str = "", *args, **kwargs):
        self.route_log(("info", "light_blue"), msg, *args, **kwargs)

    def line(self, msg: str = "", *args, **kwargs):
        self.route_log(("info", "white"), msg, *args, **kwargs)

    def success(self, msg: str = "", *args, **kwargs):
        self.route_log(("info", "light_green"), msg, *args, **kwargs)

    def fail(self, msg: str = "", *args, **kwargs):
        self.route_log(("info", "light_red"), msg, *args, **kwargs)

    def back(self, msg: str = "", *args, **kwargs):
        self.route_log(("debug", "light_cyan"), msg, *args, **kwargs)


logger = OSLogger()


def shell_cmd(cmd, getoutput=False, showcmd=True, env=None):
    if showcmd:
        logger.info(colored(f"\n$ [{os.getcwd()}]", "light_blue"))
        logger.info(colored(f"  $ {cmd}\n", "light_cyan"))
    if getoutput:
        output = subprocess.getoutput(cmd, env=env)
        return output
    else:
        subprocess.run(cmd, shell=True, env=env)


class Runtimer:
    def __enter__(self):
        self.t1, _ = self.start_time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.t2, _ = self.end_time()
        self.elapsed_time(self.t2 - self.t1)

    def start_time(self):
        t1 = datetime.datetime.now()
        self.logger_time("start", t1)
        return t1, self.time2str(t1)

    def end_time(self):
        t2 = datetime.datetime.now()
        self.logger_time("end", t2)
        return t2, self.time2str(t2)

    def elapsed_time(self, dt=None):
        if dt is None:
            dt = self.t2 - self.t1
        self.logger_time("elapsed", dt)
        return dt, self.time2str(dt)

    def logger_time(self, time_type, t):
        time_types = {
            "start": "Start",
            "end": "End",
            "elapsed": "Elapsed",
        }
        time_str = add_fillers(
            colored(
                f"{time_types[time_type]} time: [ {self.time2str(t)} ]",
                "light_magenta",
            ),
            fill_side="both",
        )
        logger.line(time_str)

    # Convert time to string
    def time2str(self, t):
        datetime_str_format = "%Y-%m-%d %H:%M:%S"
        if isinstance(t, datetime.datetime):
            return t.strftime(datetime_str_format)
        elif isinstance(t, datetime.timedelta):
            hours = t.seconds // 3600
            hour_str = f"{hours} hr" if hours > 0 else ""
            minutes = (t.seconds // 60) % 60
            minute_str = f"{minutes:>2} min" if minutes > 0 else ""
            seconds = t.seconds % 60
            second_str = f"{seconds:>2} s"
            time_str = " ".join([hour_str, minute_str, second_str]).strip()
            return time_str
        else:
            return str(t)