structio/tests/queue.py

import time
import threading
import structio


async def producer(q: structio.Queue, n: int):
    for i in range(n):
        # This will wait until the
        # queue is emptied by the
        # consumer
        await q.put(i)
        print(f"Produced {i}")
    await q.put(None)
    print("Producer done")


async def consumer(q: structio.Queue):
    while True:
        # Hangs until there is
        # something on the queue
        item = await q.get()
        if item is None:
            print("Consumer done")
            break
        print(f"Consumed {item}")
        # Simulates some work so the
        # producer waits before putting
        # the next value
        await structio.sleep(1)


def threaded_consumer(q: structio.thread.AsyncThreadQueue):
    while True:
        # Hangs until there is
        # something on the queue
        item = q.get_sync()
        if item is None:
            print("Consumer done")
            break
        print(f"Consumed {item}")
        # Simulates some work so the
        # producer waits before putting
        # the next value
        time.sleep(1)


async def main(q: structio.Queue, n: int):
    print("Starting consumer and producer")
    async with structio.create_pool() as ctx:
        ctx.spawn(producer, q, n)
        ctx.spawn(consumer, q)
    print("Bye!")


async def main_threaded(q: structio.thread.AsyncThreadQueue, n: int):
    print("Starting consumer and producer")
    async with structio.create_pool() as pool:
        pool.spawn(producer, q, n)
        await structio.thread.run_in_worker(threaded_consumer, q)
    print("Bye!")


if __name__ == "__main__":
    queue = structio.Queue(2)  # Queue has size limit of 2
    structio.run(main, queue, 5)
    queue = structio.thread.AsyncThreadQueue(2)
    structio.run(main_threaded, queue, 5)