diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4332c706..66200c55 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -41,6 +41,7 @@ Changes: You can keep using ``colorama`` to customize colors, of course. - The final processor can now return a ``bytearray`` (additionally to ``str`` and ``bytes``). `#344 `_ +- ``structlog.contextvars.bind_contextvars()`` now returns a mapping of keys to ``contextvars.Token``\s, allowing you to reset values using the new ``structlog.contextvars.reset_contextvars()``. ---- diff --git a/docs/api.rst b/docs/api.rst index 38af5cda..c5f46165 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -176,6 +176,7 @@ Please see :doc:`thread-local` for details. .. autofunction:: merge_contextvars .. autofunction:: clear_contextvars .. autofunction:: unbind_contextvars +.. autofunction:: reset_contextvars .. _procs: diff --git a/docs/conf.py b/docs/conf.py index 5a14e6d6..c07640f3 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -71,7 +71,10 @@ def find_version(*file_paths): ("py:class", "PlainFileObserver"), ("py:class", "structlog.threadlocal.TLLogger"), ("py:class", "TextIO"), + ("py:class", "TLLogger"), + ("py:class", "Token"), ("py:class", "traceback"), + ("py:class", "contextvars.Token"), ("py:class", "structlog._base.BoundLoggerBase"), ("py:class", "structlog.dev._Styles"), ("py:class", "structlog.types.EventDict"), @@ -163,4 +166,4 @@ def find_version(*file_paths): # Twisted's trac tends to be slow linkcheck_timeout = 300 -intersphinx_mapping = {"https://docs.python.org/3": None} +intersphinx_mapping = {"python": ("https://docs.python.org/3", None)} diff --git a/docs/contextvars.rst b/docs/contextvars.rst index 58d69605..d3649400 100644 --- a/docs/contextvars.rst +++ b/docs/contextvars.rst @@ -47,6 +47,7 @@ The general flow is: >>> # values: >>> clear_contextvars() >>> bind_contextvars(a=1, b=2) + {'a': at ...>, 'b': at ...>} >>> # Then use loggers as per normal >>> # (perhaps by using structlog.get_logger() to create them). >>> log.msg("hello") @@ -61,3 +62,18 @@ The general flow is: >>> clear_contextvars() >>> log.msg("hi there") event='hi there' a=None + + +If e.g. your request handler calls a helper function that needs to temporarily override some contextvars before restoring them back to their original values, you can use the :class:`~contextvars.contextvars.Token`\s returned by :func:`~structlog.contextvars.bind_contextvars` along with :func:`~structlog.contextvars.reset_contextvars` to accomplish this (much like how :meth:`contextvars.ContextVar.reset` works): + +.. code-block:: python + + def foo(): + bind_contextvars(a=1) + _helper() + log.msg("a is restored!") # a=1 + + def _helper(): + tokens = bind_contextvars(a=2) + log.msg("a is overridden") # a=2 + reset_contextvars(**tokens) diff --git a/src/structlog/contextvars.py b/src/structlog/contextvars.py index 91ea223d..9a05c47f 100644 --- a/src/structlog/contextvars.py +++ b/src/structlog/contextvars.py @@ -16,7 +16,7 @@ import contextvars -from typing import Any, Dict +from typing import Any, Dict, Mapping import structlog @@ -98,16 +98,22 @@ def clear_contextvars() -> None: k.set(Ellipsis) -def bind_contextvars(**kw: Any) -> None: - """ +def bind_contextvars(**kw: Any) -> "Mapping[str, contextvars.Token[Any]]": + r""" Put keys and values into the context-local context. Use this instead of :func:`~structlog.BoundLogger.bind` when you want some context to be global (context-local). + Return the mapping of ``contextvars.Token``\s resulting + from setting the backing ``ContextVar``\s. + Suitable for passing to :func:`reset_contextvars`. + .. versionadded:: 20.1.0 - .. versionchanged:: 21.1.0 See toplevel note. + .. versionchanged:: 21.1.0 Return the ``contextvars.Token`` mapping + rather than None. See also the toplevel note. """ + rv = {} for k, v in kw.items(): structlog_k = f"{STRUCTLOG_KEY_PREFIX}{k}" try: @@ -116,7 +122,21 @@ def bind_contextvars(**kw: Any) -> None: var = contextvars.ContextVar(structlog_k, default=Ellipsis) _CONTEXT_VARS[structlog_k] = var - var.set(v) + rv[k] = var.set(v) + + return rv + + +def reset_contextvars(**kw: "contextvars.Token[Any]") -> None: + r""" + Reset contextvars corresponding to the given Tokens. + + .. versionadded:: 21.1.0 + """ + for k, v in kw.items(): + structlog_k = f"{STRUCTLOG_KEY_PREFIX}{k}" + var = _CONTEXT_VARS[structlog_k] + var.reset(v) def unbind_contextvars(*keys: str) -> None: diff --git a/tests/test_contextvars.py b/tests/test_contextvars.py index 4cf3b998..01539dcd 100644 --- a/tests/test_contextvars.py +++ b/tests/test_contextvars.py @@ -16,6 +16,7 @@ get_contextvars, get_merged_contextvars, merge_contextvars, + reset_contextvars, unbind_contextvars, ) @@ -63,6 +64,30 @@ async def coro(): "d": 4, } == await event_loop.create_task(coro()) + async def test_reset(self, event_loop): + """ + reset_contextvars allows resetting contexvars to + previously-set values. + """ + + async def coro(): + bind_contextvars(a=1) + + assert {"a": 1} == get_contextvars() + + await event_loop.create_task(nested_coro()) + + async def nested_coro(): + tokens = bind_contextvars(a=2, b=3) + + assert {"a": 2, "b": 3} == get_contextvars() + + reset_contextvars(**tokens) + + assert {"a": 1} == get_contextvars() + + await event_loop.create_task(coro()) + async def test_nested_async_bind(self, event_loop): """ Context is passed correctly between "nested" concurrent operations.