new features: spawn, spawn_n; map, map_n, itermap; usage code instead of proper doc for now

Author: gistart

README.md

@@ -1,35 +1,124 @@
 # asyncio-pool
 
-TODO: cancelled, timeouts, callbacks, features, tests, readme
+###### Supports python 3.5+ (including PyPy 6+, which is also 3.5)
 
-Example (more in `tests/` and `examples/` dirs): # TODO
+
+## Usage
+(more in `tests/` and `examples/` dirs): # TODO
 
 ```python
-import asyncio as aio
-from asyncio_pool import AioPool
+async def worker(n):  # dummy worker
+    await aio.sleep(1 / n)
+    return n
 
 
-async def worker(n):
-    await aio.sleep(1 / n)
+async def spawn_n_usage(todo=[range(1,51), range(51,101), range(101,200)]):
+    futures = []
+    async with AioPool(size=20) as pool:
+        for tasks in todo:
+            for i in tasks:  # too many tasks
+                # Returns quickly for all tasks, does not wait for pool space.
+                # Workers are not spawned, they wait for pool space in their
+                # own background tasks.
+                fut = await pool.spawn_n(worker(i))
+                futures.append(fut)
+        # At this point not a single worker should start.
+
+        # Context manager calls `join` at exit, so this will finish when all
+        # workers return, crash or cancelled.
+
+    assert sum(itertools.chain.from_iterable(todo)) == \
+        sum(f.result() for f in futures)
 
 
-async def run_in_pool():
+async def spawn_usage(todo=range(1,4)):
+    futures = []
+    async with AioPool(size=2) as pool:
+        for i in todo:  # 1, 2, 3
+            # Returns quickly for 1 and 2, then waits for empty space for 3,
+            # spawns 3 and returns. Can save some resources I guess.
+            fut = await pool.spawn(worker(i))
+            futures.append(fut)
+        # At this point some of the workers already started.
 
-    async with AioPool(size=10) as pool:  # no more than 10 concurrent coroutines
-        results = await pool.map(worker, range(1, 100))
+        # Context manager calls `join()` at exit, so this will finish when all
+        # workers return, crash or cancelled.
 
-### OR
+    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,
+    # 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)
+
+    assert isinstance(results[0], ZeroDivisionError) \
+        and sum(results[1:]) == sum(todo)
+
+
+async def itermap_usage(todo=range(1,11)):
+    # Python 3.6+
+    result = 0
+    async with AioPool(size=10) as pool:
+        # Combines spawn_n and iterwait, which is a wrapper for asyncio.wait,
+        # which yields results of finished workers according to `timeout` and
+        # `yield_when` params passed to asyncio.wait (see it's docs for details)
+        async for res in pool.itermap(worker, todo, timeout=0.5):
+            result += res
+        # technically, you can skip join call
+
+    assert result == sum(todo)
 
-    # generator returning futures for each worker result
-    futures = await pool.itermap(worker, range(1,100))
-    # or spawning manually: list of futures for each worker result
-    futures = [await pool.spawn(worker(i)) for i in range(1,100)]
 
+async def callbacks_usage():
+    pass  # TODO
+
+
+async def details(todo=range(1,11)):
+    pool = AioPool(size=5)
+
+    # This code:
+    f1 = []
+    for i in todo:
+        f1.append(await pool.spawn_n(worker(i)))
+    # is equivalent to one call of `map_n`:
+    f2 = await pool.map_n(worker, todo)
+
+    # Afterwards you can await for any given future:
+    try:
+        assert 3 == await f1[2]  # result of spawn_n(worker(3))
+    except Exception as e:
+        # exception happened in worker (or CancelledError) will be re-raised
+        pass
+
+    # Or use `asyncio.wait` to handle results in batches (see `iterwait` also):
+    important_res = 0
+    more_important = [f1[1], f2[1], f2[2]]
+    while more_important:
+        done, more_important = await aio.wait(more_important, timeout=0.5)
+        # handle result, note it will re-raise exceptions
+        important_res += sum(f.result() for f in done)
+
+    assert important_res == 2 + 2 + 3
+
+    # But you need to join, to allow all spawned workers to finish
+    # (of course you can `asyncio.wait` all of the futures if you want to)
     await pool.join()
-    print [fut.result() for fut in batch]  # will re-raise exceptions
 
+    assert all(f.done() for f in itertools.chain(f1,f2))  # this is guaranteed
 
-### OR moar later
 ```
+
+
+## TODO:
+
+- [ ] cancelled, timeouts
+- [ ] callbacks
+- [x] tests
+- [x] usage
+- [ ] examples
+- [ ] docs
+- [ ] readme

asyncio_pool/__init__.py

@@ -1,3 +1,14 @@
 # coding: utf8
 
-from .pool import AioPool
+import sys
+from .base_pool import BaseAioPool
+
+
+if sys.version_info < (3, 6):  # this means 3.5  # TODO test 3.4?
+
+    class AioPool(BaseAioPool): pass
+
+else:
+    from .mx_asynciter import MxAsyncIterPool
+
+    class AioPool(MxAsyncIterPool, BaseAioPool): pass

