CODEWITHSUNDEEP
python-3.xpython-asynciocoroutine
user2283347
user2283347

Reputation: 779

Context switching between python coroutines

I have a Python application in which multiple "tasks" will call exec function to evaluate multiple Python statements using a global "context" dictionary. The application would benefit tremendously by utilizing coroutines, but I could not find a way to create coroutine-specific context (global variables).

A simple example looks like this:

import asyncio

context = {}

async def async_exec(statements):
    global context
    for stmt in statements:
        exec(stmt, context, context)
        await asyncio.sleep(0)


def sync_main():
    asyncio.run(async_exec(['a=1', 'print(f"{a}==1")']))
    asyncio.run(async_exec(['a=2', 'print(f"{a}==2")']))

async def async_main():
    return await asyncio.gather(
        asyncio.create_task(async_exec(['a=1', 'print(f"{a}==1")'])),
        asyncio.create_task(async_exec(['a=2', 'print(f"{a}==2")'])),
    )


sync_main()
asyncio.run(async_main())

When I execute the code, I get

1==1
2==2
2==1 <- problem here
2==2

because the actual async execution sequence is

a=1
a=2
print(f"{a}==1")
print(f"{a}==2")

Is there any easy way to save and restore context variables for coroutines?

Working Code using Andrej's suggestion

Many thanks for the suggested use of contextvars, I am not sure how costly it would be to maintain fairly large dictionaries as contextvars but the following works for my application:

import asyncio
import contextvars

context = contextvars.ContextVar("global")

async def async_exec(statements):
    global context
    for stmt in statements:
        # retrieving the global dictionary
        gv = context.get({})
        # use the dictionary to evaluate statement
        exec(stmt, gv, gv)
        # set global dictionary back to the context
        context.set(gv)
        await asyncio.sleep(0)


def sync_main():
    asyncio.run(async_exec(['a=1', 'print(f"{a}==1")']))
    asyncio.run(async_exec(['a=2', 'print(f"{a}==2")']))

async def async_main():
    return await asyncio.gather(
        asyncio.create_task(async_exec(['a=1', 'print(f"{a}==1")'])),
        asyncio.create_task(async_exec(['a=2', 'print(f"{a}==2")'])),
    )

sync_main()
asyncio.run(async_main())

with output

1==1
2==2
1==1
2==2

Upvotes: 1

Views: 763

Answers (1)

Andrej Kesely
Andrej Kesely

Reputation: 195448

As far as I understand, you can use contextvars built-in module:

import asyncio
import contextvars

context = contextvars.ContextVar("Some context")


async def async_exec(statements):
    global context
    for stmt in statements:
        exec(stmt, {"context": context}, {"context": context})
        await asyncio.sleep(0)


def sync_main():
    asyncio.run(async_exec(["context.set(1)", 'print(f"{context.get()}==1")']))
    asyncio.run(async_exec(["context.set(2)", 'print(f"{context.get()}==2")']))


async def async_main():
    return await asyncio.gather(
        asyncio.create_task(
            async_exec(["context.set(1)", 'print(f"{context.get()}==1")'])
        ),
        asyncio.create_task(
            async_exec(["context.set(2)", 'print(f"{context.get()}==2")'])
        ),
    )


sync_main()
asyncio.run(async_main())

Prints correctly:

1==1
2==2
1==1
2==2

Upvotes: 1

Related Questions