|
1 | 1 | import logging |
2 | 2 | import os |
3 | 3 | import sys |
| 4 | +import threading |
| 5 | +from enum import Enum |
| 6 | +from typing import Optional, Union |
| 7 | + |
| 8 | +logging.addLevelName(logging.WARNING, "WARN") |
4 | 9 |
|
5 | 10 | loglevel = os.getenv("LOGLEVEL", "INFO") |
6 | 11 |
|
|
11 | 16 | use_colors = sys.stderr.isatty() |
12 | 17 |
|
13 | 18 |
|
| 19 | +class HandlerMode(Enum): |
| 20 | + NOWHERE = 0 |
| 21 | + IN_HIVEMIND = 1 |
| 22 | + IN_ROOT_LOGGER = 2 |
| 23 | + |
| 24 | + |
| 25 | +_init_lock = threading.RLock() |
| 26 | +_current_mode = HandlerMode.IN_HIVEMIND |
| 27 | +_default_handler = None |
| 28 | + |
| 29 | + |
14 | 30 | class TextStyle: |
15 | 31 | """ |
16 | 32 | ANSI escape codes. Details: https://en.wikipedia.org/wiki/ANSI_escape_code#Colors |
@@ -60,23 +76,82 @@ def format(self, record: logging.LogRecord) -> str: |
60 | 76 | return super().format(record) |
61 | 77 |
|
62 | 78 |
|
63 | | -def get_logger(module_name: str) -> logging.Logger: |
64 | | - # trim package name |
65 | | - name_without_prefix = ".".join(module_name.split(".")[1:]) |
| 79 | +def _initialize_if_necessary(): |
| 80 | + global _current_mode, _default_handler |
66 | 81 |
|
67 | | - logging.addLevelName(logging.WARNING, "WARN") |
68 | | - formatter = CustomFormatter( |
69 | | - fmt="{asctime}.{msecs:03.0f} [{bold}{levelcolor}{levelname}{reset}] [{bold}{caller}{reset}] {message}", |
70 | | - style="{", |
71 | | - datefmt="%b %d %H:%M:%S", |
72 | | - ) |
73 | | - handler = logging.StreamHandler() |
74 | | - handler.setFormatter(formatter) |
75 | | - logger = logging.getLogger(name_without_prefix) |
76 | | - logger.setLevel(loglevel) |
77 | | - logger.addHandler(handler) |
| 82 | + with _init_lock: |
| 83 | + if _default_handler is not None: |
| 84 | + return |
| 85 | + |
| 86 | + formatter = CustomFormatter( |
| 87 | + fmt="{asctime}.{msecs:03.0f} [{bold}{levelcolor}{levelname}{reset}] [{bold}{caller}{reset}] {message}", |
| 88 | + style="{", |
| 89 | + datefmt="%b %d %H:%M:%S", |
| 90 | + ) |
| 91 | + _default_handler = logging.StreamHandler() |
| 92 | + _default_handler.setFormatter(formatter) |
| 93 | + |
| 94 | + _enable_default_handler("hivemind") |
| 95 | + |
| 96 | + |
| 97 | +def get_logger(name: Optional[str] = None) -> logging.Logger: |
| 98 | + """ |
| 99 | + Same as ``logging.getLogger()`` but ensures that the default log handler is initialized. |
| 100 | + """ |
| 101 | + |
| 102 | + _initialize_if_necessary() |
| 103 | + return logging.getLogger(name) |
| 104 | + |
| 105 | + |
| 106 | +def _enable_default_handler(name: str) -> None: |
| 107 | + logger = get_logger(name) |
| 108 | + logger.addHandler(_default_handler) |
78 | 109 | logger.propagate = False |
79 | | - return logger |
| 110 | + logger.setLevel(loglevel) |
| 111 | + |
| 112 | + |
| 113 | +def _disable_default_handler(name: str) -> None: |
| 114 | + logger = get_logger(name) |
| 115 | + logger.removeHandler(_default_handler) |
| 116 | + logger.propagate = True |
| 117 | + logger.setLevel(logging.NOTSET) |
| 118 | + |
| 119 | + |
| 120 | +def use_hivemind_log_handler(where: Union[HandlerMode, str]) -> None: |
| 121 | + """ |
| 122 | + Choose loggers where the default hivemind log handler is applied. Options for the ``where`` argument are: |
| 123 | +
|
| 124 | + * "in_hivemind" (default): Use the hivemind log handler in the loggers of the ``hivemind`` package. |
| 125 | + Don't propagate their messages to the root logger. |
| 126 | + * "nowhere": Don't use the hivemind log handler anywhere. |
| 127 | + Propagate the ``hivemind`` messages to the root logger. |
| 128 | + * "in_root_logger": Use the hivemind log handler in the root logger |
| 129 | + (that is, in all application loggers until they disable propagation to the root logger). |
| 130 | + Propagate the ``hivemind`` messages to the root logger. |
| 131 | +
|
| 132 | + The options may be defined as strings (case-insensitive) or values from the HandlerMode enum. |
| 133 | + """ |
| 134 | + |
| 135 | + global _current_mode |
| 136 | + |
| 137 | + if isinstance(where, str): |
| 138 | + # We allow `where` to be a string, so a developer does not have to import the enum for one usage |
| 139 | + where = HandlerMode[where.upper()] |
| 140 | + |
| 141 | + if where == _current_mode: |
| 142 | + return |
| 143 | + |
| 144 | + if _current_mode == HandlerMode.IN_HIVEMIND: |
| 145 | + _disable_default_handler("hivemind") |
| 146 | + elif _current_mode == HandlerMode.IN_ROOT_LOGGER: |
| 147 | + _disable_default_handler(None) |
| 148 | + |
| 149 | + _current_mode = where |
| 150 | + |
| 151 | + if _current_mode == HandlerMode.IN_HIVEMIND: |
| 152 | + _enable_default_handler("hivemind") |
| 153 | + elif _current_mode == HandlerMode.IN_ROOT_LOGGER: |
| 154 | + _enable_default_handler(None) |
80 | 155 |
|
81 | 156 |
|
82 | 157 | def golog_level_to_python(level: str) -> int: |
|
0 commit comments