From 18830a9348b16eb9a3020e121da027ad30513532 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 11:59:26 +0100 Subject: [PATCH] Use contextvars for the log context and add log_context ## Summary The global log context used `threading.local`, which does not follow `asyncio` tasks: concurrent tasks on one thread shared, and overwrote, each other's context. ## Changes - Store the context in a `contextvars.ContextVar`. Threads and tasks each get their own value. A new thread still starts empty, as before, and `set_log_context` now stores a copy of its arguments. - Add `log_context(**fields)`, a context manager that merges fields over the active context and restores the previous one on exit, including when the block raises. Blocks nest. - Export it from the package root, document it in the README and the changelog. - Tests: restoration, exceptions, nesting, records inside and outside the block, isolation between concurrent `asyncio` tasks, and the stored context being a copy. ## Verification - `pytest`: 111 passed (latest dependencies) and 93 passed with 18 skipped (no extras). `ruff` and `mypy --strict` are clean. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01U2nYHVGDdthQpcwNBPbp8r --- CHANGELOG.md | 6 +++ README.md | 18 +++++++- pylogkit/__init__.py | 2 + pylogkit/log_setup.py | 44 ++++++++++++++++---- tests/test_log_context.py | 86 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 146 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0c33e8b..d6beb4b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,6 +38,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). distribution name is settled (#20). - `setup_logging(use_emoji=...)` turns the level emoji in the verbose console layout off (#25). +- `log_context(**fields)` is a context manager that adds fields to the log + context for a `with` block and restores the previous context on exit, + including on exceptions (#21). ### Changed @@ -71,6 +74,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). `log_exception()`, `log_duration()` and `tqdm_logging()` are typed to accept a `ContextualLoggerAdapter` as well as a `Logger`. The `ImportError` raised when colorlog is missing now chains the original error (#17). +- The global log context is held in a `contextvars.ContextVar` instead of + `threading.local`, so it follows `asyncio` tasks as well as threads. A new + thread still starts with an empty context (#21). ### Deprecated diff --git a/README.md b/README.md index ddd42ef..e952db0 100644 --- a/README.md +++ b/README.md @@ -91,8 +91,8 @@ left out of the text layouts. Context is applied in three layers. Later layers win: -1. The global context of the current thread: `set_log_context(**fields)`, - `clear_log_context()` and `get_log_context()`. +1. The global context of the current thread or `asyncio` task: `set_log_context(**fields)`, + `clear_log_context()`, `get_log_context()` and the `log_context(**fields)` context manager. 2. The fields of a `ContextualLoggerAdapter`, added with `log.with_context(**fields)`. This returns a new adapter and never changes the global context. 3. `extra={...}` passed to an individual logging call. @@ -104,6 +104,20 @@ log.info("this line is alice's again") log.info("this one is carol's", extra={"user_id": "carol"}) ``` +`log_context` adds fields for the length of a `with` block and restores the previous context on +exit, even when the block raises. The context is held in a `contextvars.ContextVar`, so concurrent +`asyncio` tasks and threads do not see each other's values: + +```python +from pylogkit import log_context + +with log_context(request_id="abc"): + log.info("handled") # carries request_id=abc + with log_context(user_id="bob"): + log.info("nested") # carries both +log.info("outside") # carries neither +``` + ## Configuring `setup_logging()` `setup_logging()` returns a standard `logging.Logger`. Every argument is optional. diff --git a/pylogkit/__init__.py b/pylogkit/__init__.py index 1faac4a..db67aa6 100644 --- a/pylogkit/__init__.py +++ b/pylogkit/__init__.py @@ -14,6 +14,7 @@ ContextualLoggerAdapter, clear_log_context, get_log_context, + log_context, log_duration, log_exception, set_log_context, @@ -27,6 +28,7 @@ "ContextualLoggerAdapter", "clear_log_context", "get_log_context", + "log_context", "log_duration", "log_exception", "set_log_context", diff --git a/pylogkit/log_setup.py b/pylogkit/log_setup.py index 5c9834e..c46a2f9 100644 --- a/pylogkit/log_setup.py +++ b/pylogkit/log_setup.py @@ -3,7 +3,8 @@ :func:`setup_logging` configures console, file, JSON-lines and syslog output for a logger. Context fields (user, session, request, host, environment and process id) are attached to every record, globally through -:func:`set_log_context` or per logger through :class:`ContextualLoggerAdapter`. +:func:`set_log_context` and :func:`log_context`, or per logger through +:class:`ContextualLoggerAdapter`. Dependencies: @@ -27,10 +28,11 @@ import os import socket import sys -import threading import time import warnings from collections.abc import Callable, Iterable, Iterator, Sized +from contextlib import contextmanager +from contextvars import ContextVar from functools import cache, wraps from logging.handlers import RotatingFileHandler, SysLogHandler, TimedRotatingFileHandler from typing import IO, TYPE_CHECKING, Any, ParamSpec, TypeVar @@ -111,17 +113,43 @@ except ImportError: tqdm = None -_log_context = threading.local() +_log_context: ContextVar[dict[str, Any] | None] = ContextVar("pylogkit_log_context", default=None) def set_log_context(**kwargs: Any) -> None: - """Replace the global log context of the current thread with ``kwargs``.""" - _log_context.data = kwargs + """Replace the global log context of the current thread or task with ``kwargs``. + + The context is stored in a :class:`contextvars.ContextVar`, so each thread + and each :mod:`asyncio` task sees its own value. A new thread starts with + an empty context, and a task starts with a copy of its parent's. + """ + _log_context.set(dict(kwargs)) def clear_log_context() -> None: - """Remove all fields from the global log context of the current thread.""" - _log_context.data = {} + """Remove all fields from the global log context of the current thread or task.""" + _log_context.set({}) + + +@contextmanager +def log_context(**kwargs: Any) -> Iterator[None]: + """Add ``kwargs`` to the log context for the duration of a ``with`` block. + + Fields are merged over the active context. On exit, including when the + block raises, the previous context is restored exactly, so blocks nest and + concurrent :mod:`asyncio` tasks do not affect each other:: + + with log_context(request_id="abc"): + logger.info("handled") # carries request_id="abc" + + Args: + **kwargs: Fields to add to, or override in, the active context. + """ + token = _log_context.set({**(_log_context.get() or {}), **kwargs}) + try: + yield + finally: + _log_context.reset(token) @cache @@ -143,7 +171,7 @@ def get_log_context() -> dict[str, Any]: every call because they are cheap and may change (for example after a ``fork``). """ - context = getattr(_log_context, "data", {}).copy() + context = dict(_log_context.get() or {}) context.setdefault("user_id", "-") context.setdefault("session_id", "-") context.setdefault("request_id", "-") diff --git a/tests/test_log_context.py b/tests/test_log_context.py index 48362a4..b891a7f 100644 --- a/tests/test_log_context.py +++ b/tests/test_log_context.py @@ -79,3 +79,89 @@ def test_context_argument_of_setup_logging_sets_the_global_context(): setup_logging("ctx-arg", to_console=False, context={"user_id": "seeded"}) assert get_log_context()["user_id"] == "seeded" + + +def test_log_context_adds_fields_and_restores_them(): + """Fields apply inside the block and the previous context returns afterwards.""" + from pylogkit import log_context + + set_log_context(user_id="alice") + + with log_context(request_id="abc", user_id="bob"): + inside = get_log_context() + + after = get_log_context() + assert inside["request_id"] == "abc" + assert inside["user_id"] == "bob" + assert after["request_id"] == "-" + assert after["user_id"] == "alice" + + +def test_log_context_restores_on_exception(): + """The previous context is restored when the block raises.""" + import pytest + + from pylogkit import log_context + + set_log_context(user_id="alice") + + with pytest.raises(RuntimeError), log_context(user_id="bob"): + raise RuntimeError("boom") + + assert get_log_context()["user_id"] == "alice" + + +def test_log_context_nests(): + """Inner blocks add to the outer one and leave it intact on exit.""" + from pylogkit import log_context + + with log_context(request_id="outer"): + with log_context(user_id="bob"): + assert get_log_context()["request_id"] == "outer" + assert get_log_context()["user_id"] == "bob" + assert get_log_context()["user_id"] == "-" + assert get_log_context()["request_id"] == "outer" + + +def test_log_context_reaches_log_records(capsys): + """Records logged inside the block carry its fields.""" + import logging + + from pylogkit import log_context, setup_logging + + setup_logging("ctx-manager") + + with log_context(request_id="req-7"): + logging.getLogger("ctx-manager").info("inside") + logging.getLogger("ctx-manager").info("outside") + + inside, outside = capsys.readouterr().out.strip().splitlines() + assert "[request_id=req-7]" in inside + assert "request_id" not in outside + + +def test_context_is_isolated_between_asyncio_tasks(): + """Concurrent tasks must each keep their own context (``threading.local`` could not).""" + import asyncio + + from pylogkit import log_context + + async def handle(name: str) -> str: + with log_context(request_id=name): + await asyncio.sleep(0.01) + return str(get_log_context()["request_id"]) + + async def main() -> list[str]: + return await asyncio.gather(*(handle(f"req-{i}") for i in range(5))) + + assert asyncio.run(main()) == [f"req-{i}" for i in range(5)] + assert get_log_context()["request_id"] == "-" + + +def test_set_log_context_copies_its_arguments(): + """Later changes to the returned context must not leak into the stored one.""" + set_log_context(user_id="alice") + + get_log_context()["user_id"] = "mallory" + + assert get_log_context()["user_id"] == "alice"