automatically infer closure arguments

Author: rmorshea

src/idom/core/hooks.py

@@ -8,6 +8,7 @@
 import asyncio
 from logging import getLogger
 from threading import get_ident as get_thread_id
+from types import FunctionType
 from typing import (
     Any,
     Awaitable,
@@ -101,21 +102,23 @@ def dispatch(
 
 @overload
 def use_effect(
-    function: None = None, args: Optional[Sequence[Any]] = None
+    function: None = None,
+    args: Sequence[Any] | "ellipsis" | None = ...,
 ) -> Callable[[_EffectApplyFunc], None]:
     ...
 
 
 @overload
 def use_effect(
-    function: _EffectApplyFunc, args: Optional[Sequence[Any]] = None
+    function: _EffectApplyFunc,
+    args: Sequence[Any] | "ellipsis" | None = ...,
 ) -> None:
     ...
 
 
 def use_effect(
     function: Optional[_EffectApplyFunc] = None,
-    args: Optional[Sequence[Any]] = None,
+    args: Sequence[Any] | "ellipsis" | None = ...,
 ) -> Optional[Callable[[_EffectApplyFunc], None]]:
     """See the full :ref:`Use Effect` docs for details
 
@@ -130,6 +133,8 @@ def use_effect(
         If not function is provided, a decorator. Otherwise ``None``.
     """
     hook = current_hook()
+
+    args = _try_to_infer_closure_args(function, args)
     memoize = use_memo(args=args)
     last_clean_callback: Ref[Optional[_EffectCleanFunc]] = use_ref(None)
 
@@ -209,21 +214,23 @@ def dispatch(action: _ActionType) -> None:
 
 @overload
 def use_callback(
-    function: None = None, args: Optional[Sequence[Any]] = None
+    function: None = None,
+    args: Sequence[Any] | "ellipsis" | None = ...,
 ) -> Callable[[_CallbackFunc], _CallbackFunc]:
     ...
 
 
 @overload
 def use_callback(
-    function: _CallbackFunc, args: Optional[Sequence[Any]] = None
+    function: _CallbackFunc,
+    args: Sequence[Any] | "ellipsis" | None = ...,
 ) -> _CallbackFunc:
     ...
 
 
 def use_callback(
     function: Optional[_CallbackFunc] = None,
-    args: Optional[Sequence[Any]] = (),
+    args: Sequence[Any] | "ellipsis" | None = ...,
 ) -> Union[_CallbackFunc, Callable[[_CallbackFunc], _CallbackFunc]]:
     """See the full :ref:`Use Callback` docs for details
 
@@ -234,6 +241,7 @@ def use_callback(
     Returns:
         The current function
     """
+    args = _try_to_infer_closure_args(function, args)
     memoize = use_memo(args=args)
 
     def setup(function: _CallbackFunc) -> _CallbackFunc:
@@ -254,21 +262,23 @@ def __call__(self, func: Callable[[], _StateType]) -> _StateType:
 
 @overload
 def use_memo(
-    function: None = None, args: Optional[Sequence[Any]] = None
+    function: None = None,
+    args: Sequence[Any] | "ellipsis" | None = ...,
 ) -> _LambdaCaller:
     ...
 
 
 @overload
 def use_memo(
-    function: Callable[[], _StateType], args: Optional[Sequence[Any]] = None
+    function: Callable[[], _StateType],
+    args: Sequence[Any] | "ellipsis" | None = ...,
 ) -> _StateType:
     ...
 
 
 def use_memo(
     function: Optional[Callable[[], _StateType]] = None,
-    args: Optional[Sequence[Any]] = None,
+    args: Sequence[Any] | "ellipsis" | None = ...,
 ) -> Union[_StateType, Callable[[Callable[[], _StateType]], _StateType]]:
     """See the full :ref:`Use Memo` docs for details
 
@@ -279,6 +289,8 @@ def use_memo(
     Returns:
         The current state
     """
+    args = _try_to_infer_closure_args(function, args)
+
     memo: _Memo[_StateType] = _use_const(_Memo)
 
     if memo.empty():
@@ -350,6 +362,23 @@ def _use_const(function: Callable[[], _StateType]) -> _StateType:
     return current_hook().use_state(function)
 
 
+def _try_to_infer_closure_args(
+    func: Callable[..., Any] | None,
+    args: Sequence[Any] | "ellipsis" | None,
+) -> Sequence[Any] | None:
+    if args is ...:
+        if isinstance(func, FunctionType):
+            return (
+                [cell.cell_contents for cell in func.__closure__]
+                if func.__closure__
+                else []
+            )
+        else:
+            return None
+    else:
+        return cast("Sequence[Any] | None", args)
+
+
 _current_life_cycle_hook: Dict[int, "LifeCycleHook"] = {}
 
 

tests/test_core/test_hooks.py

@@ -331,7 +331,7 @@ async def test_use_effect_cleanup_occurs_before_next_effect():
     @idom.component
     @component_hook.capture
     def ComponentWithEffect():
-        @idom.hooks.use_effect
+        @idom.hooks.use_effect(args=None)
         def effect():
             if cleanup_triggered.current:
                 cleanup_triggered_before_next_effect.current = True
@@ -485,7 +485,7 @@ async def effect():
 
     with idom.Layout(ComponentWithAsyncEffect()) as layout:
         await layout.render()
-        await effect_ran.wait()
+        await asyncio.wait_for(effect_ran.wait(), 1)
 
 
 async def test_use_async_effect_cleanup():
@@ -496,7 +496,7 @@ async def test_use_async_effect_cleanup():
     @idom.component
     @component_hook.capture
     def ComponentWithAsyncEffect():
-        @idom.hooks.use_effect
+        @idom.hooks.use_effect(args=None)  # force this to run every time
         async def effect():
             effect_ran.set()
             return cleanup_ran.set
@@ -506,7 +506,7 @@ async def effect():
     with idom.Layout(ComponentWithAsyncEffect()) as layout:
         await layout.render()
 
-        await effect_ran.wait()
+        cleanup_ran.wait()
         component_hook.latest.schedule_render()
 
         await layout.render()
@@ -524,7 +524,7 @@ async def test_use_async_effect_cancel(caplog):
     @idom.component
     @component_hook.capture
     def ComponentWithLongWaitingEffect():
-        @idom.hooks.use_effect
+        @idom.hooks.use_effect(args=None)  # force this to run every time
         async def effect():
             effect_ran.set()
             try:
@@ -575,7 +575,7 @@ async def test_error_in_effect_cleanup_is_gracefully_handled(caplog):
     @idom.component
     @component_hook.capture
     def ComponentWithEffect():
-        @idom.hooks.use_effect
+        @idom.hooks.use_effect(args=None)  # force this to run every time
         def ok_effect():
             def bad_cleanup():
                 raise ValueError("Something went wong :(")
@@ -765,7 +765,7 @@ async def test_use_memo_always_runs_if_args_are_none():
     @idom.component
     @component_hook.capture
     def ComponentWithMemo():
-        value = idom.hooks.use_memo(lambda: next(iter_values))
+        value = idom.hooks.use_memo(lambda: next(iter_values), args=None)
         used_values.append(value)
         return idom.html.div()
 
@@ -860,3 +860,35 @@ def bad_callback():
 
     first_log_line = next(iter(caplog.records)).msg.split("\n", 1)[0]
     assert re.match(f"Failed to schedule render via {bad_callback}", first_log_line)
+
+
+async def test_use_effect_automatically_infers_closure_args():
+    set_count = idom.Ref()
+    did_effect = asyncio.Event()
+
+    @idom.component
+    def CounterWithEffect():
+        count, set_count.current = idom.hooks.use_state(0)
+
+        @idom.hooks.use_effect
+        def some_effect_that_uses_count():
+            """should automatically trigger on count change"""
+            count  # use count in this closure
+            did_effect.set()
+
+        return idom.html.div()
+
+    with idom.Layout(CounterWithEffect()) as layout:
+        await layout.render()
+        await did_effect.wait()
+        did_effect.clear()
+
+        for i in range(1, 3):
+            set_count.current(i)
+            await layout.render()
+            await did_effect.wait()
+            did_effect.clear()
+
+
+def test_use_memo_automatically_infers_closure_args():
+    ...