async_with_context.py

import contextvars

import asyncio
from contextlib import asynccontextmanager
from typing import TypeVar

T = TypeVar("T")


@asynccontextmanager
async def with_context(ctx: contextvars.ContextVar[T], value: T):
    token = ctx.set(value)
    try:
        yield value
    finally:
        ctx.reset(token)


async def test_with_context(sleep=0.05, iterations=5):
    _ctx = contextvars.ContextVar("test_context")

    async def _worker(name):
        assert _ctx.get() == "main"
        async with with_context(_ctx, name) as hey:
            for index in range(5):
                assert hey == name
                assert _ctx.get() == name
                await asyncio.sleep(sleep)

    async with with_context(_ctx, "main"):
        assert _ctx.get() == "main"
        tasks = asyncio.gather(_worker(name="foo"), _worker(name="bar"))
        await asyncio.sleep(sleep * (iterations / 2.0))
        assert _ctx.get() == "main"
        await tasks
        assert _ctx.get() == "main"


if __name__ == "__main__":
    asyncio.run(test_with_context())