Add reset_contextvars and update bind_contextvars

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 <https://github.com/hynek/structlog/issues/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()``.
 
 
 ----

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:

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)}

docs/contextvars.rst

@@ -47,6 +47,7 @@ The general flow is:
    >>> # values:
    >>> clear_contextvars()
    >>> bind_contextvars(a=1, b=2)
+   {'a': <Token var=<ContextVar name='structlog_a' default=Ellipsis at ...> at ...>, 'b': <Token var=<ContextVar name='structlog_b' default=Ellipsis at ...> 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)

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:

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.