Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix context vars #302

Merged
merged 6 commits into from
Feb 18, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.rst
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@ Changes:
- ``structlog.threadlocal.wrap_dict()`` now has a correct type annotation.
`#290 <https://github.com/hynek/structlog/pull/290>`_

- Fixed bug with ``structlog.contextvars`` impl


----

Expand Down
4 changes: 2 additions & 2 deletions docs/contextvars.rst
Original file line number Diff line number Diff line change
Expand Up @@ -50,11 +50,11 @@ The general flow is:
>>> # Then use loggers as per normal
>>> # (perhaps by using structlog.get_logger() to create them).
>>> log.msg("hello")
a=1 b=2 event='hello'
event='hello' a=1 b=2
>>> # Use unbind_contextvars to remove a variable from the context
>>> unbind_contextvars("b")
>>> log.msg("world")
a=1 event='world'
event='world' a=1
>>> # And when we clear the threadlocal state again, it goes away.
>>> clear_contextvars()
>>> log.msg("hi there")
Expand Down
55 changes: 34 additions & 21 deletions src/structlog/contextvars.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
Python 3.7 as :mod:`contextvars`.

.. versionadded:: 20.1.0
.. versionchanged:: 14.0.0
Reimplement code without dict

See :doc:`contextvars`.
"""
Expand All @@ -15,12 +17,11 @@

from typing import Any, Dict

from .types import Context, EventDict, WrappedLogger
from .types import EventDict, WrappedLogger


_CONTEXT: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
"structlog_context"
)
STRUCTLOG_KEY_PREFIX = "structlog_"
_CONTEXT_VARS: Dict[str, contextvars.ContextVar[Any]] = {}


def merge_contextvars(
Expand All @@ -33,11 +34,17 @@ def merge_contextvars(
context-local context is included in all log calls.

.. versionadded:: 20.1.0
.. versionchanged:: 20.2.0 See toplevel note
"""
ctx = _get_context().copy()
ctx.update(event_dict)
ctx = contextvars.copy_context()

return ctx
for k in ctx:
if k.name.startswith(STRUCTLOG_KEY_PREFIX) and ctx[k] is not Ellipsis:
event_dict.setdefault(
k.name[len(STRUCTLOG_KEY_PREFIX) :], ctx[k]
) # noqa

return event_dict


def clear_contextvars() -> None:
Expand All @@ -48,9 +55,12 @@ def clear_contextvars() -> None:
handling code.

.. versionadded:: 20.1.0
.. versionchanged:: 20.2.0 See toplevel note
"""
ctx = _get_context()
ctx.clear()
ctx = contextvars.copy_context()
for k in ctx:
if k.name.startswith(STRUCTLOG_KEY_PREFIX):
k.set(Ellipsis)


def bind_contextvars(**kw: Any) -> None:
Expand All @@ -61,8 +71,17 @@ def bind_contextvars(**kw: Any) -> None:
context to be global (context-local).

.. versionadded:: 20.1.0
.. versionchanged:: 20.2.0 See toplevel note
"""
_get_context().update(kw)
for k, v in kw.items():
structlog_k = f"{STRUCTLOG_KEY_PREFIX}{k}"
try:
var = _CONTEXT_VARS[structlog_k]
except KeyError:
var = contextvars.ContextVar(structlog_k, default=Ellipsis)
_CONTEXT_VARS[structlog_k] = var

var.set(v)


def unbind_contextvars(*keys: str) -> None:
Expand All @@ -73,15 +92,9 @@ def unbind_contextvars(*keys: str) -> None:
remove keys from a global (context-local) context.

.. versionadded:: 20.1.0
.. versionchanged:: 20.2.0 See toplevel note
"""
ctx = _get_context()
for key in keys:
ctx.pop(key, None)


def _get_context() -> Context:
try:
return _CONTEXT.get()
except LookupError:
_CONTEXT.set({})
return _CONTEXT.get()
for k in keys:
structlog_k = f"{STRUCTLOG_KEY_PREFIX}{k}"
if structlog_k in _CONTEXT_VARS:
_CONTEXT_VARS[structlog_k].set(Ellipsis)
8 changes: 4 additions & 4 deletions src/structlog/stdlib.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@


try:
from . import contextvars
import contextvars
except ImportError:
contextvars = None # type: ignore

Expand Down Expand Up @@ -380,6 +380,7 @@ class AsyncBoundLogger:
It is useful to be able to log synchronously occasionally.

.. versionadded:: 20.2.0
.. versionchanged:: 20.2.0 fix _dispatch_to_sync contextvars usage
"""

__slots__ = ["sync_bl", "_loop"]
Expand Down Expand Up @@ -474,11 +475,10 @@ async def _dispatch_to_sync(
"""
Merge contextvars and log using the sync logger in a thread pool.
"""
ctx = contextvars._get_context().copy()
ctx.update(kw)
ctx = contextvars.copy_context()

await self._loop.run_in_executor(
self._executor, partial(meth, event, *args, **ctx)
self._executor, partial(ctx.run, partial(meth, event, *args, **kw))
)

async def debug(self, event: str, *args: Any, **kw: Any) -> None:
Expand Down
39 changes: 37 additions & 2 deletions tests/test_contextvars.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
# 2.0, and the MIT License. See the LICENSE file in the root of this
# repository for complete details.

import asyncio

import pytest

from structlog.contextvars import (
Expand Down Expand Up @@ -54,11 +56,11 @@ async def test_nested_async_bind(self, event_loop):

async def coro():
bind_contextvars(a=1)
await event_loop.create_task(nested_coro())
return merge_contextvars(None, None, {"b": 2})
return await event_loop.create_task(nested_coro())

async def nested_coro():
bind_contextvars(c=3)
return merge_contextvars(None, None, {"b": 2})

assert {"a": 1, "b": 2, "c": 3} == await event_loop.create_task(coro())

Expand Down Expand Up @@ -133,3 +135,36 @@ async def coro():
return merge_contextvars(None, None, {"b": 2})

assert {"b": 2} == await event_loop.create_task(coro())

async def test_parallel_binds(self, event_loop):
"""
Binding a variable causes it to be included in the result of
merge_contextvars.
"""

coro1_bind = asyncio.Event()
coro2_bind = asyncio.Event()

bind_contextvars(c=3)

async def coro1():
bind_contextvars(a=1)

coro1_bind.set()
await coro2_bind.wait()

return merge_contextvars(None, None, {"b": 2})

async def coro2():
bind_contextvars(a=2)

await coro1_bind.wait()
coro2_bind.set()

return merge_contextvars(None, None, {"b": 2})

coro1_task = event_loop.create_task(coro1())
coro2_task = event_loop.create_task(coro2())

assert {"a": 1, "b": 2, "c": 3} == await coro1_task
assert {"a": 2, "b": 2, "c": 3} == await coro2_task