Skip to content

Commit

Permalink
Add reset_contextvars and update bind_contextvars
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
jab committed Aug 2, 2021
1 parent c5fe2e1 commit 4aa209b
Show file tree
Hide file tree
Showing 5 changed files with 67 additions and 6 deletions.
1 change: 1 addition & 0 deletions docs/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,7 @@ Please see :doc:`thread-local` for details.
.. autofunction:: merge_contextvars
.. autofunction:: clear_contextvars
.. autofunction:: unbind_contextvars
.. autofunction:: reset_contextvars


.. _procs:
Expand Down
2 changes: 1 addition & 1 deletion docs/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)}
21 changes: 21 additions & 0 deletions docs/contextvars.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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)
29 changes: 24 additions & 5 deletions src/structlog/contextvars.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@

import contextvars

from typing import Any, Dict
from typing import Any, Dict, Mapping

import structlog

Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down
20 changes: 20 additions & 0 deletions tests/test_contextvars.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
get_contextvars,
get_merged_contextvars,
merge_contextvars,
reset_contextvars,
unbind_contextvars,
)

Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit 4aa209b

Please sign in to comment.