patientsim/utils/__init__.py

import os
import logging.config


LOGGING_NAME = "PatientSim"
VERBOSE = True


def set_logging(name=LOGGING_NAME, verbose=True):
    """Sets up logging for the given name."""
    rank = int(os.getenv('RANK', -1))
    level = logging.INFO if verbose and rank in {-1, 0} else logging.ERROR

    class ColorFormatter(logging.Formatter):
        def format(self, record):
            if record.levelname == "ERROR":
                record.msg = colorstr("red", record.msg)
            elif record.levelname == "WARNING":
                record.msg = colorstr("yellow", record.msg)
            elif record.levelname == "DEBUG":
                record.msg = colorstr("blue", record.msg)
            return super().format(record)

    logging.config.dictConfig({
        'version': 1,
        'disable_existing_loggers': False,
        'formatters': {
            name: {
                '()': ColorFormatter,
                'format': '[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s'
            }
        },
        'handlers': {
            name: {
                'class': 'logging.StreamHandler',
                'formatter': name,
                'level': level
            }
        },
        'loggers': {
            name: {
                'level': level,
                'handlers': [name],
                'propagate': False
            }
        }
    })


set_logging(LOGGING_NAME, verbose=VERBOSE)
LOGGER = logging.getLogger(LOGGING_NAME)


def colorstr(*input):
    *args, string = input if len(input) > 1 else ('blue', 'bold', input[0])
    colors = {
        'black': '\033[30m',
        'red': '\033[31m',
        'green': '\033[32m',
        'yellow': '\033[33m',
        'blue': '\033[34m',
        'magenta': '\033[35m',
        'cyan': '\033[36m',
        'white': '\033[37m',
        'bright_black': '\033[90m',
        'bright_red': '\033[91m',
        'bright_green': '\033[92m',
        'bright_yellow': '\033[93m',
        'bright_blue': '\033[94m',
        'bright_magenta': '\033[95m',
        'bright_cyan': '\033[96m',
        'bright_white': '\033[97m',
        'end': '\033[0m',
        'bold': '\033[1m',
        'underline': '\033[4m',
    }
    return ''.join(colors[x] for x in args) + f'{string}' + colors['end']


def log(message, level='info', color=False):
    if level.lower() == 'warning':
        LOGGER.warning(message)
    elif level.lower() == 'error':
        LOGGER.error(message)
    else:
        if color:
            LOGGER.info(colorstr(message))
        else:
            LOGGER.info(message)