async effects get prior task and interupt event

docs/source/about/changelog.rst

@@ -23,7 +23,21 @@ more info, see the :ref:`Contributor Guide <Creating a Changelog Entry>`.
 Unreleased
 ----------
 
-No changes.
+**Fixed**
+
+- :issue:`956` - Async effects now accept two arguments - the prior effect's
+  ``asyncio.Task`` (or ``None``) and an interupt ``asyncio.Event``. The prior effect's
+  task allows effect authors to await or cancel it to ensure that it has completed
+  before executing the next effect. The interupt event is used to signal that the effect
+  should stop and clean up any resources it may have allocated. This is useful for
+  effects that may be long running and need to be stopped when the component is
+  unmounted.
+
+**Deprecated**
+
+- :pull:`957` - Async effects that do not accept any arguments are now deprecated and
+  will be disallowed in a future release. All async effects should accept two
+  arguments - the prior effect's task and an interupt event.
 
 
 v1.0.0

docs/source/reference/_examples/simple_dashboard.py

@@ -97,7 +97,7 @@ async def interval() -> None:
         await asyncio.sleep(rate - (time.time() - usage_time.current))
         usage_time.current = time.time()
 
-    return asyncio.ensure_future(interval())
+    return asyncio.create_task(interval())
 
 
 reactpy.run(RandomWalk)

docs/source/reference/_examples/snake_game.py

@@ -135,7 +135,7 @@ async def interval() -> None:
         await asyncio.sleep(rate - (time.time() - usage_time.current))
         usage_time.current = time.time()
 
-    return asyncio.ensure_future(interval())
+    return asyncio.create_task(interval())
 
 
 def create_grid(grid_size, block_scale):

src/reactpy/backend/flask.py

@@ -242,7 +242,7 @@ async def main() -> None:
                 async_recv_queue.get,
             )
 
