new feature: task cancellation; license and setup.py added

Author: gistart

.gitignore

@@ -8,6 +8,9 @@
 .pypyenv*
 .pytest_*
 .mypy_cache
+dist
+build
+*.egg-info
 __pycache__
 local_settings.py
 

LICENSE.txt

@@ -0,0 +1,19 @@
+Copyright (c) 2018 Oleg Mihaylov (gistart)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
+of the Software, and to permit persons to whom the Software is furnished to do
+so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -41,20 +41,22 @@ async def spawn_usage(todo=range(1,4)):
  futures.append(fut)
  # At this point some of the workers already started.
 
- # Context manager calls `join()` at exit, so this will finish when all
+ # Context manager calls `join` at exit, so this will finish when all
  # workers return, crash or cancelled.
 
  assert sum(todo) == sum(fut.result() for fut in futures) # all done
 
 
 async def map_usage(todo=range(100)):
  pool = AioPool(size=10)
- # Joins internally, collects results from all spawned workers,
+ # Waits and collects results from all spawned workers,
  # returns them in same order as `todo`, if worker crashes or cancelled:
  # returns exception object as a result.
  # Basically, it wraps `spawn_usage` code into one call.
  results = await pool.map(worker, todo)
 
+ # await pool.join() # is not needed here, bcs no other tasks were spawned
+
  assert isinstance(results[0], ZeroDivisionError) \
  and sum(results[1:]) == sum(todo)
 
@@ -77,6 +79,61 @@ async def callbacks_usage():
  pass # TODO
 
 
+async def exec_usage(todo=range(1,11)):
+ async with AioPool(size=4) as pool:
+ futures = await pool.map_n(worker, todo)
+
+ # While other workers are waiting or active, you can "synchronously"
+ # execute one task. It does not interrupt others, just waits for pool
+ # space, then waits for task to finish and then returns it's result.
+ important_res = await pool.exec(worker(2))
+ assert 2 == important_res
+
+ # You can continue working as usual:
+ moar = await pool.spawn(worker(10))
+
+ assert sum(todo) == sum(f.result() for f in futures)
+
+
+async def cancel_usage():
+
+ async def wrk(*arg, **kw):
+ await aio.sleep(0.5)
+ return 1
+
+ pool = AioPool(size=2)
+
+ f_quick = await pool.spawn_n(aio.sleep(0.1))
+ f12 = await pool.spawn(wrk()), await pool.spawn_n(wrk())
+ f35 = await pool.map_n(wrk, range(3))
+
+ # At this point, if you cancel futures, returned by pool methods,
+ # you just won't be able to retrieve spawned task results, task
+ # themselves will continue working. Don't do this:
+ # f_quick.cancel()
+ # use `pool.cancel` instead:
+
+ # cancel some
+ await aio.sleep(0.1)
+ cancelled, results = await pool.cancel(f12[0], f35[2]) # running and waiting
+ assert 2 == cancelled # none of them had time to finish
+ assert 2 == len(results) and \
+ all(isinstance(res, aio.CancelledError) for res in results)
+
+ # cancel all others
+ await aio.sleep(0.1)
+
+ # not interrupted and finished successfully
+ assert f_quick.done() and f_quick.result() is None
+
+ cancelled, results = await pool.cancel() # all
+ assert 3 == cancelled
+ assert len(results) == 3 and \
+ all(isinstance(res, aio.CancelledError) for res in results)
+
+ assert await pool.join() # joins successfully
+
+
 async def details(todo=range(1,11)):
  pool = AioPool(size=5)
 
@@ -109,15 +166,15 @@ async def details(todo=range(1,11)):
  await pool.join()
 
  assert all(f.done() for f in itertools.chain(f1,f2)) # this is guaranteed
