libs/utils/logging.py

import os
import sys
import errno


def get_logger(logs_dir: str, file_name: str = "log.txt"):
    logger = PrintLogger(os.path.join(logs_dir, file_name))
    sys.stdout = logger  # record all python print
    return logger


class PrintLogger(object):

    def __init__(self, fpath=None):
        """
        python standard input/output records
        """
        self.console = sys.stdout
        self.file = None
        if fpath is not None:
            mkdir_if_missing(os.path.dirname(fpath))
            self.file = open(fpath, 'w')

    def __del__(self):
        self.close()

    def __enter__(self):
        pass

    def __exit__(self, *args):
        self.close()

    def write(self, msg):
        self.console.write(msg)
        if self.file is not None:
            self.file.write(msg)

    def write_in(self, msg):
        """write in log only, not console"""
        if self.file is not None:
            self.file.write(msg)

    def flush(self):
        self.console.flush()
        if self.file is not None:
            self.file.flush()
            os.fsync(self.file.fileno())

    def close(self):
        self.console.close()
        if self.file is not None:
            self.file.close()


def mkdir_if_missing(dir_path):
    try:
        os.makedirs(dir_path)
    except OSError as e:
        if e.errno != errno.EEXIST:
            raise