-        main_future = asyncio.ensure_future(main(), loop=loop)
+        main_future = asyncio.create_task(main(), loop=loop)
 
         dispatch_thread_info_ref.current = _DispatcherThreadInfo(
             dispatch_loop=loop,

src/reactpy/backend/tornado.py

@@ -190,7 +190,7 @@ async def recv() -> Any:
             return json.loads(await message_queue.get())
 
         self._message_queue = message_queue
-        self._dispatch_future = asyncio.ensure_future(
+        self._dispatch_future = asyncio.create_task(
             serve_layout(
                 Layout(
                     ConnectionContext(

src/reactpy/core/hooks.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 import asyncio
+from asyncio import Event, Task, create_task
+from inspect import iscoroutinefunction, isfunction, signature
 from logging import getLogger
 from types import FunctionType
 from typing import (
     TYPE_CHECKING,
     Any,
-    Awaitable,
     Callable,
     Generic,
     NewType,
@@ -15,14 +16,15 @@
     cast,
     overload,
 )
+from warnings import warn
 
-from typing_extensions import Protocol, TypeAlias
+from typing_extensions import Literal, Protocol, TypeAlias, TypedDict
 
 from reactpy.config import REACTPY_DEBUG_MODE
 from reactpy.utils import Ref
 
 from ._thread_local import ThreadLocal
-from .types import ComponentType, Key, State, VdomDict
+from .types import AsyncEffect, ComponentType, Key, State, SyncEffect, VdomDict
 
 
 if not TYPE_CHECKING:
@@ -96,32 +98,26 @@ def dispatch(new: _Type | Callable[[_Type], _Type]) -> None:
         self.dispatch = dispatch
 
 
-_EffectCleanFunc: TypeAlias = "Callable[[], None]"
-_SyncEffectFunc: TypeAlias = "Callable[[], _EffectCleanFunc | None]"
-_AsyncEffectFunc: TypeAlias = "Callable[[], Awaitable[_EffectCleanFunc | None]]"
-_EffectApplyFunc: TypeAlias = "_SyncEffectFunc | _AsyncEffectFunc"
-
-
 @overload
 def use_effect(
     function: None = None,
     dependencies: Sequence[Any] | ellipsis | None = ...,
-) -> Callable[[_EffectApplyFunc], None]:
+) -> Callable[[SyncEffect | AsyncEffect], None]:
     ...
 
 
 @overload
 def use_effect(
-    function: _EffectApplyFunc,
+    function: SyncEffect | AsyncEffect,
     dependencies: Sequence[Any] | ellipsis | None = ...,
 ) -> None:
     ...
 
 
 def use_effect(
-    function: _EffectApplyFunc | None = None,
+    function: SyncEffect | AsyncEffect | None = None,
     dependencies: Sequence[Any] | ellipsis | None = ...,
-) -> Callable[[_EffectApplyFunc], None] | None:
+) -> Callable[[SyncEffect | AsyncEffect], None] | None:
     """See the full :ref:`Use Effect` docs for details
 
     Parameters:
@@ -140,42 +136,65 @@ def use_effect(
 
     dependencies = _try_to_infer_closure_values(function, dependencies)
     memoize = use_memo(dependencies=dependencies)
-    last_clean_callback: Ref[_EffectCleanFunc | None] = use_ref(None)
+    state: _EffectState = _use_const(
+        lambda: {"prior_task": None, "prior_callback": None}
+    )
 
-    def add_effect(function: _EffectApplyFunc) -> None:
-        if not asyncio.iscoroutinefunction(function):
-            sync_function = cast(_SyncEffectFunc, function)
-        else:
-            async_function = cast(_AsyncEffectFunc, function)
+    def add_effect(function: SyncEffect | AsyncEffect) -> None:
+        memoize(lambda: _add_effect(hook, state, function))
+        return None
 
-            def sync_function() -> _EffectCleanFunc | None:
-                future = asyncio.ensure_future(async_function())
+    if function is not None:
+        add_effect(function)
+        return None
+    else:
+        return add_effect
 
-                def clean_future() -> None:
-                    if not future.cancel():
-                        clean = future.result()
-                        if clean is not None:
-                            clean()
 
-                return clean_future
+def _add_effect(
+    hook: LifeCycleHook, state: _EffectState, function: SyncEffect | AsyncEffect
+) -> None:
+    sync_function: SyncEffect
+
+    if iscoroutinefunction(function):
+        if not signature(function).parameters:  # pragma: no cover
+            warn(
+                f"Async effect functions {function} should accept two arguments - the "
+                "prior task and an interrupt event. This will be required in a future "
+                "release.",
+                DeprecationWarning,
+            )
+            original_function = function
 
-        def effect() -> None:
-            if last_clean_callback.current is not None:
-                last_clean_callback.current()
+            def function(prior_task: Task | None, _: Event) -> None:
+                if prior_task is not None:
+                    prior_task.cancel()
+                return original_function()
 
-            clean = last_clean_callback.current = sync_function()
-            if clean is not None:
-                hook.add_effect(COMPONENT_WILL_UNMOUNT_EFFECT, clean)
+        def sync_function() -> Callable[[], None]:
+            interupt = Event()
+            state["prior_task"] = create_task(function(state["prior_task"], interupt))
+            return interupt.set
 
-            return None
+    elif isfunction(function):
+        sync_function = function
+    else:
+        raise TypeError(f"Expected a function, not {function!r}")
 
-        return memoize(lambda: hook.add_effect(LAYOUT_DID_RENDER_EFFECT, effect))
+    def effect() -> None:
+        prior_callback = state["prior_callback"]
+        if prior_callback is not None:
+            prior_callback()
+        next_callback = state["prior_callback"] = sync_function()
+        if next_callback is not None:
+            hook.add_effect(COMPONENT_WILL_UNMOUNT_EFFECT, next_callback)
 
-    if function is not None:
-        add_effect(function)
-        return None
-    else:
-        return add_effect
+    hook.add_effect(LAYOUT_DID_RENDER_EFFECT, effect)
+
+
+class _EffectState(TypedDict):
+    prior_task: Task | None
+    prior_callback: Callable[[], None] | None
 
 
 def use_debug_value(

src/reactpy/core/types.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+import asyncio
 import sys
 from collections import namedtuple
 from collections.abc import Sequence
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
     Any,
+    Awaitable,
     Callable,
     Generic,
     Mapping,
@@ -233,3 +235,12 @@ class LayoutEventMessage(TypedDict):
     """The ID of the event handler."""
     data: Sequence[Any]
     """A list of event data passed to the event handler."""
+
+
+SyncEffect: TypeAlias = "Callable[[], None | Callable[[], None]]"
+"""A synchronous function which can be run by the :func:`use_effect` hook"""
+
+AsyncEffect: TypeAlias = (
+    "Callable[[asyncio.Task | None, asyncio.Event], Awaitable[None]]"
+)
+"""A asynchronous function which can be run by the :func:`use_effect` hook"""

tests/test_core/test_hooks.py

@@ -479,7 +479,7 @@ async def test_use_async_effect():
     @reactpy.component
     def ComponentWithAsyncEffect():
         @reactpy.hooks.use_effect
-        async def effect():
+        async def effect(prior, interupt):
             effect_ran.set()
 
         return reactpy.html.div()
@@ -498,9 +498,10 @@ async def test_use_async_effect_cleanup():
     @component_hook.capture
     def ComponentWithAsyncEffect():
         @reactpy.hooks.use_effect(dependencies=None)  # force this to run every time
-        async def effect():
+        async def effect(prior, interupt):
             effect_ran.set()
-            return cleanup_ran.set
+            await interupt.wait()
+            cleanup_ran.set()
 
         return reactpy.html.div()
 
@@ -514,7 +515,63 @@ async def effect():
     await asyncio.wait_for(cleanup_ran.wait(), 1)
 
 
-async def test_use_async_effect_cancel(caplog):
+async def test_use_async_effect_cancel():
+    component_hook = HookCatcher()
+    effect_ran = asyncio.Event()
+    effect_was_cancelled = asyncio.Event()
+
+    event_that_never_occurs = asyncio.Event()
+
+    @reactpy.component
+    @component_hook.capture
+    def ComponentWithLongWaitingEffect():
+        @reactpy.hooks.use_effect(dependencies=None)  # force this to run every time
+        async def effect(prior, interupt):
+            if prior is not None:
+                prior.cancel()
+            effect_ran.set()
+            try:
+                await event_that_never_occurs.wait()
+            except asyncio.CancelledError:
+                effect_was_cancelled.set()
+                raise
+
+        return reactpy.html.div()
+
+    async with reactpy.Layout(ComponentWithLongWaitingEffect()) as layout:
+        await layout.render()
+
+        await effect_ran.wait()
+        component_hook.latest.schedule_render()
+
+        await layout.render()
+
+    await asyncio.wait_for(effect_was_cancelled.wait(), 1)
+
+    # So I know we said the event never occurs but... to ensure the effect's future is
+    # cancelled before the test is cleaned up we need to set the event. This is because
+    # the cancellation doesn't propogate before the test is resolved which causes
+    # delayed log messages that impact other tests.
+    event_that_never_occurs.set()
+
+
+async def test_deprecated_use_async_effect_no_arguments():
+    effect_ran = asyncio.Event()
+
+    @reactpy.component
+    def ComponentWithAsyncEffect():
+        @reactpy.hooks.use_effect
+        async def effect():
+            effect_ran.set()
+
+        return reactpy.html.div()
+
+    async with reactpy.Layout(ComponentWithAsyncEffect()) as layout:
+        await layout.render()
+        await asyncio.wait_for(effect_ran.wait(), 1)
+
+
+async def test_deprecated_use_async_effect_cancel_no_arguments():
     component_hook = HookCatcher()
     effect_ran = asyncio.Event()
     effect_was_cancelled = asyncio.Event()

tests/test_core/test_serve.py

@@ -122,7 +122,7 @@ async def handle_event():
     send_queue = asyncio.Queue()
     recv_queue = asyncio.Queue()
 
-    asyncio.ensure_future(
+    asyncio.create_task(
         serve_layout(
             reactpy.Layout(ComponentWithTwoEventHandlers()),
             send_queue.put,