support of multiple joins

Author: gistart

.gitignore

@@ -7,7 +7,7 @@
 .env*
 .pypyenv*
 .pytest_*
-.git
+.mypy_cache
 __pycache__
 local_settings.py
 

README.md

@@ -1,3 +1,35 @@
 # asyncio-pool
 
-TODO
+TODO: cancelled, timeouts, callbacks, features, tests, readme
+
+Example (more in `tests/` and `examples/` dirs): # TODO
+
+```python
+import asyncio as aio
+from asyncio_pool import AioPool
+
+
+async def worker(n):
+    await aio.sleep(1 / n)
+
+
+async def run_in_pool():
+
+    async with AioPool(size=10) as pool:  # no more than 10 concurrent coroutines
+        results = await pool.map(worker, range(1, 100))
+
+### OR
+
+    pool = AioPool(size=10)
+
+    # 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)]
+
+    await pool.join()
+    print [fut.result() for fut in batch]  # will re-raise exceptions
+
+
+### OR moar later
+```

asyncio_pool/pool.py

@@ -1,45 +1,67 @@
 # coding: utf8
+'''Pool of asyncio coroutines with familiar interface'''
 
 import traceback
+import collections
 import asyncio as aio
 
 
 class AioPool(object):
 
-    def __init__(self, size=1024, loop=None):
-        self._size = size
+    def __init__(self, size=1024, *, loop=None):
+        self.loop = loop or aio.get_event_loop()
+
+        self.size = size
         self._waiting = 0
         self._executed = 0
-        self.semaphore = aio.Semaphore(value=self._size)
-        self.loop = loop or aio.get_event_loop()
-        self._all_done = self.loop.create_future()
+        self._joined = 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 == self._waiting == (self._size - self.semaphore._value)
+        return 0 == self._waiting == self.n_active
 
     @property
     def is_full(self):
-        return self._waiting + (self._size - self.semaphore._value) >= self._size
+        return self.size <= self._waiting + self.n_active
 
     async def join(self):
-        await self._all_done
+        if self.is_empty:
+            return True
 
-    async def __aenter__(self):
-        return self
+        fut = self.loop.create_future()
+        self._joined.append(fut)
+        try:
+            return await fut
+        finally:
+            self._joined.remove(fut)
 
-    async def __aexit__(self, ext_type, exc, tb):
-        await self.join()
+    def _release_joined(self):
+        if not self.is_empty:
+            raise RuntimeError()  # TODO
+
+        for fut in self._joined:
+            if not fut.done():
+                fut.set_result(True)
 
     async def _acquire(self):
-        if self._all_done.done():
-            self._all_done = self.loop.create_future()
         self._waiting += 1
         await self.semaphore.acquire()
         self._waiting -= 1
 
     async def _wrap(self, coro, future, cb=None, ctx=None):
         res, exc, tb = None, None, None
+
         try:
             res = await coro
             future.set_result(res)
@@ -50,10 +72,14 @@ async def _wrap(self, coro, future, cb=None, ctx=None):
             self.semaphore.release()
             self._executed += 1
 
-        if cb:
-            await self.spawn(cb(res, (exc, tb), ctx))
-        elif self.is_empty:
-            self._all_done.set_result('done')
+        try:
+            if cb:
+                await self.spawn(cb(res, (exc, tb), ctx))
+        except Exception as exc_cb:
+            pass  # TODO
+        finally:
+            if self.is_empty:
+                self._release_joined()
 
     async def spawn(self, coro, cb=None, ctx=None):
         await self._acquire()

reqs-dev.txt

@@ -1,3 +1,5 @@
+asyncio
 ipython
 pytest
 pytest-asyncio
+async-timeout

run_tests.sh

@@ -1,23 +1,26 @@
 #! /bin/sh
 
-py35=python3.5
 py36=python3.6
 py37=python3.7
-pypy3=/opt/pypy3/pypy3/bin/pypy3
-
+## 3.5 and pypy (which is also 3.5) are disalbed because of async generators
+# py35=python3.5
+# pypy3=/opt/pypy3/pypy3/bin/pypy3
 default_env=$py36
 
+todo=${@:-"./tests"}
 
 for py in $py35 $py36 $py37 $pypy3
 do
+    echo ""
     if [ -x "$(command -v $py)" ]; then
         pyname="$(basename $py)"
         envname=".env_$pyname"
 
         if ! [ -d $envname ]; then
             echo "$pyname: virtual env does not exist"
         else
-            $envname/bin/pytest tests
+            echo "$pyname: running for $todo"
+            $envname/bin/pytest $todo
         fi
     else
         echo "$py: not found"

tests/old.py

@@ -1,6 +1,12 @@
 # coding: utf8
 
-from .asyncio_pool import AioPool
+import os
+import sys
+curr_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.split(curr_dir)[0])
+
+import asyncio as aio
+from asyncio_pool import AioPool
 
 
 async def test(n):

tests/test_base.py

@@ -0,0 +1,67 @@
+# coding: utf8
+
+import pytest
+import asyncio as aio
+from asyncio_pool import AioPool
+from async_timeout import timeout
+
+
+@pytest.mark.asyncio
+async def test_timeout_cancel():
+    async def wrk(sem):
+        async with sem:
+            await aio.sleep(1)
+
+    sem = aio.Semaphore(value=2)
+
+    async with timeout(0.2):
+        with pytest.raises(aio.CancelledError):
+            await aio.gather(*[wrk(sem) for _ in range(3)])
+
+
+@pytest.mark.asyncio
+async def test_outer_join():
+
+    todo, to_release = range(1,15), range(10)
+    done, released = [], []
+
+    async def inner(n):
+        nonlocal done
+        await aio.sleep(1 / n)
+        done.append(n)
+
+    async def outer(n, pool):
+        nonlocal released
+        await pool.join()
+        released.append(n)
+
+    loop = aio.get_event_loop()
+    pool = AioPool(size=100)
+
+    tasks = [await pool.spawn(inner(i)) for i in todo]
+    joined = [loop.create_task(outer(j, pool)) for j in to_release]
+    await pool.join()
+
+    assert len(released) < len(to_release)
+    await aio.wait(joined)
+    assert len(todo) == len(done) and len(released) == len(to_release)
+
+
+@pytest.mark.asyncio
+async def test_internal_join():
+    async def wrk(n, pool):
+        aio.sleep(1 / n)
+        if n == 3:
+            await pool.join()
+        else:
+            await pool.spawn(wrk(n + 1, pool))
+        return n
+
+    return True
+    pool = AioPool(size=3)
+    await pool.spawn(wrk(1, pool))
+
+    async with timeout(1.5) as tm:
+        await pool.join()
+
+    assert tm.expired