|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved |
| 2 | +import logging |
| 3 | +import os |
| 4 | + |
| 5 | +LOG_LEVELS = { |
| 6 | + "DEBUG": logging.DEBUG, |
| 7 | + "INFO": logging.INFO, |
| 8 | + "WARNING": logging.WARNING, |
| 9 | + "ERROR": logging.ERROR, |
| 10 | + "CRITICAL": logging.CRITICAL, |
| 11 | +} |
| 12 | + |
| 13 | + |
| 14 | +class ColoredFormatter(logging.Formatter): |
| 15 | + """A command line formatter with different colors for each level.""" |
| 16 | + |
| 17 | + def __init__(self): |
| 18 | + super().__init__() |
| 19 | + reset = "\033[0m" |
| 20 | + colors = { |
| 21 | + logging.DEBUG: f"{reset}\033[36m", # cyan, |
| 22 | + logging.INFO: f"{reset}\033[32m", # green |
| 23 | + logging.WARNING: f"{reset}\033[33m", # yellow |
| 24 | + logging.ERROR: f"{reset}\033[31m", # red |
| 25 | + logging.CRITICAL: f"{reset}\033[35m", # magenta |
| 26 | + } |
| 27 | + fmt_str = "{color}%(levelname)s %(asctime)s %(process)d %(filename)s:%(lineno)4d:{reset} %(message)s" |
| 28 | + self.formatters = { |
| 29 | + level: logging.Formatter(fmt_str.format(color=color, reset=reset)) |
| 30 | + for level, color in colors.items() |
| 31 | + } |
| 32 | + self.default_formatter = self.formatters[logging.INFO] |
| 33 | + |
| 34 | + def format(self, record): |
| 35 | + formatter = self.formatters.get(record.levelno, self.default_formatter) |
| 36 | + return formatter.format(record) |
| 37 | + |
| 38 | + |
| 39 | +def get_logger(name, level=logging.INFO): |
| 40 | + """A command line logger.""" |
| 41 | + if "LOG_LEVEL" in os.environ: |
| 42 | + level = os.environ["LOG_LEVEL"].upper() |
| 43 | + assert ( |
| 44 | + level in LOG_LEVELS |
| 45 | + ), f"Invalid LOG_LEVEL: {level}, must be one of {list(LOG_LEVELS.keys())}" |
| 46 | + level = LOG_LEVELS[level] |
| 47 | + logger = logging.getLogger(name) |
| 48 | + logger.setLevel(level) |
| 49 | + logger.propagate = False |
| 50 | + ch = logging.StreamHandler() |
| 51 | + ch.setLevel(level) |
| 52 | + ch.setFormatter(ColoredFormatter()) |
| 53 | + logger.addHandler(ch) |
| 54 | + return logger |
0 commit comments