Accept explicit contextvars.Context in asyncio create_task() API

[asvetlov]

BPO	46994
Nosy	@vstinner, @asvetlov, @1st1
PRs	bpo-46994: Accept explicit contextvars.Context in asyncio create_task() API #31837

^{Note: these values reflect the state of the issue at the time it was migrated and might not reflect the current state.}

Show more details

GitHub fields:

assignee = None
closed_at = <Date 2022-03-14.11:54:32.003>
created_at = <Date 2022-03-12.10:33:08.202>
labels = ['3.11', 'type-feature', 'expert-asyncio']
title = 'Accept explicit contextvars.Context in asyncio create_task() API'
updated_at = <Date 2022-03-18.10:38:39.442>
user = 'https://github.com/asvetlov'

bugs.python.org fields:

activity = <Date 2022-03-18.10:38:39.442>
actor = 'vstinner'
assignee = 'none'
closed = True
closed_date = <Date 2022-03-14.11:54:32.003>
closer = 'asvetlov'
components = ['asyncio']
creation = <Date 2022-03-12.10:33:08.202>
creator = 'asvetlov'
dependencies = []
files = []
hgrepos = []
issue_num = 46994
keywords = ['patch']
message_count = 4.0
messages = ['414988', '415101', '415130', '415484']
nosy_count = 3.0
nosy_names = ['vstinner', 'asvetlov', 'yselivanov']
pr_nums = ['31837']
priority = 'normal'
resolution = 'fixed'
stage = 'resolved'
status = 'closed'
superseder = None
type = 'enhancement'
url = 'https://bugs.python.org/issue46994'
versions = ['Python 3.11']

[asvetlov]

Now asyncio creates a new context copy on task creation.

It is the perfect behavior *by default* and should stay as is.

However, sometimes passing an explicit context into a task and using back the context modified by task code is desired.

The most obvious use-case is testing: both unittest and pytest have multi-phase test initialization (asyncSetUp() methods and async fixtures correspondingly). If asyncSetUp() updates context variable -- test code expects to see this change.

Currently, unittest runs the setup-test-cleanup chain in a single task and uses an internal queue to push a new coroutine for execution and get results back. It works but is cumbersome.

Another solution is to create a task per test execution step and wrap the task creation with Context.run(). The problem is in getting the updated context back. A task creates a context copy on starting, thus the modified context is stored in the task internal attribute only. To get it back a trampoline async function should be used, e.g.

async def wrapper(coro):
    try:
        return await coro
    finally:
        context = contextvars.copy_context()
        # store the context copy somewhere

Again, it looks more complicated as it should be.

The proposal is:

1. Add 'context' keyword-only argument to asyncio.create_task() and loop.create_task().
2. Use this context if explicitly passed, create a copy of the current context otherwise.

The proposal is backward-compatible. Low-level API (call_soon(), call_later() etc.) already accept 'context' argument.

The attached PR demonstrates how the proposed API simplifies unittest.IsolatedAsyncioTestCase internals.

[1st1]

Yeah, +1 to add a parameter. Fwiw it was on my idea list when i was working on the pep

[asvetlov]

New changeset 9523c0d by Andrew Svetlov in branch 'main':
bpo-46994: Accept explicit contextvars.Context in asyncio create_task() API (GH-31837)
9523c0d

[vstinner]

Nice enhancement!

[nzig]

Any idea how this could be backported to older Python versions? I'm using 3.8. The best I've come up with is

lock = threading.Lock()

def create_task_with_context(coro, *, name= None, context = None):
    if context is None:
        return asyncio.create_task(coro, name=name)

    this_thread = threading.get_ident()

    def patched():
        if threading.get_ident() == this_thread:
            return context
        else:
            return original_copy_context()

    with lock:
        original_copy_context = contextvars.copy_context
        contextvars.copy_context = patched
        try:
            task = asyncio.tasks._PyTask(coro, name=name)
        finally:
            contextvars.copy_context = original_copy_context

    return task

Which is a horrible hack. Any better ideas?

Edit: Added thread safety check.

[agronholm]

Which is a horrible hack. Any better ideas?

Never mind the hackiness, this is outright dangerous. What if there's another event loop running in another thread, doing the same exact thing?