async effect context

Author: rmorshea

pyproject.toml

@@ -130,8 +130,9 @@ ignore = [
   "PLR0915",
 ]
 unfixable = [
-  # Don't touch unused imports
+  # Don't touch unused imports or unused variables
   "F401",
+  "F841",
 ]
 
 [tool.ruff.isort]

src/py/reactpy/reactpy/core/_life_cycle_hook.py

@@ -3,50 +3,30 @@
 import asyncio
 import logging
 from collections.abc import Coroutine
-from dataclasses import dataclass
 from typing import Any, Callable, TypeVar
-from weakref import WeakSet
-
-from typing_extensions import TypeAlias
 
 from reactpy.core._thread_local import ThreadLocal
 from reactpy.core.types import ComponentType, Context, ContextProviderType
 
 T = TypeVar("T")
+
+StopEffect = Callable[[], Coroutine[None, None, None]]
+StartEffect = Callable[[], Coroutine[None, None, StopEffect]]
+
 logger = logging.getLogger(__name__)
 
+_HOOK_STATE: ThreadLocal[list[LifeCycleHook]] = ThreadLocal(list)
+
 
 def current_hook() -> LifeCycleHook:
     """Get the current :class:`LifeCycleHook`"""
-    hook_stack = _hook_stack.get()
+    hook_stack = _HOOK_STATE.get()
     if not hook_stack:
         msg = "No life cycle hook is active. Are you rendering in a layout?"
         raise RuntimeError(msg)
     return hook_stack[-1]
 
 
