test async effect

Author: rmorshea

docs/source/life-cycle-hooks.rst

@@ -160,6 +160,12 @@ the effect will only occur when the given state changes:
 Now a new connection will only be estalished if a new ``url`` is provided.
 
 
+Async Effects
+.............
+
+under construction...
+
+
 **Supplementary Hooks**
 -----------------------
 
@@ -272,18 +278,6 @@ hook alongside :ref:`use_effect` or in response to element event handlers.
 :ref:`The Game Snake` provides a good use case for ``use_ref``.
 
 
-**Unique Hooks**
-----------------
-
-Hooks which are specific to IDOM and not present in React.
-
-
-use_async
----------
-
-...
-
-
 **Rules of Hooks**
 ------------------
 

idom/core/hooks.py

@@ -14,6 +14,7 @@
     NamedTuple,
     List,
     overload,
+    cast,
 )
 from typing_extensions import Protocol
 
@@ -87,7 +88,9 @@ def dispatch(
 
 
 _EffectCleanFunc = Callable[[], None]
-_EffectApplyFunc = Callable[[], Optional[_EffectCleanFunc]]
+_SyncEffectFunc = Callable[[], Optional[_EffectCleanFunc]]
+_AsyncEffectFunc = Callable[[], Awaitable[Optional[_EffectCleanFunc]]]
+_EffectApplyFunc = Union[_SyncEffectFunc, _AsyncEffectFunc]
 
 
 @overload
@@ -124,10 +127,32 @@ def use_effect(
     memoize = use_memo(args=args)
 
     def add_effect(function: _EffectApplyFunc) -> None:
+
+        if not asyncio.iscoroutinefunction(function):
+            sync_function = cast(_SyncEffectFunc, function)
+        else:
+
+            async_function = cast(_AsyncEffectFunc, function)
+
+            def sync_function() -> Optional[_EffectCleanFunc]:
+                future = asyncio.ensure_future(async_function())
+
+                def clean_future() -> None:
+                    try:
+                        clean = future.result()
+                    except asyncio.InvalidStateError:
+                        future.cancel()
+                    else:
+                        if clean is not None:
+                            clean()
+
+                return clean_future
+
         def effect() -> None:
-            clean = function()
+            clean = sync_function()
             if clean is not None:
                 hook.add_effect("will_render will_unmount", clean)
+            return None
 
         return memoize(lambda: hook.add_effect("did_render", effect))
 
@@ -313,53 +338,6 @@ def use_ref(initial_value: _StateType) -> Ref[_StateType]:
     return _use_const(lambda: Ref(initial_value))
 
 
-_AsyncCleanFunc = Callable[[], None]
-_AsyncApplyFunc = Callable[[], Awaitable[Optional[_EffectCleanFunc]]]
-
-
-@overload
-def use_async(
-    function: None = None, args: Optional[Sequence[Any]] = None
-) -> Callable[[_AsyncApplyFunc], None]:
-    ...
-
-
-@overload
-def use_async(function: _AsyncApplyFunc, args: Optional[Sequence[Any]] = None) -> None:
-    ...
-
-
-def use_async(
-    function: Optional[_AsyncApplyFunc] = None,
-    args: Optional[Sequence[Any]] = None,
-) -> Optional[Callable[[_AsyncApplyFunc], None]]:
-    effect = use_effect(args=args)
-
-    def add_asyc_effect(function: _AsyncApplyFunc) -> None:
-        def create_future() -> Callable[[], None]:
-            future = asyncio.ensure_future(function())
-
-            def cleanup_future() -> None:
-                try:
-                    cleanup = future.result()
-                except asyncio.InvalidStateError:
-                    future.cancel()
-                else:
-                    if cleanup is not None:
-                        cleanup()
-
-            return cleanup_future
-
-        effect(create_future)
-        return None
-
-    if function is not None:
-        add_asyc_effect(function)
-        return None
-    else:
-        return add_asyc_effect
-
-
 def _use_const(function: Callable[[], _StateType]) -> _StateType:
     return current_hook().use_state(function)
 

tests/test_core/test_hooks.py

@@ -1,4 +1,5 @@
 import re
+import asyncio
 
 import pytest
 
@@ -391,6 +392,79 @@ def effect():
         assert effect_run_count.current == 2
 
 
+async def test_use_async_effect():
+    effect_ran = asyncio.Event()
+
+    @idom.element
+    def ElementWithAsyncEffect():
+        @idom.hooks.use_effect
+        async def effect():
+            effect_ran.set()
+
+        return idom.html.div()
+
+    async with idom.Layout(ElementWithAsyncEffect()) as layout:
+        await layout.render()
+        await effect_ran.wait()
+
+
+async def test_use_async_effect_cleanup():
+    element_hook = HookCatcher()
+    effect_ran = asyncio.Event()
+    cleanup_ran = asyncio.Event()
+
+    @idom.element
+    @element_hook.capture
+    def ElementWithAsyncEffect():
+        @idom.hooks.use_effect
+        async def effect():
+            effect_ran.set()
+            return cleanup_ran.set
+
+        return idom.html.div()
+
+    async with idom.Layout(ElementWithAsyncEffect()) as layout:
+        await layout.render()
+
+        await effect_ran.wait()
+        element_hook.schedule_render()
+
+        await layout.render()
+
+    await asyncio.wait_for(cleanup_ran.wait(), 1)
+
+
+async def test_use_async_effect_cancel():
+    element_hook = HookCatcher()
+    effect_ran = asyncio.Event()
+    effect_was_cancelled = asyncio.Event()
+
+    event_that_is_never_set = asyncio.Event()
+
+    @idom.element
+    @element_hook.capture
+    def ElementWithLongWaitingEffect():
+        @idom.hooks.use_effect
+        async def effect():
+            effect_ran.set()
+            try:
+                await event_that_is_never_set.wait()
+            except asyncio.CancelledError:
+                effect_was_cancelled.set()
+
+        return idom.html.div()
+
+    async with idom.Layout(ElementWithLongWaitingEffect()) as layout:
+        await layout.render()
+
+        await effect_ran.wait()
+        element_hook.schedule_render()
+
+        await layout.render()
+
+    await asyncio.wait_for(effect_was_cancelled.wait(), 1)
+
+
 async def test_error_in_effect_is_gracefully_handled(caplog):
     @idom.element
     def ElementWithEffect():