asyncio_pool/base_pool.py

@@ -0,0 +1,157 @@
+# coding: utf8
+'''Pool of asyncio coroutines with familiar interface, python3.5+ friendly'''
+
+import traceback
+import collections
+import asyncio as aio
+from .utils import _get_future_result
+
+
+class BaseAioPool(object):
+
+    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.semaphore = aio.Semaphore(value=self.size, loop=self.loop)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, ext_type, exc, tb):
+        await self.join()
+
+    @property
+    def n_active(self):
+        return self.size - self.semaphore._value
+
+    @property
+    def is_empty(self):
+        return 0 == len(self._waiting) == self.n_active
+
+    @property
+    def is_full(self):
+        return self.size <= len(self._waiting) + self.n_active
+
+    async def join(self):
+        if self.is_empty:
+            return True
+
+        fut = self.loop.create_future()
+        self._joined.append(fut)
+        try:
+            return await fut
+        finally:
+            self._joined.remove(fut)
+
+    def _release_joined(self):
+        if not self.is_empty:
+            raise RuntimeError()  # TODO better message
+
+        for fut in self._joined:
+            if not fut.done():
+                fut.set_result(True)
+
+    async def _wrap(self, coro, future, cb=None, ctx=None):
+        res, exc, tb = None, None, None
+        try:
+            res = await coro
+        except Exception as _exc:
+            exc = _exc
+            tb = traceback.format_exc()  # TODO tb object instead of text
+        finally:
+            self._executed += 1
+
+        if cb:
+            err = None if exc is None else (exc, tb)
+            wrapped = self._wrap(cb(res, err, ctx), future)
+            self.loop.create_task(wrapped)
+            return
+
+        self.semaphore.release()
+        if not exc:
+            future.set_result(res)
+        else:
+            future.set_exception(exc)
+
+        if self.is_empty:
+            self._release_joined()
+
+    async def _spawn(self, future, coro, cb=None, ctx=None):
+        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)
+        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))
+        return future
+
+    async def spawn(self, coro, cb=None, ctx=None):
+        future = self.loop.create_future()
+        self._waiting.append(future)
+        return await self._spawn(future, coro, cb=cb, ctx=ctx)
+
+    async def exec(self, coro, cb=None, ctx=None):
+        return await (await self.spawn(coro, cb=cb, ctx=ctx))
+
+    async def map_n(self, fn, iterable):
+        futures = []
+        for it in iterable:
+            fut = await self.spawn_n(fn(it))
+            futures.append(fut)
+        return futures
+
+    async def map(self, fn, iterable, exc_as_result=True):
+        futures = await self.map_n(fn, iterable)
+        await self.join()
+
+        result = []
+        for fut in futures:
+            if fut.exception():
+                res = fut.exception() if exc_as_result else None
+            else:
+                res = fut.result()
+            result.append(res)
+        return result
+
+    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')
+
+    '''
+    if sys.version_info >= (3, 6):  # supports async generators
+
+        async def iterwait(self, futures, *, flat=True, exc_as_result=True,
+                timeout=None, yield_when=aio.ALL_COMPLETED):
+
+            _futures = futures[:]
+            while _futures:
+                done, _futures = await aio.wait(_futures, loop=self.loop,
+                        timeout=timeout, return_when=yield_when)
+                if flat:
+                    for fut in done:
+                        yield _get_future_result(fut, exc_as_result)
+                else:
+                    yield [_get_future_result(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):
+
+            futures = await self.map_n(fn, iterable)
+            generator = self.iterwait(futures, flat=flat, timeout=timeout,
+                    exc_as_result=exc_as_result, yield_when=yield_when)
+            async for batch in generator:
+                yield batch  # TODO is it possible to return a generator?
+    '''

asyncio_pool/mx_asynciter.py

@@ -0,0 +1,30 @@
+# coding: utf8
+'''Mixin for BaseAioPool with async iterator features python3.6+'''
+
+import asyncio as aio
+from .utils import _get_future_result
+
+
+class MxAsyncIterPool(object):
+
+    async def iterwait(self, futures, *, flat=True, exc_as_result=True,
+            timeout=None, yield_when=aio.ALL_COMPLETED):
+
+        _futures = futures[:]
+        while _futures:
+            done, _futures = await aio.wait(_futures, loop=self.loop,
+                    timeout=timeout, return_when=yield_when)
+            if flat:
+                for fut in done:
+                    yield _get_future_result(fut, exc_as_result)
+            else:
+                yield [_get_future_result(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):
+
+        futures = await self.map_n(fn, iterable)
+        generator = self.iterwait(futures, flat=flat, timeout=timeout,
+                exc_as_result=exc_as_result, yield_when=yield_when)
+        async for batch in generator:
+            yield batch  # TODO is it possible to return a generator?