contextuallogging

 17def context(
 18    function: Callable[..., Any] | None = None,
 19    /,
 20    *,
 21    keyword: str,
 22    key: str | None = None,
 23) -> Callable[..., Any]:
 24    """Decorator which adds logging context during a function call.
 25
 26    The logging context is added prior to calling the wrapped function and is reset
 27    after the function returns, creating a stack-like context in which the context added
 28    here will propagate to all logging calls made in all sub-functions.
 29
 30    Existing keys in the logging context will be overwritten by this decorator if they
 31    are already present, but the context will be reset to its original state once the
 32    wrapped function returns.
 33
 34    This decorator supports decorating async functions, as well. As this package uses
 35    ContextVars behind-the-scenes, the logging context can be propagated to async
 36    function calls in some cases and follows the same rules as ContextVars as to when
 37    that happens.
 38
 39    The created wrapper function raises a RuntimeError if the provided keyword name is
 40    not actually present in the passed-in kwargs.
 41
 42    Example usage:
 43
 44    >>> import sys
 45    >>> import logging
 46    >>> from contextuallogging import context, ContextualFormatter
 47    >>>
 48    >>> logger = logging.getLogger()
 49    >>> handler = logging.StreamHandler(sys.stdout)
 50    >>> handler.setFormatter(ContextualFormatter())
 51    >>> logger.addHandler(handler)
 52    >>> logger.setLevel(logging.INFO)
 53    >>>
 54    >>> @context(keyword="username", key="user")
 55    ... def inner_function(username: str) -> None:
 56    ...     logger.info("inner_function")
 57    ...
 58    >>> @context(keyword="request_id")
 59    ... def outer_function(request_id: str) -> None:
 60    ...     logger.info("outer_function")
 61    ...     inner_function(username="cfandrews")
 62    ...     logger.info("outer_function again")
 63    ...
 64    >>> outer_function(request_id="7fb9b341")
 65    {"level": "INFO", "logger": "root", "message": "outer_function", "request_id": "7fb9b341", "timestamp": "2023-11-25T20:56:41.796564Z"}
 66    {"level": "INFO", "logger": "root", "message": "inner_function", "request_id": "7fb9b341", "timestamp": "2023-11-25T20:56:41.797024Z", "user": "cfandrews"}
 67    {"level": "INFO", "logger": "root", "message": "outer_function again", "request_id": "7fb9b341", "timestamp": "2023-11-25T20:56:41.797075Z"}
 68
 69    Args:
 70        function (Callable[..., Any]): The function wrapped by this decorator.
 71        keyword (str): The name of the keyword argument to add to the logging context.
 72        key (str | None): The key to use in the logging context.
 73    """  # noqa: E501
 74    if function is None:
 75        return partial(context, keyword=keyword, key=key)
 76
 77    if inspect.iscoroutinefunction(function):
 78
 79        @wraps(function)
 80        async def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
 81            if keyword not in kwargs:
 82                message: Final[str] = f"Keyword argument not provided: {keyword}"
 83                raise RuntimeError(message)
 84            token: Final[Token[dict[str, object]]] = set_context_key(
 85                key=keyword if key is None else key,
 86                value=kwargs[keyword],
 87            )
 88            result: Final[Any] = await function(*args, **kwargs)
 89            reset_context(token=token)
 90            return result
 91
 92    else:
 93
 94        @wraps(function)
 95        def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
 96            if keyword not in kwargs:
 97                message: Final[str] = f"Keyword argument not provided: {keyword}"
 98                raise RuntimeError(message)
 99            token: Final[Token[dict[str, object]]] = set_context_key(
100                key=keyword if key is None else key,
101                value=kwargs[keyword],
102            )
103            result: Final[Any] = function(*args, **kwargs)
104            reset_context(token=token)
105            return result
106
107    return wrapper

 15class ContextualFormatter(Formatter):
 16    """Logging Formatter which can impart contextual information onto a LogRecord.
 17
 18    This class can be used in the same way as any regular logging formatter but will add
 19    any fields set by the contextual logging decorators.
 20
 21    This class is opinionated on its format. Specifically, it outputs log messages as
 22    JSON key-value mappings and all logged objects must be encode-able by the provided
 23    JSON encoder.
 24
 25    By default, the formatted logs will include the fields:
 26    * `timestamp`
 27    * `level`
 28    * `logger`
 29    * `message`
 30    * `exception`
 31     * If `exc_info` present on the LogRecord
 32    * `stack`
 33     * If `stack_info` present on the LogRecord
 34
 35    Example usage:
 36
 37    >>> import logging
 38    >>> import sys
 39    >>> from contextuallogging import ContextualFormatter
 40    >>>
 41    >>> logger = logging.getLogger()
 42    >>> handler = logging.StreamHandler(sys.stdout)
 43    >>> handler.setFormatter(ContextualFormatter())
 44    >>> logger.addHandler(handler)
 45    >>> logger.setLevel(logging.INFO)
 46    >>>
 47    >>> logger.info("Lorem ipsum dolor sit amet")
 48    {"level": "INFO", "logger": "root", "message": "Lorem ipsum dolor sit amet", "timestamp": "2023-11-25T22:25:42.803579Z"}
 49    """  # noqa: E501
 50
 51    def __init__(
 52        self,
 53        encoder: JSONEncoder | None = None,
 54    ) -> None:
 55        """Constructor.
 56
 57        The default encoder is instantiated as:
 58
 59        ```python
 60        JSONEncoder(
 61            skipkeys=True,
 62            ensure_ascii=False,
 63            sort_keys=True,
 64        )
 65        ```
 66
 67        Args:
 68            encoder (JSONEncoder | None): An optional JSONEncoder to use for formatting
 69                log messages.
 70        """
 71        self.__encoder: Final[JSONEncoder] = (
 72            JSONEncoder(skipkeys=True, ensure_ascii=False, sort_keys=True)
 73            if encoder is None
 74            else encoder
 75        )
 76        super().__init__()
 77
 78    def format(self, record: LogRecord) -> str:  # noqa: A003
 79        """Format the given record as a string.
 80
 81        Args:
 82            record (LogRecord): The LogRecord to format.
 83
 84        Returns:
 85            str: The formatted LogRecord.
 86        """
 87        data: Final[dict[str, object]] = {}
 88        timestamp: Final[datetime] = datetime.fromtimestamp(
 89            record.created,
 90            tz=timezone.utc,
 91        )
 92        data["timestamp"] = (
 93            f"{timestamp.year:04}-{timestamp.month:02}-{timestamp.day:02}T"
 94            f"{timestamp.hour:02}:{timestamp.minute:02}:{timestamp.second:02}."
 95            f"{timestamp.microsecond:06}Z"
 96        )
 97        data["level"] = record.levelname
 98        data["logger"] = record.name
 99        data["message"] = record.getMessage()