-_hook_stack: ThreadLocal[list[LifeCycleHook]] = ThreadLocal(list)
-
-
-@dataclass(frozen=True)
-class EffectInfo:
-    task: asyncio.Task[None]
-    stop: asyncio.Event
-
-    async def signal_stop(self, timeout: float) -> None:
-        """Signal the effect to stop and wait for it to complete."""
-        self.stop.set()
-        try:
-            await asyncio.wait_for(self.task, timeout=timeout)
-        finally:
-            # a no-op if the task has already completed
-            if self.task.cancel():
-                try:
-                    await self.task
-                except asyncio.CancelledError:
-                    logger.exception("Effect failed to stop after %s seconds", timeout)
-
-
 class LifeCycleHook:
     """Defines the life cycle of a layout component.
 
@@ -114,7 +94,8 @@ class LifeCycleHook:
         "_context_providers",
         "_current_state_index",
         "_effect_funcs",
-        "_effect_infos",
+        "_effect_starts",
+        "_effect_stops",
         "_is_rendering",
         "_rendered_atleast_once",
         "_schedule_render_callback",
@@ -136,8 +117,8 @@ def __init__(
         self._rendered_atleast_once = False
         self._current_state_index = 0
         self._state: tuple[Any, ...] = ()
-        self._effect_funcs: list[_EffectStarter] = []
-        self._effect_infos: WeakSet[EffectInfo] = WeakSet()
+        self._effect_starts: list[StartEffect] = []
+        self._effect_stops: list[StopEffect] = []
 
     def schedule_render(self) -> None:
         if self._is_rendering:
@@ -156,9 +137,9 @@ def use_state(self, function: Callable[[], T]) -> T:
         self._current_state_index += 1
         return result
 
-    def add_effect(self, start_effect: _EffectStarter) -> None:
-        """Trigger a function on the occurrence of the given effect type"""
-        self._effect_funcs.append(start_effect)
+    def add_effect(self, start_effect: StartEffect) -> None:
+        """Add an effect to this hook"""
+        self._effect_starts.append(start_effect)
 
     def set_context_provider(self, provider: ContextProviderType[Any]) -> None:
         self._context_providers[provider.type] = provider
@@ -184,40 +165,39 @@ async def affect_component_did_render(self) -> None:
 
     async def affect_layout_did_render(self) -> None:
         """The layout completed a render"""
-        for start_effect in self._effect_funcs:
-            effect_info = await start_effect()
-            self._effect_infos.add(effect_info)
-        self._effect_funcs.clear()
+        self._effect_stops.extend(
+            await asyncio.gather(*[start() for start in self._effect_starts])
+        )
+        self._effect_starts.clear()
 
         if self._schedule_render_later:
             self._schedule_render()
         self._schedule_render_later = False
 
     async def affect_component_will_unmount(self) -> None:
         """The component is about to be removed from the layout"""
-        for infos in self._effect_infos:
-            infos.stop.set()
         try:
-            await asyncio.gather(*[i.task for i in self._effect_infos])
+            await asyncio.gather(*[stop() for stop in self._effect_stops])
         except Exception:
             logger.exception("Error during effect cancellation")
-        self._effect_infos.clear()
+        finally:
+            self._effect_stops.clear()
 
     def set_current(self) -> None:
         """Set this hook as the active hook in this thread
 
         This method is called by a layout before entering the render method
         of this hook's associated component.
         """
-        hook_stack = _hook_stack.get()
+        hook_stack = _HOOK_STATE.get()
         if hook_stack:
             parent = hook_stack[-1]
             self._context_providers.update(parent._context_providers)
         hook_stack.append(self)
 
     def unset_current(self) -> None:
         """Unset this hook as the active hook in this thread"""
-        if _hook_stack.get().pop() is not self:
+        if _HOOK_STATE.get().pop() is not self:
             raise RuntimeError("Hook stack is in an invalid state")  # nocov
 
     def _schedule_render(self) -> None:
@@ -227,6 +207,3 @@ def _schedule_render(self) -> None:
             logger.exception(
                 f"Failed to schedule render via {self._schedule_render_callback}"
             )
-
-
-_EffectStarter: TypeAlias = "Callable[[], Coroutine[None, None, EffectInfo]]"

src/py/reactpy/reactpy/core/hooks.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import inspect
+import sys
 import warnings
 from collections.abc import Coroutine, Sequence
 from logging import getLogger
@@ -17,10 +18,10 @@
     overload,
 )
 
-from typing_extensions import TypeAlias
+from typing_extensions import Self, TypeAlias
 
-from reactpy.config import REACTPY_DEBUG_MODE, REACTPY_EFFECT_DEFAULT_STOP_TIMEOUT
-from reactpy.core._life_cycle_hook import EffectInfo, current_hook
+from reactpy.config import REACTPY_DEBUG_MODE
+from reactpy.core._life_cycle_hook import StopEffect, current_hook
 from reactpy.core.types import Context, Key, State, VdomDict
 from reactpy.utils import Ref
 
@@ -96,15 +97,14 @@ def dispatch(new: _Type | Callable[[_Type], _Type]) -> None:
 
 _EffectCleanFunc: TypeAlias = "Callable[[], None]"
 _SyncEffectFunc: TypeAlias = "Callable[[], _EffectCleanFunc | None]"
-_AsyncEffectFunc: TypeAlias = "Callable[[asyncio.Event], Coroutine[None, None, None]]"
+_AsyncEffectFunc: TypeAlias = "Callable[[Effect], Coroutine[None, None, None]]"
 _EffectFunc: TypeAlias = "_SyncEffectFunc | _AsyncEffectFunc"
 
 
 @overload
 def use_effect(
     function: None = None,
     dependencies: Sequence[Any] | ellipsis | None = ...,
-    stop_timeout: float = ...,
 ) -> Callable[[_EffectFunc], None]:
     ...
 
@@ -113,15 +113,13 @@ def use_effect(
 def use_effect(
     function: _EffectFunc,
     dependencies: Sequence[Any] | ellipsis | None = ...,
-    stop_timeout: float = ...,
 ) -> None:
     ...
 
 
 def use_effect(
     function: _EffectFunc | None = None,
     dependencies: Sequence[Any] | ellipsis | None = ...,
-    stop_timeout: float = REACTPY_EFFECT_DEFAULT_STOP_TIMEOUT.current,
 ) -> Callable[[_EffectFunc], None] | None:
     """See the full :ref:`Use Effect` docs for details
 
