more flexible results extraction

Author: gistart

README.md

@@ -8,39 +8,39 @@ AioPool makes sure _no more_ and _no less_ (if possible) than `size` spawned cor
 
 Read [code doctrings](../master/asyncio_pool/base_pool.py) for details.
 
-> `AioPool(size=4, *, loop=None)`
+#### AioPool(size=4, *, loop=None)
 
 Creates pool of `size` concurrent tasks. Supports async context manager interface.
 
-> `spawn(coro, cb=None, ctx=None)`
+#### spawn(coro, cb=None, ctx=None)
 
 Waits for pool space, then creates task for `coro` coroutine, returning future for it's result. Can spawn coroutine, created by `cb` with result of `coro` as first argument. `ctx` context is passed to callback as third positinal argument.
 
-> `exec(coro, cb=None, ctx=None)`
+#### exec(coro, cb=None, ctx=None)
 
 Waits for pool space, then creates task for `coro`, then waits for it to finish, then returns result of `coro` if no callback is provided, otherwise creates task for callback, waits for it and returns result of callback.
 
-> `spawn_n(coro, cb=None, ctx=None)`
+#### spawn_n(coro, cb=None, ctx=None)
 
 Creates waiting task for `coro`, returns future without waiting for pool space. Task is executed "in pool" when pool space is available.
 
-> `join()`
+#### join()
 
 Waits for all spawned (active and waiting) tasks to finish. Joining pool from coroutine, spawned by the same pool leads to *deadlock*.
 
-> `cancel(*futures)`
+#### cancel(*futures)
 
 Cancels spawned tasks (active and waiting), finding them by provided `futures`. If no futures provided -- cancels all spawned tasks.
 
-> `map(fn, iterable, cb=None, ctx=None, *, exc_as_result=True)`
+#### map(fn, iterable, cb=None, ctx=None, *, get_result=getres.flat)
 
 Spawns coroutines created by `fn` function for each item in `iterable` with `spawn`, waits for all of them to finish (including callbacks), returns results maintaining order of `iterable`.
 
-> `map_n(fn, iterable, cb=None, ctx=None, *, exc_as_result=True)`
+#### map_n(fn, iterable, cb=None, ctx=None, *, get_result=getres.flat)
 
 Spawns coroutines created by `fn` function for each item in `iterable` with `spawn_n`, returns futures for task results maintaining order of `iterable`.
 
-> `itermap(fn, iterable, cb=None, ctx=None, *, flat=True, exc_as_result=True, timeout=None, yield_when=asyncio.ALL_COMPLETED)`
+#### itermap(fn, iterable, cb=None, ctx=None, *, flat=True, get_result=getres.flat, timeout=None, yield_when=asyncio.ALL_COMPLETED)
 
 Spawns tasks with `map_n(fn, iterable, cb, ctx)`, then waits for results with `asyncio.wait` function, yielding ready results one by one if `flat` == True, otherwise yielding list of ready results.
 
@@ -114,7 +114,6 @@ async def map_usage(todo=range(100)):
 
 
 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,
@@ -153,9 +152,10 @@ async def callbacks_usage():
 
     results = []
     for fut in futures:
-        # there's a helper with this code:
-        #   from asyncio_pool import result_noraise
-        #   results.append(result_noraise(fut))
+        # there are helpers for result extraction. `flat` one will do
+        # exactly what's written below
+        #   from asyncio_pool import getres
+        #   results.append(getres.flat(fut))
         try:
             results.append(fut.result())
         except Exception as e:

asyncio_pool/__init__.py

@@ -2,7 +2,7 @@
 
 import sys
 
-from .utils import result_noraise
+from .results import getres
 from .base_pool import BaseAioPool
 
 

asyncio_pool/base_pool.py

@@ -3,7 +3,7 @@
 
 import traceback
 import asyncio as aio
-from .utils import result_noraise
+from .results import getres
 
 
 class BaseAioPool(object):
@@ -212,13 +212,20 @@ async def map_n(self, fn, iterable, cb=None, ctx=None):
             futures.append(fut)
         return futures
 
