[python][asyncio] Producer-Consumerパターンのサンプル

asyncio_evloop_producer_consumer.py

"""
asyncio のサンプルです。
asyncio.Queueを使った producer-consumer パターンのサンプル。
"""
import asyncio
import contextlib
import itertools
import random


async def make_producer(queue: asyncio.Queue):
    rnd = random.Random()
    for i in itertools.count():
        await asyncio.sleep(rnd.random())

        item = f'Item {i}'

        await queue.put(item)
        print(f'[producer] put {item}')


async def make_consumer(queue: asyncio.Queue, producer: asyncio.Future):
    rnd = random.Random()
    while True:
        if producer.cancelled() and queue.empty():
            break

        await asyncio.sleep(rnd.random())

        with contextlib.suppress(asyncio.TimeoutError):
            item = await asyncio.wait_for(queue.get(), timeout=1.0)
            print(f'====> [consumer] get {item}')


def all_cancel():
    for task in asyncio.Task.all_tasks():  # type: asyncio.Task
        task.cancel()
        with contextlib.suppress(asyncio.CancelledError):
            loop.run_until_complete(task)
        print(f'task done: {task}')


if __name__ == '__main__':
    loop = asyncio.get_event_loop()  # type: asyncio.events.AbstractEventLoop

    loop.set_debug(True)

    queue = asyncio.Queue(maxsize=5, loop=loop)
    producer = asyncio.ensure_future(make_producer(queue))
    consumer = asyncio.ensure_future(make_consumer(queue, producer))

    loop.call_later(10, loop.stop)
    loop.run_forever()

    # producer側をキャンセルして、consumerの終了を待機
    producer.cancel()
    print(f'CANCEL ==> producer')

    for task in (producer, consumer):
        with contextlib.suppress(asyncio.CancelledError):
            loop.run_until_complete(task)

    # 一気に全部キャンセルする場合
    # all_cancel()