CODEWITHSUNDEEP
pythonasync-awaitpython-asynciocoroutinepython-contextvars
Rodrigo Oliveira
Rodrigo Oliveira

Reputation: 1582

How to run a coroutine inside a context?

In the Python docs about Context Vars a Context::run method is described to enable executing a callable inside a context so changes that the callable perform to the context are contained inside the copied Context. Though what if you need to execute a coroutine? What are you supposed to do in order to achieve the same behavior?

In my case, what I wanted was something like this to handle a transactional context with possible nested transactions:

my_ctxvar = ContextVar("my_ctxvar")

async def coro(func, transaction):
    token = my_ctxvar.set(transaction)
    r = await func()
    my_ctxvar.reset(token)  # no real need for this, but why not either
    return r

async def foo():
    ctx = copy_context()
    # simplification to one case here: let's use the current transaction if there is one
    if tx_owner := my_ctxvar not in ctx:
        tx = await create_transaction()
    else:
        tx = my_ctxvar.get()
    
    try:
        r = await ctx.run(coro)  # not actually possible
        if tx_owner:
            await tx.commit()
    except Exception as e:
        if tx_owner:
            await tx.rollback()
        raise from e
    return r

Upvotes: 7

Views: 1737

Answers (1)

alex_noname
alex_noname

Reputation: 32163

As I already pointed out here, context variables are natively supported by asyncio and are ready to be used without any extra configuration. It should be noted that:

  • Сoroutines executed by the current task by means of await share the same context
  • New spawned tasks by create_task are executed in the copy of parent task context.

Therefore, in order to execute a coroutine in a copy of the current context, you can execute it as a task:

await asyncio.create_task(coro())

Small example:

import asyncio
from contextvars import ContextVar

var = ContextVar('var')


async def foo():
    await asyncio.sleep(1)
    print(f"var inside foo {var.get()}")
    var.set("ham")  # change copy


async def main():
    var.set('spam')
    await asyncio.create_task(foo())
    print(f"var after foo {var.get()}")


asyncio.run(main())

var inside foo spam
var after foo spam

Upvotes: 10

Related Questions