Allow creating new tasks while __aexit__ is waiting

We stop when there are no unfinished tasks, and then no new tasks can be created. (TODO: more thorough testing of edge cases?)
python · gvanrossum · Feb 15, 2022 · Feb 10, 2022 · Feb 10, 2022 · Feb 10, 2022
commit 0e1355d55509d083ec911995ef2d7837ff888017
@@ -170,8 +170,8 @@ async def __aexit__(self, et, exc, tb):
     def create_task(self, coro):
         if not self._entered:
             raise RuntimeError(f"TaskGroup {self!r} has not been entered")
-        if self._exiting:
-            raise RuntimeError(f"TaskGroup {self!r} is awaiting in exit")
+        if self._exiting and self._unfinished_tasks == 0:
+            raise RuntimeError(f"TaskGroup {self!r} is finished")
         task = self._loop.create_task(coro)
         task.add_done_callback(self._on_task_done)
         self._unfinished_tasks += 1
@@ -217,7 +217,7 @@ def _on_task_done(self, task):
         self._unfinished_tasks -= 1
         assert self._unfinished_tasks >= 0
 
-        if self._exiting and not self._unfinished_tasks:
+        if self._on_completed_fut is not None and not self._unfinished_tasks:
             if not self._on_completed_fut.done():
                 self._on_completed_fut.set_result(True)
 

@@ -670,3 +670,45 @@ async def do_job(delay):
                     self.assertLess(len(g._tasks), 5)
             await asyncio.sleep(1.35)
             self.assertEqual(len(g._tasks), 0)
+
+    async def test_taskgroup_24(self):
+
+        async def root(g):
+            await asyncio.sleep(0.1)
+            g.create_task(coro1(0.1))
+            g.create_task(coro1(0.2))
+
+        async def coro1(delay):
+            await asyncio.sleep(delay)
+
+        async def runner():
+            async with taskgroups.TaskGroup() as g:
+                g.create_task(root(g))
+
+        await runner()
+
+    async def test_taskgroup_25(self):
+        nhydras = 0
+
+        async def hydra(g):
+            nonlocal nhydras
+            nhydras += 1
+            await asyncio.sleep(0.01)
+            g.create_task(hydra(g))
+            g.create_task(hydra(g))
+
+        async def hercules():
+            while nhydras < 10:
+                await asyncio.sleep(0.015)
+            1 / 0
+
+        async def runner():
+            async with taskgroups.TaskGroup() as g:
+                g.create_task(hydra(g))
+                g.create_task(hercules())
+
+        with self.assertRaises(ExceptionGroup) as cm:
+            await runner()
+
+        self.assertEqual(get_error_types(cm.exception), {ZeroDivisionError})
+        self.assertGreaterEqual(nhydras, 10)