Implement __aiter__ protocol on a Queue

asyncio/queues.py

@@ -10,6 +10,8 @@
 from . import locks
 from .coroutines import coroutine
 
+END_QUEUE = object()
+
 
 class QueueEmpty(Exception):
     """Exception raised when Queue.get_nowait() is called on a Queue object
@@ -81,6 +83,16 @@ def __repr__(self):
     def __str__(self):
         return '<{} {}>'.format(type(self).__name__, self._format())
 
+    def __aiter__(self):
+        return self
+
+    @coroutine
+    def __anext__(self):
+        val = yield from self.get()
+        if val is END_QUEUE:
+            raise StopAsyncIteration
+        return val
+
     def _format(self):
         result = 'maxsize={!r}'.format(self._maxsize)
         if getattr(self, '_queue', None):

tests/test_pep492.py

@@ -227,5 +227,44 @@ async def runner():
             self.loop.run_until_complete(runner())
 
 
+class QueueAsyncIteratorTests(BaseTest):
+    def test_get_iter(self):
+        async def consumer(queue, num_expected):
+            cnt = 0
+            async for item in queue:
+                cnt += 1
+                if cnt == num_expected:
+                    break
+
+        async def producer(queue, num_items):
+            for i in range(num_items):
+                await queue.put(i)
+
+        q = asyncio.Queue(loop=self.loop)
+        num_items = 5
+        self.loop.run_until_complete(
+            asyncio.gather(producer(q, 5),
+                           consumer(q, 5),
+                           loop=self.loop),
+            )
+
+    def test_stops_on_sentinel(self):
+        async def consumer(queue):
+            async for item in queue:
+                pass
+
+        async def producer(queue):
+            for i in range(5):
+                await queue.put(i)
+            await queue.put(asyncio.queues.END_QUEUE)
+
+        q = asyncio.Queue(loop=self.loop)
+        self.loop.run_until_complete(
+            asyncio.gather(producer(q),
+                           asyncio.wait_for(consumer(q), 5, loop=self.loop),
+                           loop=self.loop)
+            )
+
+
 if __name__ == '__main__':
     unittest.main()