simpler add_effect interface

Author: rmorshea

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

@@ -2,7 +2,7 @@
 
 import logging
 from asyncio import gather
-from collections.abc import AsyncGenerator
+from collections.abc import Awaitable
 from typing import Any, Callable, TypeVar
 
 from anyio import Semaphore
@@ -41,7 +41,7 @@ class LifeCycleHook:
         .. testcode::
 
             from reactpy.core._life_cycle_hook import LifeCycleHook
-            from reactpy.core.hooks import current_hook, COMPONENT_DID_RENDER_EFFECT
+            from reactpy.core.hooks import current_hook
 
             # this function will come from a layout implementation
             schedule_render = lambda: ...
@@ -63,7 +63,11 @@ class LifeCycleHook:
 
                 # and save state or add effects
                 current_hook().use_state(lambda: ...)
-                current_hook().add_effect(COMPONENT_DID_RENDER_EFFECT, lambda: ...)
+
+                async def effect():
+                    yield
+
+                current_hook().add_effect(effect)
             finally:
                 await hook.affect_component_did_render()
 
@@ -88,10 +92,10 @@ class LifeCycleHook:
         "__weakref__",
         "_context_providers",
         "_current_state_index",
-        "_pending_effects",
+        "_effect_cleanups",
+        "_effect_startups",
         "_render_access",
         "_rendered_atleast_once",
-        "_running_effects",
         "_schedule_render_callback",
         "_schedule_render_later",
         "_state",
@@ -110,8 +114,8 @@ def __init__(
         self._rendered_atleast_once = False
         self._current_state_index = 0
         self._state: tuple[Any, ...] = ()
-        self._pending_effects: list[AsyncGenerator[None, None]] = []
-        self._running_effects: list[AsyncGenerator[None, None]] = []
+        self._effect_startups: list[Callable[[], Awaitable[None]]] = []
+        self._effect_cleanups: list[Callable[[], Awaitable[None]]] = []
         self._render_access = Semaphore(1)  # ensure only one render at a time
 
     def schedule_render(self) -> None:
@@ -131,9 +135,14 @@ def use_state(self, function: Callable[[], T]) -> T:
         self._current_state_index += 1
         return result
 
-    def add_effect(self, effect_func: Callable[[], AsyncGenerator[None, None]]) -> None:
+    def add_effect(
+        self,
+        start_effect: Callable[[], Awaitable[None]],
+        clean_effect: Callable[[], Awaitable[None]],
+    ) -> None:
         """Add an effect to this hook"""
-        self._pending_effects.append(effect_func())
+        self._effect_startups.append(start_effect)
+        self._effect_cleanups.append(clean_effect)
 
     def set_context_provider(self, provider: ContextProviderType[Any]) -> None:
         self._context_providers[provider.type] = provider
@@ -155,29 +164,28 @@ async def affect_component_did_render(self) -> None:
         self._rendered_atleast_once = True
         self._current_state_index = 0
         self._render_access.release()
+        del self.component
 
     async def affect_layout_did_render(self) -> None:
         """The layout completed a render"""
         try:
-            await gather(*[g.asend(None) for g in self._pending_effects])
-            self._running_effects.extend(self._pending_effects)
+            await gather(*[start() for start in self._effect_startups])
         except Exception:
             logger.exception("Error during effect startup")
         finally:
-            self._pending_effects.clear()
-        if self._schedule_render_later:
-            self._schedule_render()
-        self._schedule_render_later = False
-        del self.component
+            self._effect_startups.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"""
         try:
-            await gather(*[g.aclose() for g in self._running_effects])
+            await gather(*[clean() for clean in self._effect_cleanups])
         except Exception:
             logger.exception("Error during effect cleanup")
         finally:
-            self._running_effects.clear()
+            self._effect_cleanups.clear()
 
     def set_current(self) -> None:
         """Set this hook as the active hook in this thread

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

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import AsyncGenerator, Awaitable, Sequence
+from collections.abc import Coroutine, Sequence
 from logging import getLogger
 from types import FunctionType
 from typing import (
@@ -95,7 +95,9 @@ def dispatch(new: _Type | Callable[[_Type], _Type]) -> None:
 
 _EffectCleanFunc: TypeAlias = "Callable[[], None]"
 _SyncEffectFunc: TypeAlias = "Callable[[], _EffectCleanFunc | None]"
-_AsyncEffectFunc: TypeAlias = "Callable[[], Awaitable[_EffectCleanFunc | None]]"
+_AsyncEffectFunc: TypeAlias = (
+    "Callable[[], Coroutine[None, None, _EffectCleanFunc | None]]"
+)
 _EffectApplyFunc: TypeAlias = "_SyncEffectFunc | _AsyncEffectFunc"
 
 
@@ -146,36 +148,28 @@ def add_effect(function: _EffectApplyFunc) -> None:
             async_function = cast(_AsyncEffectFunc, function)
 
             def sync_function() -> _EffectCleanFunc | None:
-                future = asyncio.ensure_future(async_function())
+                task = asyncio.create_task(async_function())
 
                 def clean_future() -> None:
-                    if not future.cancel():
-                        clean = future.result()
+                    if not task.cancel():
+                        clean = task.result()
                         if clean is not None:
                             clean()
 
                 return clean_future
 
-        async def effect() -> AsyncGenerator[None, None]:
+        async def start_effect() -> None:
             if last_clean_callback.current is not None:
                 last_clean_callback.current()
+                last_clean_callback.current = None
+            last_clean_callback.current = sync_function()
 
-            cleaned = False
-            clean = sync_function()
-
-            def callback() -> None:
-                nonlocal cleaned
-                if clean and not cleaned:
-                    cleaned = True
-                    clean()
-
-            last_clean_callback.current = callback
-            try:
-                yield
-            finally:
-                callback()
+        async def clean_effect() -> None:
+            if last_clean_callback.current is not None:
+                last_clean_callback.current()
+                last_clean_callback.current = None
 
-        return memoize(lambda: hook.add_effect(effect))
+        return memoize(lambda: hook.add_effect(start_effect, clean_effect))
 
     if function is not None:
         add_effect(function)