gh-119154: Simplify consumers by adding *.Queue.iter() and *.Queue.iter_nowait()

Doc/library/asyncio-queue.rst

@@ -70,6 +70,26 @@ Queue
       Return an item if one is immediately available, else raise
       :exc:`QueueEmpty`.
 
+   .. method:: iter()
+
+      Return an :term:`asynchronous iterator` which iterates over the queue
+      of items. If queue is empty, wait until the next item is available.
+
+      Stops iteration if the queue has been shut down and is empty, or if
+      the queue has been shut down immediately.
+
+      .. versionadded:: 3.14
+
+   .. method:: iter_nowait()
+
+      Return an :term:`iterator` which iterates over the queue of items
+      without blocking.
+
+      Only iterate over the items which are immediately available, but raise
+      the :exc:`QueueEmpty` exception if none are.
+
+      .. versionadded:: 3.14
+
    .. coroutinemethod:: join()
 
       Block until all items in the queue have been received and processed.
@@ -191,16 +211,11 @@ concurrent tasks::
 
 
    async def worker(name, queue):
-       while True:
-           # Get a "work item" out of the queue.
-           sleep_for = await queue.get()
-
+       # Get a "work item" out of the queue.
+       async for sleep_for in queue.iter():
            # Sleep for the "sleep_for" seconds.
            await asyncio.sleep(sleep_for)
 
-           # Notify the queue that the "work item" has been processed.
-           queue.task_done()
-
            print(f'{name} has slept for {sleep_for:.2f} seconds')
 
 
@@ -215,22 +230,16 @@ concurrent tasks::
            total_sleep_time += sleep_for
            queue.put_nowait(sleep_for)
 
-       # Create three worker tasks to process the queue concurrently.
-       tasks = []
-       for i in range(3):
-           task = asyncio.create_task(worker(f'worker-{i}', queue))
-           tasks.append(task)
+       # All tasks have been queued
+       queue.shutdown()
 
-       # Wait until the queue is fully processed.
+       # Create three worker tasks to process the queue concurrently.
        started_at = time.monotonic()
-       await queue.join()
-       total_slept_for = time.monotonic() - started_at
+       async with asyncio.TaskGroup() as tg:
+           for i in range(3):
+               tg.create_task(worker(f'worker-{i}', queue))
 
-       # Cancel our worker tasks.
-       for task in tasks:
-           task.cancel()
-       # Wait until all worker tasks are cancelled.
-       await asyncio.gather(*tasks, return_exceptions=True)
+       total_slept_for = time.monotonic() - started_at
 
        print('====')
        print(f'3 workers slept in parallel for {total_slept_for:.2f} seconds')

Doc/library/queue.rst

@@ -132,6 +132,49 @@ provide the public methods described below.
    will not block.  Similarly, if full() returns ``False`` it doesn't
    guarantee that a subsequent call to put() will not block.
 
+.. method:: Queue.iter(block=True, timeout=None)
+
+   Return an :term:`iterator` which iterates over the queue of items.  If optional
+   args *block* is true and *timeout* is ``None`` (the default), block if necessary
+   until the next item is available.  If *timeout* is a positive number, it blocks
+   at most *timeout* seconds and stops iteration if no item was available within that
+   time, but if the first item is not available the :exc:`Empty` exception is raised.
+   Otherwise (*block* is false), iterate over all item which are immediately
+   available, but raise the :exc:`Empty` exception if none are (*timeout* is ignored
+   in that case).
+
+   Stops iteration if the queue has been shut down and is empty, or if the queue has
+   been shut down immediately.
+
+   Example::
+
+      import concurrent.futures
+      import queue
+      import time
+
+      def worker(name, q):
+          for item in q.iter():
+              time.sleep(.01)
+              print(f'{name} finished {item}')
+
+      q = queue.Queue()
+      for item in range(30):
+          q.put(item)
+
+      q.shutdown()
+      with concurrent.futures.ThreadPoolExecutor() as tp:
+          for i in range(3):
+              tp.submit(worker, f'worker-{i}', q)
+
+      print('All work completed')
+
+   .. versionadded:: 3.14
+
+.. method:: Queue.iter_nowait()
+
+   Equivalent to ``iter(False)``.
+
+   .. versionadded:: 3.14
 
 .. method:: Queue.put(item, block=True, timeout=None)
 
@@ -212,8 +255,7 @@ Example of how to wait for enqueued tasks to be completed::
     q = queue.Queue()
 
     def worker():
-        while True:
-            item = q.get()
+        for item in q:
             print(f'Working on {item}')
             print(f'Finished {item}')
             q.task_done()

Doc/whatsnew/3.14.rst

@@ -92,6 +92,12 @@ ast
 Added :func:`ast.compare` for comparing two ASTs.
 (Contributed by Batuhan Taskaya and Jeremy Hylton in :issue:`15987`.)
 
+asyncio
+-------
+
+Add :meth:`asyncio.Queue.iter` and :meth:`asyncio.Queue.iter_nowait`.
+(Contributed by Wannes Boeykens in :gh:`119154`.)
+
 os
 --
 
@@ -107,6 +113,12 @@ pathlib
   another, like :func:`shutil.copyfile`.
   (Contributed by Barney Gale in :gh:`73991`.)
 
+queue
+-----
+
+Add :meth:`queue.Queue.iter` and :meth:`queue.Queue.iter_nowait`.
+(Contributed by Wannes Boeykens in :gh:`119154`.)
+
 symtable
 --------
 

Lib/asyncio/queues.py

@@ -30,6 +30,20 @@ class QueueShutDown(Exception):
     pass
 
 
