More robust async event shielding

Author: Archmonger

src/reactpy/core/hooks.py

@@ -231,9 +231,13 @@ def decorator(func: _AsyncEffectFunc) -> None:
         async def effect(stop: asyncio.Event) -> None:
             # Make sure we always clean up the previous effect's resources
             if pending_task.current:
-                pending_task.current.cancel()
+                previous_task = pending_task.current
+                if not shield:
+                    previous_task.cancel()
                 with contextlib.suppress(asyncio.CancelledError):
-                    await pending_task.current
+                    await previous_task
+                if pending_task.current is previous_task:
+                    pending_task.current = None
 
             run_effect_cleanup(cleanup_func)
 
@@ -267,7 +271,9 @@ async def effect(stop: asyncio.Event) -> None:
                 if not shield:
                     task.cancel()
                 with contextlib.suppress(asyncio.CancelledError):
-                    await task
+                    cleanup_func.current = await task
+                if pending_task.current is task:
+                    pending_task.current = None
 
             # Run the clean-up function when the effect is stopped,
             # if it hasn't been run already by a new effect

tests/test_core/test_hooks.py

@@ -609,6 +609,7 @@ async def effect():
 async def test_use_async_effect_shield():
     component_hook = HookCatcher()
     effect_ran = asyncio.Event()
+    effect_was_cancelled = asyncio.Event()
     effect_finished = asyncio.Event()
     stop_waiting = asyncio.Event()
 
@@ -618,7 +619,11 @@ def ComponentWithShieldedEffect():
         @reactpy.hooks.use_async_effect(dependencies=None, shield=True)
         async def effect():
             effect_ran.set()
-            await stop_waiting.wait()
+            try:
+                await stop_waiting.wait()
+            except asyncio.CancelledError:
+                effect_was_cancelled.set()
+                raise
             effect_finished.set()
 
         return reactpy.html.div()
@@ -636,6 +641,7 @@ async def effect():
 
         # Verify effect hasn't finished yet but also wasn't cancelled
         assert not effect_finished.is_set()
+        assert not effect_was_cancelled.is_set()
 
         # Now allow the effect to finish
         stop_waiting.set()
@@ -644,6 +650,7 @@ async def effect():
         await layout.render()
 
     await asyncio.wait_for(effect_finished.wait(), 1)
+    assert not effect_was_cancelled.is_set()
 
 
 async def test_async_effect_sleep_is_cancelled_on_re_render():