|
7 | 7 |
|
8 | 8 | import logging |
9 | 9 |
|
| 10 | +class ColoredFormatter(logging.Formatter): |
| 11 | + """A logging formatter that adds ANSI color codes to log levels.""" |
| 12 | + |
| 13 | + # ANSI color codes matching pytest's colors |
| 14 | + COLORS = { |
| 15 | + logging.CRITICAL: "\033[91m", # Red |
| 16 | + logging.ERROR: "\033[91m", # Red |
| 17 | + logging.WARNING: "\033[93m", # Yellow |
| 18 | + logging.INFO: "\033[92m", # Green |
| 19 | + logging.DEBUG: "\033[95m", # Magenta/Purple |
| 20 | + } |
| 21 | + RESET = "\033[0m" |
| 22 | + |
| 23 | + def format(self, record: logging.LogRecord) -> str: |
| 24 | + color = self.COLORS.get(record.levelno, "") |
| 25 | + original_levelname = record.levelname |
| 26 | + if color: |
| 27 | + record.levelname = f'{color}{record.levelname}{self.RESET}' |
| 28 | + result = super().format(record) |
| 29 | + record.levelname = original_levelname |
| 30 | + return result |
| 31 | + |
| 32 | + |
| 33 | +def setup_colored_logging(level: int = logging.INFO) -> None: |
| 34 | + """Set up logging with colored formatter.""" |
| 35 | + logger = logging.getLogger() |
| 36 | + logger.setLevel(level) |
| 37 | + for handler in logger.handlers[:]: |
| 38 | + logger.removeHandler(handler) |
| 39 | + handler = logging.StreamHandler() |
| 40 | + formatter = ColoredFormatter( |
| 41 | + fmt="%(asctime)s.%(msecs)03d %(levelname)s %(message)s", datefmt="%b %d %H:%M:%S" |
| 42 | + ) |
| 43 | + handler.setFormatter(formatter) |
| 44 | + logger.addHandler(handler) |
| 45 | + |
| 46 | + |
10 | 47 | logger = logging.getLogger() |
11 | | -logging.basicConfig( |
12 | | - format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s", datefmt="%b %d %H:%M:%S", level=logging.INFO |
13 | | -) |
| 48 | +setup_colored_logging() |
0 commit comments