add a version of the test with weakrefs

python · kumaraditya303 · Jan 7, 2025 · Jan 6, 2025 · Jan 6, 2025 · Jan 6, 2025
commit 7c966d0bcd5900b9f3e6a98f40233c6ef4d3b191
@@ -1,6 +1,7 @@
 # Adapted with permission from the EdgeDB project;
 # license: PSFL.
 
+import weakref
 import sys
 import gc
 import asyncio
@@ -33,6 +34,24 @@ def no_other_refs():
     return [coro]
 
 
+def set_gc_state(enabled: bool) -> bool:
+    was_enabled = gc.isenabled()
+    if enabled:
+        gc.enable()
+    else:
+        gc.disable()
+    return was_enabled
+
+
+@contextlib.contextmanager
+def disable_gc() -> Generator[None]:
+    was_enabled = set_gc_state(enabled=False)
+    try:
+        yield
+    finally:
+        set_gc_state(enabled=was_enabled)
+
+
 class BaseTestTaskGroup:
 
     async def test_taskgroup_01(self):
@@ -960,6 +979,30 @@ async def coro_fn():
         self.assertIsInstance(exc, _Done)
         self.assertListEqual(gc.get_referrers(exc), no_other_refs())
 
+
+    async def test_exception_refcycles_parent_task_wr(self):
+        """Test that TaskGroup deletes self._parent_task and create_task() deletes task"""
+        tg = asyncio.TaskGroup()
+        exc = None
+
+        class _Done(Exception):
+            pass
+
+        async def coro_fn():
+            async with tg:
+                raise _Done
+
+        with disable_gc():
+            try:
+                async with asyncio.TaskGroup() as tg2:
+                    task_wr = weakref.ref(tg2.create_task(coro_fn()))
+            except* _Done as excs:
+                exc = excs.exceptions[0].exceptions[0]
+
+        self.assertIsNone(task_wr())
+        self.assertIsInstance(exc, _Done)
+        self.assertListEqual(gc.get_referrers(exc), no_other_refs())
+
     async def test_exception_refcycles_propagate_cancellation_error(self):
         """Test that TaskGroup deletes propagate_cancellation_error"""
         tg = asyncio.TaskGroup()