|
4 | 4 |
|
5 | 5 |
|
6 | 6 |
|
| 7 | +from __future__ import print_function |
| 8 | + |
7 | 9 | import builtins |
8 | 10 | from contextlib import nullcontext |
9 | 11 | import logging |
10 | | -from logging import Handler, Logger |
| 12 | +from logging import Formatter, Handler, Logger |
11 | 13 | import os |
12 | 14 | from prompt_toolkit import patch_stdout |
13 | 15 | import sys |
14 | 16 |
|
15 | 17 |
|
| 18 | +class _PyflybyFormatter(Formatter): |
| 19 | + _ANSI_RESET = "\033[0m" |
| 20 | + _COLORS = { |
| 21 | + "blue": "\033[34m", |
| 22 | + "yellow": "\033[33m", |
| 23 | + "red": "\033[31m", |
| 24 | + } |
| 25 | + |
| 26 | + def _color_for_level(self, levelno): |
| 27 | + if levelno >= logging.ERROR: |
| 28 | + return self._COLORS["red"] |
| 29 | + elif levelno >= logging.WARNING: |
| 30 | + return self._COLORS["yellow"] |
| 31 | + else: |
| 32 | + return self._COLORS["blue"] |
| 33 | + |
| 34 | + def formatInteractive(self, record): |
| 35 | + color = self._color_for_level(record.levelno) |
| 36 | + prefix = "%s%s[PYFLYBY]%s " % (self._ANSI_RESET, color, self._ANSI_RESET) |
| 37 | + msg = super().format(record) |
| 38 | + return "".join(["%s%s\n" % (prefix, line) for line in msg.splitlines()]) |
| 39 | + |
| 40 | + def formatPlain(self, record): |
| 41 | + msg = super().format(record) |
| 42 | + return "".join(["[PYFLYBY] %s\n" % line for line in msg.splitlines()]) |
| 43 | + |
| 44 | + |
16 | 45 | class _PyflybyHandler(Handler): |
17 | 46 |
|
18 | 47 | _pre_log_function = None |
19 | 48 | _logged_anything_during_context = False |
20 | 49 |
|
21 | | - _interactive_prefix = "\033[0m\033[33m[PYFLYBY]\033[0m " |
22 | | - _noninteractive_prefix = "[PYFLYBY] " |
| 50 | + def __init__(self): |
| 51 | + super().__init__() |
| 52 | + self.setFormatter(_PyflybyFormatter()) |
23 | 53 |
|
24 | 54 | def emit(self, record): |
25 | 55 | try: |
| 56 | + formatter = self.formatter |
26 | 57 | if _is_ipython() or _is_interactive(sys.stderr): |
27 | | - prefix = self._interactive_prefix |
| 58 | + msg = formatter.formatInteractive(record) |
28 | 59 | patch_stdout_c = patch_stdout.patch_stdout(raw=True) |
29 | 60 | else: |
30 | | - prefix = self._noninteractive_prefix |
| 61 | + msg = formatter.formatPlain(record) |
31 | 62 | patch_stdout_c = nullcontext() |
32 | 63 |
|
33 | | - msg = self.format(record) |
34 | | - msg = ''.join(["%s%s\n" % (prefix, line) for line in msg.splitlines()]) |
35 | 64 | with patch_stdout_c: |
36 | 65 | sys.stderr.write(msg) |
37 | 66 | sys.stderr.flush() |
@@ -75,8 +104,7 @@ class PyflybyLogger(Logger): |
75 | 104 |
|
76 | 105 | def __init__(self, name, level): |
77 | 106 | Logger.__init__(self, name) |
78 | | - handler = _PyflybyHandler() |
79 | | - self.addHandler(handler) |
| 107 | + self.addHandler(_PyflybyHandler()) |
80 | 108 | self.set_level(level) |
81 | 109 |
|
82 | 110 | def set_level(self, level): |
|
0 commit comments