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 = []