100        if record.exc_info is not None:
101            data["exception"] = record.exc_text
102        if record.stack_info is not None:
103            data["stack"] = self.formatStack(stack_info=record.stack_info)
104        return self.__encoder.encode({**data, **get_context()})

16def get_context() -> dict[str, object]:
17    """Gets the current logging context.
18
19    This is an advanced utility function which is not necessary for most use-cases.
20
21    Example usage:
22
23    >>> from contextuallogging import get_context, set_context_key
24    >>> get_context()
25    {}
26    >>> _ = set_context_key(key="key", value="value")
27    >>> get_context()
28    {'key': 'value'}
29
30    Returns:
31        dict[str, object]: The current logging context.
32    """
33    return _context.get()

36def set_context_key(key: str, value: object) -> Token[dict[str, object]]:
37    """Set the given key-value pair in the logging context.
38
39    This is an advanced utility function which is not necessary for most use-cases.
40
41    The returned `contextvars.Token` object can be passed to `reset_context` to reset
42    the logging context to its state as it was prior to this operation.
43
44    Example usage:
45
46    >>> from contextuallogging import get_context, set_context_key
47    >>> get_context()
48    {}
49    >>> _ = set_context_key(key="key", value="value")
50    >>> get_context()
51    {'key': 'value'}
52
53    Args:
54        key (str): The key of the logging field.
55        value (object): The value of the logging field.
56
57    Returns:
58        Token[dict[str, object]]: The Token required to reset the logging context to the
59            previous state.
60    """
61    return _context.set({**get_context(), key: value})

64def reset_context(token: Token[dict[str, object]]) -> None:
65    """Reset the logging context with the given token.
66
67    This is an advanced utility function which is not necessary for most use-cases.
68
69    The given token should be obtained from a `set_context_key` call and should never be
70    constructed from scratch manually.
71
72    Example usage:
73
74    >>> from contextuallogging import get_context, reset_context, set_context_key
75    >>> get_context()
76    {}
77    >>> reset_token = set_context_key(key="key", value="value")
78    >>> get_context()
79    {'key': 'value'}
80    >>> reset_context(token=reset_token)
81    >>> get_context()
82    {}
83
84    Args:
85        token (Token[dict[str, object]]): The reset Token.
86    """
87    _context.reset(token)