Merge pull request #8 from gistart/pr/cancel-fix

minor cancel() improvements

Author: gistart

asyncio_pool/base_pool.py

@@ -143,6 +143,7 @@ async def _spawn(self, future, coro, cb=None, ctx=None):
             await self.semaphore.acquire()
         except BaseException as e:
             acq_error = True
+            coro.close()
             if not future.done():
                 future.set_exception(e)
         finally:
@@ -273,8 +274,10 @@ async def cancel(self, *futures, get_result=getres.flat):
                     tasks.append(task)
                     _futures.append(fut)
 
-        cancelled = sum(1 for task in tasks if task.cancel())
-        await aio.wait(tasks)  # let them actually cancel
+        cancelled = 0
+        if tasks:
+            cancelled = sum(1 for task in tasks if task.cancel())
+            await aio.wait(tasks)  # let them actually cancel
         # need to collect them anyway, to supress warnings
         return cancelled, [get_result(fut) for fut in _futures]
 

tests/test_base.py

@@ -1,17 +1,11 @@
 # coding: utf8
 
-import sys
 import pytest
-import warnings
 import asyncio as aio
-from itertools import chain
 from asyncio_pool import AioPool
 from async_timeout import timeout
 
 
-pytestmark = pytest.mark.filterwarnings('ignore:coroutine')
-
-
 @pytest.mark.asyncio
 async def test_concurrency():
     todo = range(1,21)
@@ -91,7 +85,6 @@ async def wrk_safe(*arg, **kw):
             await aio.sleep(0.5)
         except aio.CancelledError:
             await aio.sleep(0.1)  # simulate cleanup
-            pass
         return 1
 
     pool = AioPool(size=5)

tests/test_callbacks.py

@@ -1,14 +1,10 @@
 # coding: utf8
 
-import sys
 import pytest
 import asyncio as aio
 from asyncio_pool import AioPool, getres
 
 
-pytestmark = pytest.mark.filterwarnings('ignore:coroutine')
-
-
 async def cb(res, err, ctx):
     if err:
         exc, tb = err

tests/test_map.py

@@ -1,14 +1,10 @@
 # coding: utf8
 
-import sys
 import pytest
 import asyncio as aio
 from asyncio_pool import AioPool, getres
 
 
-pytestmark = pytest.mark.filterwarnings('ignore:coroutine')
-
-
 async def wrk(n):
     await aio.sleep(1 / n)
     return n*10
@@ -56,3 +52,20 @@ async def wrk(n):
             else:
                 assert False  # should not get here
             i += 1  # does not support enumerate btw (
+
+
+@pytest.mark.asyncio
+async def test_itermap_cancel():
+
+    async def wrk(n):
+        await aio.sleep(n / 100)
+        return n
+
+    todo = range(1, 101)
+
+    async with AioPool(5) as pool:
+        async for res in pool.itermap(wrk, todo, yield_when=aio.FIRST_COMPLETED):
+            if res == 13:
+                cancelled, _ = await pool.cancel()
+                break
+    assert cancelled == 100 - 13

tests/test_spawn.py

@@ -1,14 +1,10 @@
 # coding: utf8
 
-import sys
 import pytest
 import asyncio as aio
 from asyncio_pool import AioPool
 
 
-pytestmark = pytest.mark.filterwarnings('ignore:coroutine')
-
-
 @pytest.mark.asyncio
 async def test_spawns_behaviour():
     started = []