Skip to content

Commit

Permalink
Reimpl contextvars code
Browse files Browse the repository at this point in the history
The old impl was broken, since python contextvars impl use shallow copy
to copy its context, and using a dict as a contextvar type ends up
sharing the same dict among different contexts
  • Loading branch information
Elad Namdar committed Feb 4, 2021
1 parent e191df7 commit fe3176f
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 25 deletions.
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
53 changes: 32 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,15 @@ 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 +53,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 +69,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 +90,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

0 comments on commit fe3176f

Please sign in to comment.