add executor.map() functionality

Author: kristjanvalur

src/qasync/__init__.py

@@ -214,7 +214,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:
@@ -241,6 +259,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,7 @@
 import threading
 import time
 import weakref
-from concurrent.futures import CancelledError
+from concurrent.futures import CancelledError, TimeoutError
 
 import pytest
 
@@ -167,3 +167,32 @@ 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]
+
+
+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
+    assert duration < 0.05
+
+    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}