From 4aa209b7660fc124a0fcbbb83666f2421beef267 Mon Sep 17 00:00:00 2001
From: Joshua Bronson <jabronson@gmail.com>
Date: Mon, 2 Aug 2021 00:26:58 +0000
Subject: [PATCH] Add reset_contextvars and update bind_contextvars

structlog.contextvars.bind_contextvars() now returns a
Mapping[str, contextvar.Token] corresponding to the results
of setting the requested contextvars.

This return value is suitable for use with the newly-added
structlog.contextvars.reset_contextvars() function, which resets the
requested contextvars to their previous values using the given Tokens.
---
 docs/api.rst                 |  1 +
 docs/conf.py                 |  2 +-
 docs/contextvars.rst         | 21 +++++++++++++++++++++
 src/structlog/contextvars.py | 29 ++++++++++++++++++++++++-----
 tests/test_contextvars.py    | 20 ++++++++++++++++++++
 5 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/docs/api.rst b/docs/api.rst
index a4b4ae9a..d6b4ff37 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -172,6 +172,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 20ee2da1..45c005a4 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -163,4 +163,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 33bd079e..92dfc99e 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': <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,23 @@ 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..eeecb8f9 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]:
+    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,20 @@ 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) -> 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..c63f0a55 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,25 @@ 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.