How To Write Asynchronous Code With Contextvars Properly
Recently I contributed a project which using async/await syntax to write asynchronous code. After coding, testing, and fixing repeatedly, I summary out how to use write asynchronous code with contextvars properly.
What Is contextvars
Package contextvars comes from PEP-567 and provides APIs that allow developers to store, access and manage the local context. It’s similar to thread-local storage (TLS), but, unlike TLS, it also allows correctly keeping track of values per asynchronous task.
import asyncio
from contextvars import ContextVar, Context
var_what = ContextVar("what")
async def hello():
try:
return f"hello {var_what.get()}"
except LookupError:
return "hello?"
# evaluated by `python -m asyncio` instead of calling `asyncio.run`
assert await hello() == "hello?"
task_1 = asyncio.create_task(hello())
coroutine_1 = hello()
var_what.set("world")
task_2 = asyncio.create_task(hello())
coroutine_2 = hello()
assert await task_1 == "hello?"
assert await task_2 == "hello world"
assert await coroutine_1 == "hello world"
assert await coroutine_2 == "hello world"This example uses ContextVar object to store and access the value of the pre-defined field “what” from the current context.
And you see the problem came up from this example
that the objects of coroutine are unable to keep the context in track.
Because of only the objects of asyncio.Task will store the current context when they been created.
See also the document of contextvars for additional details.
Best Practices
If you want to set up a context, the first thing to do is creating some resources that stored in the context, then run some async tasks within the context, clean up resources from the context at last.
The best way to do that is by using contextlib.asynccontextmanager function as a decorator.
- It’s more pythonic that being able to use async-with syntax.
- It’s more readable that the initialization and cleanup processes are in the same function.
In the real situation, we need to keep the request-related data, such as security tokens and request data, share global resources, such as databases in web applications. So I code an example below that shows you how to use contextvars in asynchronous programming properly.
from contextlib import asynccontextmanager
from contextvars import copy_context, ContextVar
var_redis_client = ContextVar("redis_client")
@asynccontextmanager
async def create_app_context(settings):
# Initialization process
redis_client = object() # FIXME: create a Redis client
var_redis_client.set(redis_client)
try:
# Yield the current context which is set up.
yield copy_context()
finally:
# Cleanup process
# FIXME: Redis client cleanup
pass
# evaluated by `python -m asyncio` instead of calling `asyncio.run`
settings = {"redis_uri": "redis://..."}
async with create_app_context(settings) as app_ctx:
assert var_redis_client in copy_context()
# There are two ways to get the Redis client
# that been stored in the current context.
assert var_redis_client.get() is app_ctx[var_redis_client]Test Your Codes
Use pytest fixtures to set up the app context, is easy to access in test functions.
@pytest.mark.asynciodecorator comes from pytest-asyncio package.
import pytest
@pytest.fixture
def settings():
return {"redis_uri": "redis://..."}
@pytest.fixture
async def app_ctx(settings):
async with create_app_context(settings) as ctx:
# Use the yield_fixture to set up the app context.
yield ctx
@pytest.fixture
def redis_client(app_ctx):
return app_ctx[var_redis_client]
@pytest.mark.asyncio
async def test_without_context(app_ctx, redis_client):
assert var_redis_client not in copy_context()
# Although the current context does not have a Redis client,
# you still be able to access it in context from fixtures.
assert redis_client is app_ctx[var_redis_client]
def apply_context(ctx):
"""
Update the current context
"""
for var in ctx:
var.set(ctx[var])
@pytest.mark.asyncio
async def test_within_context(app_ctx, redis_client):
# Must update the current context,
# so you can use Redis client in the body of the testing function.
apply_context(app_ctx)
assert var_redis_client in copy_context()
assert redis_client is app_ctx[var_redis_client]The problem that came up from the above example is the context from the async yield_fixture is not propagated to the other fixtures or test functions. So it needs to pass the context explicitly.
I code a function called apply_context to apply the given context to the current context.
It needs to be executed first every time if need a context in fixtures or test functions.
IMHO, it is better if the context from async yield_fixtures is propagated implicitly to other fixtures or test functions when async yield_fixtures’ test function is marked by @pytest.mark.asyncio decorator.
So I create a pull request #161 to pytest-asyncio.
Hopefully, it can be merged into the master branch.
