Guangyang Li
Guangyang Li

Reputation: 2821

How to handle multiple results from a coroutine function?

I have some generators doing some searching stuff, and I use another generator wrapping them up:

def searching_stuff_1():
    # searching
    yield 1
    # and searching
    yield 2
    yield 3

def searching_stuff_2():
    yield 4
    yield 5


def gen():
    yield from searching_stuff_1()
    yield from searching_stuff_2()

for result in gen():
    print(result)

So now I am wondering how I can rewrite it into async version, that can yield multiple values in searching_stuff_1 and searching_stuff_2.

I was trying:

import asyncio

async def searching_stuff_1():
    f = asyncio.Future()
    result = []
    await asyncio.sleep(1)
    #searching
    result.append(1)
    #searching
    result.append(2)
    result.append(3)
    f.set_result(result)
    return f

async def searching_stuff_2():
    f = asyncio.Future()
    result = []
    await asyncio.sleep(1)
    result.append(4)
    result.append(5)
    f.set_result(result)
    return f

async def producer():
    coros = [searching_stuff_1(), searching_stuff_2()]
    for future in asyncio.as_completed(coros):
        yield await future

async def consumer(xs):
    async for future in xs:
        for r in future.result():
            print(r)
loop = asyncio.get_event_loop()
loop.run_until_complete(consumer(producer()))
loop.close()

However, in this version, I have to append all results into a list and wrap it in a Future instance. I am wondering if there is better way to handle multiple results from coroutine function. Is it possible that I can still yield those numbers?

Upvotes: 8

Views: 2510

Answers (1)

patpat
patpat

Reputation: 704

Yes you can still yield those numbers in async version, this is Asynchronous Generators, You can use async for(PEP492) and Asynchronous Comprehensions(PEP530), like this, rewrite from your first example. though this require python version higher than or equals to 3.6

import asyncio


async def searching_stuff_1():
    # searching
    yield 1
    # and searching
    yield 2
    yield 3


async def searching_stuff_2():
    yield 4
    yield 5


async def gen():
    async for i in searching_stuff_1():
        yield i
    async for i in searching_stuff_2():
        yield i


async def gen_all():
    return [i async for i in gen()]

if __name__ == "__main__":
    result = asyncio.get_event_loop().run_until_complete(gen_all())
    print(result)

for wishing to run two async generator async, you can use asyncio.gather.
But since asyncio.gather only collect coroutines' result in an async manner, you need to combine each yielded result from async generator separately with async def gen2 before calling asyncio.gather,

async def gen2(coro):
    return [i async for i in coro()]


# combine two async
async def gen_all2():
    return list(chain.from_iterable(await gather(gen2(searching_stuff_1), gen2(searching_stuff_2))))

To prove my point, we can change searching_stuff to:

async def searching_stuff_1():
    print("searching_stuff_1 begin")
    # searching
    yield 1
    await asyncio.sleep(1)
    # and searching
    yield 2
    yield 3
    print("searching_stuff_1 end")

async def searching_stuff_2():
    print("searching_stuff_2 begin")
    yield 4
    await asyncio.sleep(1)
    yield 5
    print("searching_stuff_2 end")

and then make the move:

result = asyncio.get_event_loop().run_until_complete(gen_all())
print(result)

> searching_stuff_2 begin
> searching_stuff_1 begin
> searching_stuff_2 end
> searching_stuff_1 end
> [1, 2, 3, 4, 5]

Upvotes: 3

Related Questions