"""Methods and classes for logging capabilities of HEROs"""
import contextlib
import contextvars
import logging
import os
import sys
from functools import wraps
SPAM = 5
logging.SPAM = 5 # ty:ignore
logging.addLevelName(levelName="SPAM", level=5)
# context variable that holds the name of the HERO whose capability is currently executed
_current_hero: contextvars.ContextVar = contextvars.ContextVar("heros.current_hero", default=None)
_registered_heros: set[str] = set()
[docs]
def register_hero(name: str) -> None:
"""Register a LocalHERO as living in this process.
Records emitted from contexts without an active HERO (e.g. background
threads of dependencies) are attributed to the registered HERO.
"""
_registered_heros.add(name)
[docs]
def unregister_hero(name: str) -> None:
"""Remove a LocalHERO from the set of HEROs living in this process."""
_registered_heros.discard(name)
[docs]
def resolve_origin() -> str:
"""Resolve the HERO name that log records emitted now belong to.
Returns:
The name of the currently active HERO if set, else the name of the
single HERO registered in this process, ``PID:<pid>`` if no or more than
one HERO is registered
"""
origin = _current_hero.get()
if origin is not None:
return origin
if len(_registered_heros) == 1:
return next(iter(_registered_heros))
return f"PID:{os.getpid()}"
[docs]
class hero_logging_context(contextlib.ContextDecorator):
"""Attribute all log records emitted within this context to a specific HERO.
While the context is active, every log record created in it is stamped with the name of the HERO, so that records
emitted by dependencies of a heroized class can be attributed to the HERO whose capability is currently executed.
Can be used in three ways:
- as a context manager:
```with hero_logging_context("my_hero"): ...```
- as a function or method decorator:
```@hero_logging_context("my_hero")``` uses a fixed name, while
```@hero_logging_context(name_attr="_name")``` reads the name at call time from the instance the decorated method
belongs to
- as a class decorator:
`@hero_logging_context(name_attr="name")` wraps all plain-function methods of the class (and its base classes).
Args:
name: name of the HERO
name_attr: name of the attribute of the instance the decorated method belongs to. Takes precedence over `name`.
"""
def __init__(self, name: str = "unknown", name_attr: str | None = None):
self.name = name
self.name_attr = name_attr
self._tokens: list[contextvars.Token] = []
[docs]
def __call__(self, func):
if func.__class__ is type:
seen = set()
for cls in func.__mro__:
for attr_name, attr_value in vars(cls).items():
if attr_name.startswith("__") or attr_name in seen:
continue
seen.add(attr_name)
if isinstance(attr_value, (staticmethod, classmethod)) or not callable(attr_value):
continue # leaves staticmethod, classmethod, as they are
if getattr(attr_value, "_heros_log_context", False):
continue # already wrapped by a decorated base class
wrapped = hero_logging_context(self.name, self.name_attr)(attr_value)
wrapped._heros_log_context = True
setattr(func, attr_name, wrapped)
return func
else:
@wraps(func)
def inner(*args, **kwargs):
name = getattr(args[0], self.name_attr) if self.name_attr is not None else self.name
token = _current_hero.set(name)
try:
return func(*args, **kwargs)
finally:
_current_hero.reset(token)
return inner
[docs]
def __enter__(self):
self._tokens.append(_current_hero.set(self.name))
return self
[docs]
def __exit__(self, *exc):
_current_hero.reset(self._tokens.pop())
[docs]
def _install_hero_record_factory() -> None:
"""Install a log record factory that stamps every record with the origin HERO name.
The factory wraps the currently installed factory (idempotently, guarded by a marker
attribute) and sets `record.origin` on every created record. Since the record factory is
invoked for each record at creation time, before any handler or propagation logic, this
also covers third-party loggers and loggers with `propagate=False`.
"""
current_factory = logging.getLogRecordFactory()
if getattr(current_factory, "_heros_origin_stamp", False):
return
def factory(*args, **kwargs):
record = current_factory(*args, **kwargs)
record.origin = resolve_origin()
return record
factory._heros_origin_stamp = True # ty:ignore
logging.setLogRecordFactory(factory)
[docs]
class Logger(logging.Logger):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if "HEROS_LOG" in os.environ:
self.setLevel(os.environ["HEROS_LOG"])
[docs]
def setLevel(self, level, globally=False):
"""Set logger level; optionally propagate to all existing loggers."""
if isinstance(level, str):
level = level.upper()
try:
level = int(level)
except ValueError:
pass
super().setLevel(level)
if globally:
os.environ["HEROS_LOG"] = str(level)
for logger in logging.root.manager.loggerDict.values():
if isinstance(logger, logging.Logger):
logger.setLevel(level)
[docs]
def spam(self, msg, *args, **kwargs):
"""Log a SPAM-level message."""
if self.isEnabledFor(SPAM):
self.log(SPAM, msg, *args, **kwargs)
# logger factory
[docs]
def get_logger(name: str = "heros") -> logging.Logger:
logging.setLoggerClass(Logger)
# stamp every log record with the origin HERO name
_install_hero_record_factory()
# Set up console handler only once
if not logging.getLogger().handlers:
console_handler = logging.StreamHandler()
console_handler.setFormatter(ColoredFormatter())
logging.getLogger().addHandler(console_handler)
logging.getLogger().setLevel(logging.INFO) # default to show SPAM messages
return logging.getLogger(name)
log = get_logger("heros")