@@ -145,22 +143,21 @@ def use_effect(
     hook = current_hook()
     dependencies = _try_to_infer_closure_values(function, dependencies)
     memoize = use_memo(dependencies=dependencies)
-    effect_info: Ref[EffectInfo | None] = use_ref(None)
+    effect_ref: Ref[Effect | None] = use_ref(None)
 
     def add_effect(function: _EffectFunc) -> None:
-        effect = _cast_async_effect(function)
+        effect_func = _cast_async_effect(function)
 
-        async def create_effect_task() -> EffectInfo:
-            if effect_info.current is not None:
-                last_effect_info = effect_info.current
-                await last_effect_info.signal_stop(stop_timeout)
+        async def start_effect() -> StopEffect:
+            if effect_ref.current is not None:
+                await effect_ref.current.stop()
 
-            stop = asyncio.Event()
-            info = EffectInfo(asyncio.create_task(effect(stop)), stop)
-            effect_info.current = info
-            return info
+            effect = effect_ref.current = Effect()
+            effect.task = asyncio.create_task(effect_func(effect))
 
-        return memoize(lambda: hook.add_effect(create_effect_task))
+            return effect.stop
+
+        return memoize(lambda: hook.add_effect(start_effect))
 
     if function is not None:
         add_effect(function)
@@ -169,47 +166,114 @@ async def create_effect_task() -> EffectInfo:
         return add_effect
 
 
+class Effect:
+    """A context manager for running asynchronous effects."""
+
+    task: asyncio.Task[Any]
+    """The task that is running the effect."""
+
+    def __init__(self) -> None:
+        self._stop = asyncio.Event()
+        self._started = False
+        self._cancel_count = 0
+
+    async def stop(self) -> None:
+        """Signal the effect to stop."""
+        if self._started:
+            self._cancel_task()
+        self._stop.set()
+        try:
+            await self.task
+        except asyncio.CancelledError:
+            pass
+        except Exception:
+            logger.exception("Error while stopping effect")
+
+    async def __aenter__(self) -> Self:
+        self._started = True
+        self._cancel_count = self.task.cancelling()
+        if self._stop.is_set():
+            self._cancel_task()
+        return self
+
+    if sys.version_info < (3, 11):  # nocov
+        # Python<3.11 doesn't have Task.cancelling so we need to track it ourselves.
+
+        _3_11__aenter__ = __aenter__
+
+        async def __aenter__(self) -> Self:
+            cancel_count = 0
+            old_cancel = self.task.cancel
+
+            def new_cancel(*a, **kw) -> None:
+                nonlocal cancel_count
+                cancel_count += 1
+                return old_cancel(*a, **kw)
+
+            self.task.cancel = new_cancel
+            self.task.cancelling = lambda: cancel_count
+
+            return await _3_11__aenter__(self)  # noqa: F821
+
+    async def __aexit__(self, exc_type: type[BaseException], *exc: Any) -> Any:
+        if exc_type is not None and not issubclass(exc_type, asyncio.CancelledError):
+            # propagate non-cancellation exceptions
+            return None
+
+        try:
+            await self._stop.wait()
+        except asyncio.CancelledError:
+            if self.task.cancelling() > self._cancel_count:
+                # Task has been cancelled by something else - propagate it
+                return None
+
+        return True
+
+    def _cancel_task(self) -> None:
+        self.task.cancel()
+        self._cancel_count += 1
+
+
 def _cast_async_effect(function: Callable[..., Any]) -> _AsyncEffectFunc:
     if inspect.iscoroutinefunction(function):
         if len(inspect.signature(function).parameters):
             return function
 
         warnings.warn(
-            'Async effect functions should accept a "stop" asyncio.Event as their '
+            "Async effect functions should accept an Effect context manager as their "
             "first argument. This will be required in a future version of ReactPy.",
             stacklevel=3,
         )
 
-        async def wrapper(stop: asyncio.Event) -> None:
-            task = asyncio.create_task(function())
-            await stop.wait()
-            if not task.cancel():
+        async def wrapper(effect: Effect) -> None:
+            cleanup = None
+            async with effect:
                 try:
-                    cleanup = await task
+                    cleanup = await function()
                 except Exception:
                     logger.exception("Error while applying effect")
-                    return
-                if cleanup is not None:
-                    try:
-                        cleanup()
-                    except Exception:
-                        logger.exception("Error while cleaning up effect")
+            if cleanup is not None:
+                try:
+                    cleanup()
+                except Exception:
+                    logger.exception("Error while cleaning up effect")
 
         return wrapper
     else:
 
-        async def wrapper(stop: asyncio.Event) -> None:
-            try:
-                cleanup = function()
-            except Exception:
-                logger.exception("Error while applying effect")
-                return
-            await stop.wait()
-            try:
-                if cleanup is not None:
+        async def wrapper(effect: Effect) -> None:
+            cleanup = None
+            async with effect:
+                try:
+                    cleanup = function()
+                except Exception:
+                    logger.exception("Error while applying effect")
+
+            if cleanup is not None:
+                try:
                     cleanup()
-            except Exception:
-                logger.exception("Error while cleaning up effect")
+                except Exception:
+                    logger.exception("Error while cleaning up effect")
 
         return wrapper