custom_logging.py

import os.path as osp
import os
import errno
import sys

def mkdir_if_missing(directory):
    if not osp.exists(directory):
        try:
            os.makedirs(directory)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise   


class Logger(object):
    """
    Write console output to external text file.
    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/utils/logging.py.
    """  
    def __init__(self, fpath=None):
        self.console = sys.stdout
        self.file = None
        if fpath is not None:
            cnt=1
            while os.path.exists(fpath):
                s = os.path.splitext(fpath)
                fpath =  s[0]+"_{}.".format(cnt)+s[1]
                cnt+=1
            if not os.path.exists(fpath):
                mkdir_if_missing(osp.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 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()