fix: as_completed with return_exceptions=True

a_sync/utils/as_completed.py

@@ -76,11 +76,12 @@ def as_completed(fs, *, timeout: Optional[float] = None, return_exceptions: bool
             ...
         ```
     """
+    if isinstance(fs, Mapping):
+        return as_completed_mapping(fs, timeout=timeout, return_exceptions=return_exceptions, aiter=aiter, tqdm=tqdm, **tqdm_kwargs) 
     if return_exceptions:
-        raise NotImplementedError
+        fs = [__exc_wrap(f) for f in fs]
     return (
-        as_completed_mapping(fs, timeout=timeout, return_exceptions=return_exceptions, aiter=aiter, tqdm=tqdm, **tqdm_kwargs) if isinstance(fs, Mapping)
-        else ASyncIterator.wrap(__yield_as_completed(fs, tqdm=tqdm, **tqdm_kwargs)) if aiter 
+        ASyncIterator.wrap(__yield_as_completed(fs, timeout=timeout, tqdm=tqdm, **tqdm_kwargs)) if aiter 
         else tqdm_asyncio.as_completed(fs, timeout=timeout, **tqdm_kwargs) if tqdm 
         else asyncio.as_completed(fs, timeout=timeout)
     )
@@ -128,3 +129,9 @@ async def __yield_as_completed(futs: Iterable[Awaitable[T]], *, timeout: Optiona
 
 async def __mapping_wrap(k: KT, v: Awaitable[VT]) -> VT:
     return k, await v
+
+async def __exc_wrap(awaitable: Awaitable[T]) -> Union[T, Exception]:
+    try:
+        return await awaitable
+    except Exception as e:
+        return e

tests/fixtures.py

@@ -135,3 +135,14 @@ def _test_kwargs(fn, default: Literal['sync','async',None]):
         assert fn() == 2
     elif default == 'async':
         assert asyncio.get_event_loop().run_until_complete(fn()) == 2
+
+async def sample_task(n):
+    await asyncio.sleep(0.01)
+    return n
+
+async def timeout_task(n):
+    await asyncio.sleep(0.1)
+    return n
+
+async def sample_exc(n):
+    raise ValueError("Sample error")

tests/test_as_completed.py

@@ -0,0 +1,79 @@
+
+import asyncio
+
+import a_sync
+import pytest
+
+from tests.fixtures import sample_exc, sample_task, timeout_task
+
+
+@pytest.mark.asyncio_cooperative
+async def test_as_completed_with_awaitables():
+    tasks = [sample_task(i) for i in range(5)]
+    results = [await result for result in a_sync.as_completed(tasks, aiter=False)]
+    assert sorted(results) == list(range(5)), "Results should be in ascending order from 0 to 4"
+
+@pytest.mark.asyncio_cooperative
+async def test_as_completed_with_awaitables_aiter():
+    tasks = [sample_task(i) for i in range(5)]
+    results = []
+    async for result in a_sync.as_completed(tasks, aiter=True):
+        results.append(result)
+    assert sorted(results) == list(range(5)), "Results should be in ascending order from 0 to 4"
+
+@pytest.mark.asyncio_cooperative
+async def test_as_completed_with_mapping():
+    tasks = {'task1': sample_task(1), 'task2': sample_task(2)}
+    results = {}
+    for result in a_sync.as_completed(tasks, aiter=False):
+        key, value = await result
+        results[key] = value
+    assert results == {'task1': 1, 'task2': 2}, "Results should match the input mapping"
+
+@pytest.mark.asyncio_cooperative
+async def test_as_completed_with_mapping_aiter():
+    tasks = {'task1': sample_task(1), 'task2': sample_task(2)}
+    results = {}
+    async for key, result in a_sync.as_completed(tasks, aiter=True):
+        results[key] = result
+    assert results == {'task1': 1, 'task2': 2}, "Results should match the input mapping"
+
+@pytest.mark.asyncio_cooperative
+async def test_as_completed_with_timeout():
+    tasks = [timeout_task(i) for i in range(2)]
+    with pytest.raises(asyncio.TimeoutError):
+        [await result for result in a_sync.as_completed(tasks, aiter=False, timeout=0.05)]
+
+@pytest.mark.asyncio_cooperative
+async def test_as_completed_with_timeout_aiter():
+    tasks = [timeout_task(i) for i in range(2)]
+    with pytest.raises(asyncio.TimeoutError):
+        [result async for result in a_sync.as_completed(tasks, aiter=True, timeout=0.05)]
+
+@pytest.mark.asyncio_cooperative
+async def test_as_completed_return_exceptions():
+    tasks = {i: sample_exc(i) for i in range(1)}
+    results = [await result for result in a_sync.as_completed(tasks, aiter=False, return_exceptions=True)]
+    assert isinstance(results[0], ValueError), "The result should be an exception"
+
+@pytest.mark.asyncio_cooperative
+async def test_as_completed_return_exceptions_aiter():
+    tasks = [sample_exc(i) for i in range(1)]
+    results = []
+    async for result in a_sync.as_completed(tasks, aiter=True, return_exceptions=True):
+        results.append(result)
+    assert isinstance(results[0], ValueError), "The result should be an exception"
+
+@pytest.mark.asyncio_cooperative
+async def test_as_completed_with_tqdm_disabled():
+    tasks = [sample_task(i) for i in range(5)]
+    results = [await result for result in a_sync.as_completed(tasks, aiter=False, tqdm=False)]
+    assert sorted(results) == list(range(5)), "Results should be in ascending order from 0 to 4"
+
+@pytest.mark.asyncio_cooperative
+async def test_as_completed_with_tqdm_disabled_aiter():
+    tasks = [sample_task(i) for i in range(5)]
+    results = []
+    async for result in a_sync.as_completed(tasks, aiter=True, tqdm=False):
+        results.append(result)
+    assert sorted(results) == list(range(5)), "Results should be in ascending order from 0 to 4"

tests/test_gather.py

@@ -0,0 +1,32 @@
+import asyncio
+
+import pytest
+
+from a_sync import gather
+from tests.fixtures import sample_exc
+
+
+async def sample_task(number):
+    await asyncio.sleep(0.1)
+    return number * 2
+
+@pytest.mark.asyncio_cooperative
+async def test_gather_with_awaitables():
+    results = await gather(sample_task(1), sample_task(2), sample_task(3))
+    assert results == [2, 4, 6]
+
+@pytest.mark.asyncio_cooperative
+async def test_gather_with_mapping():
+    tasks = {'a': sample_task(1), 'b': sample_task(2), 'c': sample_task(3)}
+    results = await gather(tasks)
+    assert results == {'a': 2, 'b': 4, 'c': 6}
+
+@pytest.mark.asyncio_cooperative
+async def test_gather_with_exceptions():
+    with pytest.raises(ValueError):
+        await gather(sample_exc(None))
+
+@pytest.mark.asyncio_cooperative
+async def test_gather_with_return_exceptions():
+    results = await gather(sample_exc(None), return_exceptions=True)
+    assert isinstance(results[0], ValueError)