Fix tests with eager task factory.

python · kumaraditya303 · Sep 26, 2024 · Sep 23, 2024 · Sep 23, 2024 · Sep 23, 2024
commit e2cf78a194d8ca8046ce52e223ddd6584c062689
@@ -72,37 +72,56 @@ def future_callback(index, future, task_group):
 
         try:
             error = future.exception()
+            exceptions[index] = error
         except exceptions_mod.CancelledError as cancelled_error:
+            # If another task finishes first and cancels this task, it
+            # is propagated here.
             exceptions[index] = cancelled_error
+            return
         else:
-            exceptions[index] = error
-            task_group._errors.remove(error)
+            if error is not None:
+                return
 
         nonlocal winner_result, winner_index
-        if (winner_result is None) and (not task_group._aborting):
-            # If this is in an eager task factory, it's possible
-            # for multiple tasks to get here. In that case, we want
-            # only the first one to win and the rest to no-op before
-            # cancellation.
+        # If this is in an eager task factory, it's possible
+        # for multiple tasks to get here. In that case, we want
+        # only the first one to win and the rest to no-op before
+        # cancellation.
+        if winner_result is None and not task_group._aborting:
             winner_result = future.result()
             winner_index = index
+
+            # Cancel all other tasks, we win!
             task_group._abort()
 
     async with taskgroups.TaskGroup() as task_group:
         for index, coro in enumerate(coro_fns):
             if task_group._aborting:
                 break
 
+            exceptions.append(None)
+            task = loop.create_task(coro())
+
+            # We don't want the task group to propagate the error. Instead,
+            # we want to put it in our special exceptions list, so we manually
+            # create the task.
+            task.add_done_callback(task_group._on_task_done_without_propagation)
+            task_group._tasks.add(task)
+
+            # We need this extra wrapper here to stop the closure from having
+            # an incorrect index.
             def wrapper(idx):
                 return lambda future: future_callback(idx, future, task_group)
 
-            exceptions.append(None)
-            task = task_group.create_task(coro())
             task.add_done_callback(wrapper(index))
 
             if delay is not None:
-                await tasks.sleep(delay)
+                await tasks.sleep(delay or 0)
             else:
-                await task
+                # We don't care about exceptions here, the callback will
+                # deal with it.
+                with contextlib.suppress(BaseException):
+                    # If there's no delay, we just wait until completion.
+                    await task
 
     return winner_result, winner_index, exceptions
diff --git a/Lib/asyncio/taskgroups.py b/Lib/asyncio/taskgroups.py
@@ -199,7 +199,8 @@ def _abort(self):
             if not t.done():
                 t.cancel()
 
-    def _on_task_done(self, task):
+    def _on_task_done_without_propagation(self, task):
+        # For staggered_race()
         self._tasks.discard(task)
 
         if self._on_completed_fut is not None and not self._tasks:
@@ -209,7 +210,10 @@ def _on_task_done(self, task):
         if task.cancelled():
             return
 
-        exc = task.exception()
+        return task.exception()
+
+    def _on_task_done(self, task):
+        exc = self._on_task_done_without_propagation(task)
         if exc is None:
             return
 

@@ -243,16 +243,17 @@ async def fail():
         async def run():
             winner, index, excs = await asyncio.staggered.staggered_race(
                 [
+                    lambda: fail(),
                     lambda: asyncio.sleep(1),
                     lambda: asyncio.sleep(0),
-                    lambda: fail()
                 ],
                 delay=None
             )
             self.assertIsNone(winner)
             self.assertEqual(index, 1)
             self.assertIsNone(excs[index])
-            self.assertIsInstance(excs[2], ValueError)
+            self.assertIsInstance(excs[0], ValueError)
+            self.assertEqual(len(excs), 2)
 
         self.run_coro(run())