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 asyncio_pool/mx_asyncgen.pyasyncio_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