-
+ assert 2 * sum(todo) == sum(f.result() for f in itertools.chain(f1,f2))
 ```
 
 
 ## TODO:
 
 - [ ] callbacks
-- [ ] setup
-- [ ] cancelled, timeouts
+- [x] setup
+- [x] cancelled, timeouts
 - [x] tests
 - [ ] loadtests
 - [x] usage

asyncio_pool/__init__.py

@@ -1,6 +1,8 @@
 # coding: utf8
 
 import sys
+
+from .utils import result_noraise
 from .base_pool import BaseAioPool
 
 
@@ -9,6 +11,6 @@
  class AioPool(BaseAioPool): pass
 
 else:
- from .mx_asynciter import MxAsyncIterPool
+ from .mx_asyncgen import MxAsyncGenPool
 
- class AioPool(MxAsyncIterPool, BaseAioPool): pass
+ class AioPool(MxAsyncGenPool, BaseAioPool): pass

asyncio_pool/base_pool.py

@@ -4,18 +4,21 @@
 import traceback
 import collections
 import asyncio as aio
-from .utils import _get_future_result
+from .utils import result_noraise
 
 
 class BaseAioPool(object):
+ ''' BaseAioPool implements features, supposed to work in all supported
+ python versions. Other features supposed to be implemented as mixins.'''
 
  def __init__(self, size=1024, *, loop=None):
  self.loop = loop or aio.get_event_loop()
 
  self.size = size
  self._executed = 0
- self._joined = collections.deque()
- self._waiting = collections.deque()
+ self._joined = set()
+ self._waiting = {} # future -> task
+ self._spawned = {} # future -> task
  self.semaphore = aio.Semaphore(value=self.size, loop=self.loop)
 
  async def __aenter__(self):
@@ -41,7 +44,7 @@ async def join(self):
  return True
 
  fut = self.loop.create_future()
- self._joined.append(fut)
+ self._joined.add(fut)
  try:
  return await fut
  finally:
@@ -72,33 +75,46 @@ async def _wrap(self, coro, future, cb=None, ctx=None):
  return
 
  self.semaphore.release()
- if not exc:
- future.set_result(res)
- else:
- future.set_exception(exc)
 
+ if not future.done():
+ if exc:
+ future.set_exception(exc)
+ else:
+ future.set_result(res)
+
+ del self._spawned[future]
  if self.is_empty:
  self._release_joined()
 
  async def _spawn(self, future, coro, cb=None, ctx=None):
+ acq_error = False
  try:
  await self.semaphore.acquire()
  except Exception as e:
- future.set_exception(e)
- self._waiting.remove(future)
- wrapped = self._wrap(coro, future, cb=cb, ctx=ctx)
- self.loop.create_task(wrapped)
+ acq_error = True
+ if not future.done():
+ future.set_exception(e)
+ finally:
+ del self._waiting[future]
+
+ if future.done():
+ if not acq_error and future.cancelled(): # outside action
+ self.semaphore.release()
+ else: # all good, can spawn now
+ wrapped = self._wrap(coro, future, cb=cb, ctx=ctx)
+ task = self.loop.create_task(wrapped)
+ self._spawned[future] = task
  return future
 
  async def spawn_n(self, coro, cb=None, ctx=None):
  future = self.loop.create_future()
- self._waiting.append(future)
- self.loop.create_task(self._spawn(future, coro, cb=cb, ctx=ctx))
+ task = self.loop.create_task(self._spawn(future, coro, cb=cb, ctx=ctx))
+ self._waiting[future] = task
  return future
 
  async def spawn(self, coro, cb=None, ctx=None):
  future = self.loop.create_future()
- self._waiting.append(future)
+ self._waiting[future] = self.loop.create_future() # TODO omg ???
  return await self._spawn(future, coro, cb=cb, ctx=ctx)
 
  async def exec(self, coro, cb=None, ctx=None):
@@ -113,16 +129,36 @@ async def map_n(self, fn, iterable):
 
  async def map(self, fn, iterable, exc_as_result=True):
  futures = await self.map_n(fn, iterable)
- await self.join()
-
- results = []
- for fut in futures:
- res = _get_future_result(fut, exc_as_result)
- results.append(res)
- return results
+ await aio.wait(futures)
+ return [result_noraise(fut, exc_as_result) for fut in futures]
 
  async def iterwait(self, *arg, **kw): # TODO there's a way to support 3.5?
  raise NotImplementedError('python3.6+ required')
 
  async def itermap(self, *arg, **kw): # TODO there's a way to support 3.5?
  raise NotImplementedError('python3.6+ required')
+
+ def _cancel(self, *futures):
+ tasks, _futures = [], []
+
+ if not len(futures): # meaning cancel all
+ tasks.extend(self._waiting.values())
+ tasks.extend(self._spawned.values())
+ _futures.extend(self._waiting.keys())
+ _futures.extend(self._spawned.keys())
+ else:
+ for fut in futures:
+ task = self._spawned.get(fut, self._waiting.get(fut))
+ if task:
+ tasks.append(task)
+ _futures.append(fut)
+
+ cancelled = sum([1 for fut in tasks if fut.cancel()])
+ return cancelled, _futures
+
+ async def cancel(self, *futures, exc_as_result=True):
+ cancelled, _futures = self._cancel(*futures)
+ await aio.sleep(0) # let them actually cancel
+ # need to collect them anyway, to supress warnings
+ results = [result_noraise(fut, exc_as_result) for fut in _futures]
+ return cancelled, results

asyncio_pool/mx_asynciter.py renamed to asyncio_pool/mx_asyncgen.py

@@ -1,11 +1,12 @@
 # coding: utf8
-'''Mixin for BaseAioPool with async iterator features python3.6+'''
+'''Mixin for BaseAioPool with async generator features, python3.6+'''
 
 import asyncio as aio
-from .utils import _get_future_result
+from .utils import result_noraise
 
 
-class MxAsyncIterPool(object):
+class MxAsyncGenPool(object):
+ ''' Asynchronous generator wrapper for asyncio.wait.'''
 
  async def iterwait(self, futures, *, flat=True, exc_as_result=True,
  timeout=None, yield_when=aio.ALL_COMPLETED):
@@ -16,9 +17,9 @@ async def iterwait(self, futures, *, flat=True, exc_as_result=True,
  timeout=timeout, return_when=yield_when)
  if flat:
  for fut in done:
- yield _get_future_result(fut, exc_as_result)
+ yield result_noraise(fut, exc_as_result)
  else:
- yield [_get_future_result(f, exc_as_result) for f in done]
+ yield [result_noraise(f, exc_as_result) for f in done]
 
  async def itermap(self, fn, iterable, *, flat=True, exc_as_result=True,
  timeout=None, yield_when=aio.ALL_COMPLETED):

asyncio_pool/utils.py

@@ -1,8 +1,10 @@
 # coding: utf8
 
+import asyncio as aio
 
-def _get_future_result(future, exc_as_result=True):
- if future.exception():
- return future.exception() if exc_as_result else None
- else:
+
+def result_noraise(future, exc_as_result=True):
+ try:
  return future.result()
+ except Exception as exc:
+ return exc if exc_as_result else None