-    async def map(self, fn, iterable, cb=None, ctx=None, *, exc_as_result=True):
+    async def map(self, fn, iterable, cb=None, ctx=None, *,
+            get_result=getres.flat):
         '''Spawns coroutines, created with `fn` function for each item in
         `iterable`, waits for all of them to finish, crash or be cancelled,
         returning resuls.
 
-        If coroutine or callback crashes or is cancelled, with `exc_as_result`
-        == True exceptions object will be returned, with == False -- just None.
+        `get_result` is function, that accepts future as only positional
+        argument, whose goal is to extract result from future. You can pass
+        your own, or use inluded `getres` object, that has 3 extractors:
+        `getres.dont` will return future untouched, `getres.flat` will return
+        exception object if coroutine crashed or was cancelled, otherwise will
+        return result of a coroutine (or of the callback), `getres.pair` will
+        return tuple of (`result', 'exception` object) with None in place of
+        missing item.
 
         Read more about callbacks in `spawn` docstring.
         '''
@@ -228,10 +235,10 @@ async def map(self, fn, iterable, cb=None, ctx=None, *, exc_as_result=True):
             futures.append(fut)
 
         await aio.wait(futures)
-        return [result_noraise(fut, exc_as_result) for fut in futures]
+        return [get_result(fut) for fut in futures]
 
     async def itermap(self, fn, iterable, cb=None, ctx=None, *, flat=True,
-            exc_as_result=True, timeout=None, yield_when=aio.ALL_COMPLETED):
+            get_result=getres.flat, timeout=None, yield_when=aio.ALL_COMPLETED):
         '''Spawns coroutines created with `fn` for each item in `iterable`, then
         waits for results with `iterwait` (implementation specific). See docs
         for `map_n` and `iterwait` (in mixins for py3.5 and py3.6+).
@@ -256,7 +263,7 @@ def _cancel(self, *futures):
         cancelled = sum([1 for fut in tasks if fut.cancel()])
         return cancelled, _futures
 
-    async def cancel(self, *futures, exc_as_result=True):
+    async def cancel(self, *futures, get_result=getres.flat):
         '''Cancels spawned or waiting tasks, found by their `futures`. If no
         `futures` are passed -- cancels all spwaned and waiting tasks.
 
@@ -269,5 +276,5 @@ 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]
+        results = [get_result(fut) for fut in _futures]
         return cancelled, results

asyncio_pool/mx_asyncgen.py

@@ -2,10 +2,10 @@
 '''Mixin for BaseAioPool with async generator features, python3.6+'''
 
 import asyncio as aio
