Use contextvars for the log context and add log_context

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
 

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.

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",

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", "-")

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"