+class _AsyncQueueIterator:
+    def __init__(self, queue):
+        self._queue = queue
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return await self._queue.get()
+        except QueueShutDown:
+            raise StopAsyncIteration
+
+
 class Queue(mixins._LoopBoundMixin):
     """A queue, useful for coordinating producer and consumer coroutines.
 
@@ -216,6 +230,33 @@ def get_nowait(self):
         self._wakeup_next(self._putters)
         return item
 
+    def iter(self):
+        """Return an asynchronous iterator which iterates over the queue of
+        items. If queue is empty, wait until the next item is available.
+
+        Stops iteration if the queue has been shut down and is empty, or if the
+        queue has been shut down immediately.
+        """
+        return _AsyncQueueIterator(self)
+
+    def iter_nowait(self):
+        """Return an iterator which iterates over the queue of items without
+        blocking.
+
+        Only iterate over the items which are immediately available, but raise
+        the QueueEmpty exception if none are.
+        """
+        try:
+            yield self.get_nowait()
+        except QueueShutDown:
+            return
+
+        try:
+            while True:
+                yield self.get_nowait()
+        except (QueueEmpty, QueueShutDown):
+            return
+
     def task_done(self):
         """Indicate that a formerly enqueued task is complete.
 

Lib/queue.py

@@ -275,6 +275,41 @@ def _put(self, item):
     def _get(self):
         return self.queue.popleft()
 
+    def iter(self, block=True, timeout=None):
+        '''Return an iterator which iterates over the queue of items.
+
+        If optional args 'block' is true and 'timeout' is None (the default),
+        block if necessary until the next item is available. If 'timeout' is a
+        positive number, it blocks at most 'timeout' seconds and stops
+        iteration if no item was available within that time, but if the first
+        item is not available the Empty exception is raised. Otherwise ('block'
+        is false), iterate over all item which are immediately available, but
+        raise the Empty exception if none are ('timeout' is ignored in that
+        case).
+
+        Stops iteration if the queue has been shut down and is empty, or if the
+        queue has been shut down immediately.
+        '''
+        try:
+            yield self.get(block=block, timeout=timeout)
+        except ShutDown:
+            return
+
+        try:
+            while True:
+                yield self.get(block=block, timeout=timeout)
+        except (Empty, ShutDown):
+            return
+
+    def iter_nowait(self):
+        '''Return an iterator which iterates over the queue of items without
+        blocking.
+
+        Only iterate over the items which are immediately available, but raise
+        the Empty exception if none are.
+        '''
+        return self.iter(block=False)
+
     __class_getitem__ = classmethod(types.GenericAlias)
 
 

Lib/test/test_asyncio/test_queues.py

@@ -116,6 +116,37 @@ async def putter():
         self.assertTrue(t.done())
         self.assertTrue(t.result())
 
+    async def test_iter(self):
+        q = asyncio.Queue()
+        accumulator = 0
+
+        async def worker():
+            nonlocal accumulator
+
+            async for item in q.iter():
+                accumulator += item
+
+        async with asyncio.TaskGroup() as tg:
+            tg.create_task(worker())
+            tg.create_task(worker())
+            for i in range(100):
+                q.put_nowait(i)
+
+            q.shutdown()
+
+        self.assertEqual(sum(range(100)), accumulator)
+
+    async def test_iter_nowait(self):
+        q = asyncio.Queue()
+        accumulator = 0
+        for i in range(100):
+            q.put_nowait(i)
+
+        for item in q.iter_nowait():
+            accumulator += item
+
+        self.assertEqual(sum(range(100)), accumulator)
+
 
 class QueueGetTests(unittest.IsolatedAsyncioTestCase):
 
@@ -167,6 +198,8 @@ def test_nonblocking_get(self):
     def test_nonblocking_get_exception(self):
         q = asyncio.Queue()
         self.assertRaises(asyncio.QueueEmpty, q.get_nowait)
+        with self.assertRaises(asyncio.QueueEmpty):
+            list(q.iter_nowait())
 
     async def test_get_cancelled_race(self):
         q = asyncio.Queue()
@@ -471,27 +504,22 @@ async def test_task_done(self):
 
         # Two workers get items from the queue and call task_done after each.
         # Join the queue and assert all items have been processed.
-        running = True
 
         async def worker():
             nonlocal accumulator
 
-            while running:
-                item = await q.get()
+            async for item in q.iter():
                 accumulator += item
                 q.task_done()
 
         async with asyncio.TaskGroup() as tg:
-            tasks = [tg.create_task(worker())
-                     for index in range(2)]
-
+            tg.create_task(worker())
+            tg.create_task(worker())
             await q.join()
             self.assertEqual(sum(range(100)), accumulator)
 
             # close running generators
-            running = False
-            for i in range(len(tasks)):
-                q.put_nowait(0)
+            q.shutdown()
 
     async def test_join_empty_queue(self):
         q = self.q_class()
@@ -566,6 +594,7 @@ async def test_shutdown_empty(self):
             await q.get()
         with self.assertRaisesShutdown():
             q.get_nowait()
+        self.assertEqual(list(q.iter_nowait()), [])
 
     async def test_shutdown_nonempty(self):
         # Test shutting down a non-empty queue
@@ -610,6 +639,7 @@ async def test_shutdown_nonempty(self):
             await q.get()
         with self.assertRaisesShutdown():
             q.get_nowait()
+        self.assertEqual(list(q.iter_nowait()), [])
 
         # Ensure there is 1 unfinished task, and join() task succeeds
         q.task_done()
@@ -652,6 +682,7 @@ async def test_shutdown_immediate(self):
             await q.get()
         with self.assertRaisesShutdown():
             q.get_nowait()
+        self.assertEqual(list(q.iter_nowait()), [])
 
         # Ensure there are no unfinished tasks
         with self.assertRaises(