-from .utils import result_noraise
+from .results import getres
 
 
-async def iterwait(futures, *, flat=True, exc_as_result=True,
+async def iterwait(futures, *, flat=True, get_result=getres.flat,
         timeout=None, yield_when=aio.ALL_COMPLETED, loop=None):
     '''Wraps `asyncio.wait` into asynchronous generator, accessible with
     `async for` syntax. May be useful in conjunction with `spawn_n`.
@@ -23,22 +23,22 @@ async def iterwait(futures, *, flat=True, exc_as_result=True,
                                         return_when=yield_when, loop=loop)
         if flat:
             for fut in done:
-                yield result_noraise(fut, exc_as_result)
+                yield get_result(fut)
         else:
-            yield [result_noraise(f, exc_as_result) for f in done]
+            yield [get_result(fut) for fut in done]
 
 
 class MxAsyncGenPool(object):
     # Asynchronous generator wrapper for asyncio.wait.
 
     async def itermap(self, fn, iterable, cb=None, ctx=None, *, flat=True,
-            exc_as_result=True, timeout=None, yield_when=aio.ALL_COMPLETED):
+            get_result=getres.flat, timeout=None,
+            yield_when=aio.ALL_COMPLETED):
         '''Spawns coroutines created with `fn` for each item in `iterable`, then
         waits for results with `iterwait`. See docs for `map_n` and `iterwait`.
         '''
         futures = await self.map_n(fn, iterable, cb, ctx)
         generator = iterwait(futures, flat=flat, timeout=timeout,
-                exc_as_result=exc_as_result, yield_when=yield_when,
-                loop=self.loop)
+                get_result=get_result, yield_when=yield_when, loop=self.loop)
         async for batch in generator:
             yield batch  # TODO is it possible to return a generator?

asyncio_pool/mx_asynciter.py

@@ -4,18 +4,18 @@
 import asyncio as aio
 from collections import deque
 from functools import partial
-from .utils import result_noraise
+from .results import getres
 
 
 class iterwait:
 
-    def __init__(self, futures, *, flat=True, exc_as_result=True,
+    def __init__(self, futures, *, flat=True, get_result=getres.flat,
             timeout=None, yield_when=aio.ALL_COMPLETED, loop=None):
 
         self.results = deque()
         self.flat = flat
         self._futures = futures
-        self._extract = partial(result_noraise, exc_as_result=exc_as_result)
+        self._getres = get_result
         self._wait = partial(aio.wait, timeout=timeout, loop=loop,
                              return_when=yield_when)
 
@@ -33,7 +33,7 @@ async def _wait_next(self):
         while True:
             done, self._futures = await self._wait(self._futures)
             if done:
-                batch = [self._extract(fut) for fut in done]
+                batch = [self._getres(fut) for fut in done]
                 if self.flat:
                     self.results.extend(batch)
                 else:
@@ -44,13 +44,14 @@ async def _wait_next(self):
 class MxAsyncIterPool(object):
 
     def itermap(self, fn, iterable, cb=None, ctx=None, *, flat=True,
-            exc_as_result=True, timeout=None, yield_when=aio.ALL_COMPLETED):
+            get_result=getres.flat, timeout=None,
+            yield_when=aio.ALL_COMPLETED):
         '''Spawns coroutines created with `fn` for each item in `iterable`, then
         waits for results with `iterwait`. See docs for `map_n` and `iterwait`.
         '''
         mk_map = partial(self.map_n, fn, iterable, cb=cb, ctx=ctx)
         mk_waiter = partial(iterwait, flat=flat, loop=self.loop,
-                            exc_as_result=exc_as_result, timeout=timeout,
+                            get_result=get_result, timeout=timeout,
                             yield_when=yield_when)
 
         class _itermap:

asyncio_pool/results.py

@@ -0,0 +1,24 @@
+# coding: utf8
+
+from functools import partial
+
+
+def result_noraise(future, flat=True):
+    '''Extracts result from future, never raising an exception.
+
+    If `flat` is True -- returns result or exception instance (including
+    CancelledError), if `flat` is False -- returns tuple of (`result`,
+    `exception` object).
+
+    If traceback is needed -- just re-raise returned exception.'''
+    try:
+        res = future.result()
+        return res if flat else (res, None)
+    except Exception as exc:
+        return exc if flat else (None, exc)
+
+
+class getres:
+    dont = lambda fut: fut
+    flat = partial(result_noraise, flat=True)
+    pair = partial(result_noraise, flat=False)

asyncio_pool/utils.py

@@ -8,39 +8,39 @@ AioPool makes sure _no more_ and _no less_ (if possible) than `size` spawned cor
 
 Read [code doctrings](../master/asyncio_pool/base_pool.py) for details.
 
-> `AioPool(size=4, *, loop=None)`
+#### AioPool(size=4, *, loop=None)
 
 Creates pool of `size` concurrent tasks. Supports async context manager interface.
 
-> `spawn(coro, cb=None, ctx=None)`
+#### spawn(coro, cb=None, ctx=None)
 
 Waits for pool space, then creates task for `coro` coroutine, returning future for it's result. Can spawn coroutine, created by `cb` with result of `coro` as first argument. `ctx` context is passed to callback as third positinal argument.
 
-> `exec(coro, cb=None, ctx=None)`
+#### exec(coro, cb=None, ctx=None)
 
 Waits for pool space, then creates task for `coro`, then waits for it to finish, then returns result of `coro` if no callback is provided, otherwise creates task for callback, waits for it and returns result of callback.
 
-> `spawn_n(coro, cb=None, ctx=None)`
+#### spawn_n(coro, cb=None, ctx=None)
 
 Creates waiting task for `coro`, returns future without waiting for pool space. Task is executed "in pool" when pool space is available.
 
-> `join()`
+#### join()
 
 Waits for all spawned (active and waiting) tasks to finish. Joining pool from coroutine, spawned by the same pool leads to *deadlock*.
 
-> `cancel(*futures)`
+#### cancel(*futures)
 
 Cancels spawned tasks (active and waiting), finding them by provided `futures`. If no futures provided -- cancels all spawned tasks.
 
-> `map(fn, iterable, cb=None, ctx=None, *, exc_as_result=True)`
+#### map(fn, iterable, cb=None, ctx=None, *, get_result=getres.flat)
 
 Spawns coroutines created by `fn` function for each item in `iterable` with `spawn`, waits for all of them to finish (including callbacks), returns results maintaining order of `iterable`.
 
-> `map_n(fn, iterable, cb=None, ctx=None, *, exc_as_result=True)`
+#### map_n(fn, iterable, cb=None, ctx=None, *, get_result=getres.flat)
 
 Spawns coroutines created by `fn` function for each item in `iterable` with `spawn_n`, returns futures for task results maintaining order of `iterable`.
 
-> `itermap(fn, iterable, cb=None, ctx=None, *, flat=True, exc_as_result=True, timeout=None, yield_when=asyncio.ALL_COMPLETED)`
+#### itermap(fn, iterable, cb=None, ctx=None, *, flat=True, get_result=getres.flat, timeout=None, yield_when=asyncio.ALL_COMPLETED)
 
 Spawns tasks with `map_n(fn, iterable, cb, ctx)`, then waits for results with `asyncio.wait` function, yielding ready results one by one if `flat` == True, otherwise yielding list of ready results.
 

docs/_readme_template.md

@@ -104,9 +104,10 @@ async def cb(res, err, ctx):  # callback
 
     results = []
     for fut in futures:
-        # there's a helper with this code:
-        #   from asyncio_pool import result_noraise
-        #   results.append(result_noraise(fut))
+        # there are helpers for result extraction. `flat` one will do
+        # exactly what's written below
+        #   from asyncio_pool import getres
+        #   results.append(getres.flat(fut))
         try:
             results.append(fut.result())
         except Exception as e:

examples/_usage.py

@@ -8,7 +8,7 @@
 import time
 import argparse
 import asyncio as aio
-from asyncio_pool import AioPool, result_noraise
+from asyncio_pool import AioPool, getres
 
 
 async def loadtest_spawn(tasks, pool_size, duration):
@@ -18,7 +18,7 @@ async def loadtest_spawn(tasks, pool_size, duration):
             fut = await pool.spawn(aio.sleep(duration))
             futures.append(fut)
 
-    return [result_noraise(f) for f in futures]
+    return [getres.flat(fut) for fut in futures]
 
 
 async def loadtest_spawn_n(tasks, pool_size, duration):
@@ -28,7 +28,7 @@ async def loadtest_spawn_n(tasks, pool_size, duration):
             fut = await pool.spawn_n(aio.sleep(duration))
             futures.append(fut)
 
-    return [result_noraise(f) for f in futures]
+    return [getres.flat(f) for f in futures]
 
 
 async def loadtest_map(tasks, pool_size, duration):

tests/loadtest.py

@@ -3,7 +3,7 @@
 import sys
 import pytest
 import asyncio as aio
-from asyncio_pool import AioPool, result_noraise
+from asyncio_pool import AioPool, getres
 
 
 pytestmark = pytest.mark.filterwarnings('ignore:coroutine')
@@ -34,7 +34,7 @@ async def test_spawn_n():
             fut = await pool.spawn_n(wrk(i), cb, ctx)
             futures.append(fut)
 
-    results = [result_noraise(f) for f in futures]
+    results = [getres.flat(f) for f in futures]
     assert all(isinstance(e, ZeroDivisionError) for e in results[:2])
     assert sum(results[2:]) == 2 * (sum(todo) - 0 - 1)
 
@@ -54,7 +54,7 @@ async def test_map_n():
     async with AioPool(size=3) as pool:
         futures = await pool.map_n(wrk, todo, cb)
 
-    results = [result_noraise(f) for f in futures]
+    results = [getres.flat(f) for f in futures]
     assert 2 * sum(todo) == sum(results)
 
 
@@ -134,7 +134,7 @@ async def wrk(n):
         await aio.sleep(0)
         return n
 
-    _r = result_noraise
+    _r = getres.flat
     _ctx = (123, 456, [1,2], {3,4})
     inst = CbCls()