add executor.map() functionality

Author: kristjanvalur

src/qasync/__init__.py

@@ -195,7 +195,25 @@ def submit(self, callback, *args, **kwargs):
             return future
 
     def map(self, func, *iterables, timeout=None):
-        raise NotImplementedError("use as_completed on the event loop")
+        deadline = time.monotonic() + timeout if timeout is not None else None
+        futures = [self.submit(func, *args) for args in zip(*iterables)]
+
+        # must have generator as a closure so that the submit occurs before first iteration
+        def generator():
+            try:
+                futures.reverse()
+                while futures:
+                    if deadline is not None:
+                        yield _result_or_cancel(
+                            futures.pop(), timeout=deadline - time.monotonic()
+                        )
+                    else:
+                        yield _result_or_cancel(futures.pop())
+            finally:
+                for future in futures:
+                    future.cancel()
+
+        return generator()
 
     def shutdown(self, wait=True, *, cancel_futures=False):
         with self.__shutdown_lock:
@@ -222,6 +240,16 @@ def __exit__(self, *args):
         self.shutdown()
 
 
+def _result_or_cancel(fut, timeout=None):
+    try:
+        try:
+            return fut.result(timeout)
+        finally:
+            fut.cancel()
+    finally:
+        del fut  # break reference cycle in exceptions
+
+
 def _format_handle(handle: asyncio.Handle):
     cb = getattr(handle, "_callback", None)
     if isinstance(getattr(cb, "__self__", None), asyncio.tasks.Task):

tests/test_qthreadexec.py

@@ -6,7 +6,8 @@
 import threading
 import time
 import weakref
-from concurrent.futures import CancelledError
+from concurrent.futures import CancelledError, TimeoutError
+from itertools import islice
 
 import pytest
 
@@ -145,3 +146,103 @@ def task():
         assert cancels > 0
     else:
         assert cancels == 0
+
+
+def test_map(executor):
+    """Basic test of executor map functionality"""
+    results = list(executor.map(lambda x: x + 1, range(10)))
+    assert results == [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+
+    results = list(executor.map(lambda x, y: x + y, range(10), range(9)))
+    assert results == [0, 2, 4, 6, 8, 10, 12, 14, 16]
+
+
+def test_map_timeout(executor):
+    """Test that map with timeout raises TimeoutError and cancels futures"""
+    results = []
+
+    def func(x):
+        nonlocal results
+        time.sleep(0.05)
+        results.append(x)
+        return x
+
+    start = time.monotonic()
+    with pytest.raises(TimeoutError):
+        list(executor.map(func, range(10), timeout=0.01))
+    duration = time.monotonic() - start
+    # this test is flaky on some platforms, so we give it a wide bearth.
+    assert duration < 0.1
+
+    executor.shutdown(wait=True)
+    # only about half of the tasks should have completed
+    # because the max number of workers is 5 and the rest of
+    # the tasks were not started at the time of the cancel.
+    assert set(results) != {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
+
+
+def test_map_error(executor):
+    """Test that map with an exception will raise, and remaining tasks are cancelled"""
+    results = []
+
+    def func(x):
+        nonlocal results
+        time.sleep(0.05)
+        if len(results) == 5:
+            raise ValueError("Test error")
+        results.append(x)
+        return x
+
+    with pytest.raises(ValueError):
+        list(executor.map(func, range(15)))
+
+    executor.shutdown(wait=True, cancel_futures=False)
+    assert len(results) <= 10, "Final 5 at least should have been cancelled"
+
+
+@pytest.mark.parametrize("cancel", [True, False])
+def test_map_shutdown(executor, cancel):
+    results = []
+
+    def func(x):
+        nonlocal results
+        time.sleep(0.05)
+        results.append(x)
+        return x
+
+    # Get the first few results.
+    # Keep the iterator alive so that it isn't closed when its reference is dropped.
+    m = executor.map(func, range(15))
+    values = list(islice(m, 5))
+    assert values == [0, 1, 2, 3, 4]
+
+    executor.shutdown(wait=True, cancel_futures=cancel)
+    if cancel:
+        assert len(results) < 15, "Some tasks should have been cancelled"
+    else:
+        assert len(results) == 15, "All tasks should have been completed"
+    m.close()
+
+
+def test_map_start(executor):
+    """Test that map starts tasks immediately, before iterating"""
+    e = threading.Event()
+    m = executor.map(lambda x: (e.set(), x), range(1))
+    e.wait(timeout=0.1)
+    assert list(m) == [(None, 0)]
+
+
+def test_map_close(executor):
+    """Test that closing a running map cancels all remaining tasks."""
+    results = []
+    def func(x):
+        nonlocal results
+        time.sleep(0.05)
+        results.append(x)
+        return x
+    m = executor.map(func, range(10))
+    # must start the generator so that close() has any effect
+    assert next(m) == 0
+    m.close()
+    executor.shutdown(wait=True, cancel_futures=False)
+    assert len(results) < 10, "Some tasks should have been cancelled"