Merge pull request #1554 from carver/wait-iter

New cancellable asynchronous iterators

Author: carver

p2p/cancellable.py

@@ -1,27 +1,28 @@
 from typing import (
+    AsyncIterator,
     Awaitable,
     TypeVar,
 )
 
 from cancel_token import CancelToken
 
+TReturn = TypeVar('TReturn')
+
 
 class CancellableMixin:
     cancel_token: CancelToken = None
 
-    _TReturn = TypeVar('_TReturn')
-
     async def wait(self,
-                   awaitable: Awaitable[_TReturn],
+                   awaitable: Awaitable[TReturn],
                    token: CancelToken = None,
-                   timeout: float = None) -> _TReturn:
+                   timeout: float = None) -> TReturn:
         """See wait_first()"""
         return await self.wait_first(awaitable, token=token, timeout=timeout)
 
     async def wait_first(self,
-                         *awaitables: Awaitable[_TReturn],
+                         *awaitables: Awaitable[TReturn],
                          token: CancelToken = None,
-                         timeout: float = None) -> _TReturn:
+                         timeout: float = None) -> TReturn:
         """
         Wait for the first awaitable to complete, unless we timeout or the token chain is triggered.
 
@@ -39,3 +40,30 @@ async def wait_first(self,
         else:
             token_chain = token.chain(self.cancel_token)
         return await token_chain.cancellable_wait(*awaitables, timeout=timeout)
+
+    async def wait_iter(
+            self,
+            aiterable: AsyncIterator[TReturn],
+            token: CancelToken = None,
+            timeout: float = None) -> AsyncIterator[TReturn]:
+        """
+        Iterate through an async iterator, raising the OperationCancelled exception if the token is
+        triggered. For example:
+
+        ::
+
+            async for val in self.wait_iter(my_async_iterator()):
+                do_stuff(val)
+
+        See :meth:`CancellableMixin.wait_first` for using arguments ``token`` and ``timeout``
+        """
+        aiter = aiterable.__aiter__()
+        while True:
+            try:
+                yield await self.wait(
+                    aiter.__anext__(),
+                    token=token,
+                    timeout=timeout,
+                )
+            except StopAsyncIteration:
+                break

tests/p2p/test_service.py

@@ -1,5 +1,6 @@
 import asyncio
 
+from cancel_token import OperationCancelled
 import pytest
 
 from p2p.service import BaseService
@@ -41,6 +42,33 @@ async def test_daemon_exit_causes_parent_cancellation():
     await asyncio.wait_for(service.events.cleaned_up.wait(), timeout=1)
 
 
+@pytest.mark.asyncio
+async def test_cancel_exits_async_generator():
+    service = WaitService()
+    asyncio.ensure_future(service.run())
+
+    async def cancel_soon():
+        await service.sleep(0.05)
+        await service.cancel()
+
+    asyncio.ensure_future(cancel_soon())
+
+    async def async_iterator():
+        yield 1
+        await asyncio.sleep(0.5)
+        assert False, "iterator should have been cancelled by now"
+
+    try:
+        async for val in service.wait_iter(async_iterator()):
+            assert val == 1
+    except OperationCancelled:
+        pass
+    else:
+        assert False, "iterator should have been cancelled during iteration"
+
+    await service.cancel()
+
+
 @pytest.mark.asyncio
 async def test_service_tasks_do_not_leak_memory():
     service = WaitService()