gh-108518: Make concurrent.futures.Executor.map() consistent with built-in map()

Lib/concurrent/futures/_base.py

@@ -314,7 +314,9 @@ def wait(fs, timeout=None, return_when=ALL_COMPLETED):
 def _result_or_cancel(fut, timeout=None):
     try:
         try:
-            return fut.result(timeout)
+            return _FutureResult.from_value(fut.result(timeout))
+        except Exception as e:
+            return _FutureResult.from_exception(e)
         finally:
             fut.cancel()
     finally:
@@ -566,6 +568,46 @@ def set_exception(self, exception):
 
     __class_getitem__ = classmethod(types.GenericAlias)
 
+
+class _FutureResult(object):
+    """
+    This is used to record the exception instead of throwing them.
+
+    _FutureResult must contain either the value of future or an exception
+    that was thrown during the computation of future. Use is_exception
+    property to determine which one it is.
+    """
+
+    def __init__(self, exception, value):
+        self._exception = exception
+        self._value = value
+
+    @classmethod
+    def from_exception(cls, exc):
+        return cls(exc, None)
+
+    @classmethod
+    def from_value(cls, value):
+        return cls(None, value)
+
+    @property
+    def exception(self):
+        if not self.is_exception:
+            raise RuntimeError("No exception thrown.")
+        return self._exception
+
+    @property
+    def value(self):
+        if self.is_exception:
+            raise RuntimeError(
+                "Cannot get result value because an exception was thrown.")
+        return self._value
+
+    @property
+    def is_exception(self):
+        return self._exception is not None
+
+
 class Executor(object):
     """This is an abstract base class for concrete asynchronous executors."""
 
@@ -602,6 +644,11 @@ def map(self, fn, *iterables, timeout=None, chunksize=1):
                 before the given timeout.
             Exception: If fn(*args) raises for any values.
         """
+        return _MapResultIterator.from_generator(
+            self._map(fn, *iterables, timeout=timeout)
+        )
+
+    def _map(self, fn, *iterables, timeout=None):
         if timeout is not None:
             end_time = timeout + time.monotonic()
 
@@ -648,6 +695,25 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         return False
 
 
+class _MapResultIterator(object):
+    """The iterator returned by map()."""
+    def __init__(self, gen):
+        self.gen = gen
+
+    @classmethod
+    def from_generator(cls, gen):
+        return cls(gen)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        result = next(self.gen)
+        if result.is_exception:
+            raise result.exception
+        return result.value
+
+
 class BrokenExecutor(RuntimeError):
     """
     Raised when a executor has become non-functional after a severe failure.

Lib/concurrent/futures/process.py

@@ -204,7 +204,14 @@ def _process_chunk(fn, chunk):
     This function is run in a separate process.
 
     """
-    return [fn(*args) for args in chunk]
+    results = []
+    for args in chunk:
+        try:
+            result = _base._FutureResult.from_value(fn(*args))
+        except Exception as e:
+            result = _base._FutureResult.from_exception(e)
+        results.append(result)
+    return results
 
 
 def _sendback_result(result_queue, work_id, result=None, exception=None,
@@ -617,6 +624,9 @@ def _chain_from_iterable_of_lists(iterable):
     careful not to keep references to yielded objects.
     """
     for element in iterable:
+        if element.is_exception:
+            raise element.exception
+        element = element.value
         element.reverse()
         while element:
             yield element.pop()
@@ -830,10 +840,11 @@ def map(self, fn, *iterables, timeout=None, chunksize=1):
         if chunksize < 1:
             raise ValueError("chunksize must be >= 1.")
 
-        results = super().map(partial(_process_chunk, fn),
+        results = super()._map(partial(_process_chunk, fn),
                               _get_chunks(*iterables, chunksize=chunksize),
                               timeout=timeout)
-        return _chain_from_iterable_of_lists(results)
+        return _base._MapResultIterator.from_generator(
+            _chain_from_iterable_of_lists(results))
 
     def shutdown(self, wait=True, *, cancel_futures=False):
         with self._shutdown_lock:

Lib/test/test_concurrent_futures/executor.py

@@ -52,6 +52,9 @@ def test_map_exception(self):
         self.assertEqual(i.__next__(), (0, 1))
         self.assertEqual(i.__next__(), (0, 1))
         self.assertRaises(ZeroDivisionError, i.__next__)
+        self.assertEqual(i.__next__(), (0, 1))
+        self.assertRaises(StopIteration, i.__next__)
+        self.assertRaises(StopIteration, i.__next__)
 
     @support.requires_resource('walltime')
     def test_map_timeout(self):

Lib/test/test_concurrent_futures/test_thread_pool.py

@@ -79,16 +79,13 @@ def log_n_wait(ident):
             # submit work to saturate the pool
             fut = pool.submit(log_n_wait, ident="first")
             try:
-                with contextlib.closing(
-                    pool.map(log_n_wait, ["second", "third"], timeout=0)
-                ) as gen:
-                    with self.assertRaises(TimeoutError):
-                        next(gen)
+                iterator = pool.map(log_n_wait, ["second"], timeout=0)
+                with self.assertRaises(TimeoutError):
+                    next(iterator)
             finally:
                 stop_event.set()
             fut.result()
         # ident='second' is cancelled as a result of raising a TimeoutError
-        # ident='third' is cancelled because it remained in the collection of futures
         self.assertListEqual(log, ["ident='first' started", "ident='first' stopped"])