CODEWITHSUNDEEP
pythonpython-asyncio
ivaigult
ivaigult

Reputation: 6667

How to check what coroutine is completed after asyncio.wait

Consider the following code:

import random
import asyncio

class RandomLife(object):
    def __init__(self, name: str):
        self.name = name
        self.coro = asyncio.sleep(random.randrange(0, 5))

    def __await__(self):
        return self.coro.__await__()


async def main():
    objects = [RandomLife("one"), RandomLife("two"), RandomLife("three")]

    finished, unfinished = await asyncio.wait(objects, return_when=asyncio.FIRST_COMPLETED)

    print(finished)

    await asyncio.wait(unfinished)

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

After then first asyncio.wait I want to know what instance of RandomLife has completed. But the finished variable is a set of Task s, rather than a RandomLife instance. How do I convert this task to a RandomLife? Is it possible?

Upvotes: 2

Views: 2336

Answers (1)

ivaigult
ivaigult

Reputation: 6667

As the documentation warns:

Note wait() schedules coroutines as Tasks automatically and later returns those implicitly created Task objects in (done, pending) sets. Therefore the following code won’t work as expected:

async def foo():
    return 42

coro = foo()
done, pending = await asyncio.wait({coro})

if coro in done:
    # This branch will never be run!

Here is how the above snippet can be fixed:

async def foo():
    return 42

task = asyncio.create_task(foo())
done, pending = await asyncio.wait({task})

if task in done:
    # Everything will work as expected now.

We can employ the same trick. First, we need to wrap all the coroutines to tasks, and then set up mapping from a task created to its RandomLife instance:

import random
import asyncio

class RandomLife(object):
    def __init__(self, name: str):
        self.name = name
        self.coro = asyncio.sleep(random.randrange(0, 5))

    def __await__(self):
        return self.coro.__await__()


async def main():
    objects = [RandomLife("one"), RandomLife("two"), RandomLife("three")]

    # Wrap all the coros to tasks, as the documentation suggests.
    tasks = [asyncio.create_task(o.coro) for o in objects]
    # Set up mapping from tasks created to RandomLife instances.
    task2life = dict(zip(tasks, objects))

    finished, unfinished = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)

    # Get first task finished.
    finised_task = list(finished)[0]
    # Map it back to the RandomLife instance.
    finished_life = task2life[finised_task]

    print(finished_life.name)

    await asyncio.wait(unfinished)

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

Upvotes: 1

Related Questions