allow arguments to be passed to rx memo event handlers (#5021)

* allow arguments to be passed

* add test case for memo components accepting event handlers

* AppHarness: quit browsers early to avoid lingering events

* fix ruff

---------

Co-authored-by: Masen Furer <[email protected]>

Author: adhami3310

reflex/compiler/utils.py

@@ -318,7 +318,7 @@ def compile_custom_component(
     }
 
     # Concatenate the props.
-    props = [prop._js_expr for prop in component.get_prop_vars()]
+    props = list(component.props)
 
     # Compile the component.
     return (

reflex/components/component.py

@@ -49,12 +49,13 @@
 from reflex.constants.compiler import SpecialAttributes
 from reflex.constants.state import FRONTEND_EVENT_STATE
 from reflex.event import (
-    EventActionsMixin,
     EventCallback,
     EventChain,
     EventHandler,
     EventSpec,
     no_args_event_spec,
+    parse_args_spec,
+    run_script,
 )
 from reflex.style import Style, format_as_emotion
 from reflex.utils import console, format, imports, types
@@ -66,7 +67,7 @@
     Var,
     cached_property_no_lock,
 )
-from reflex.vars.function import ArgsFunctionOperation, FunctionStringVar
+from reflex.vars.function import ArgsFunctionOperation, FunctionStringVar, FunctionVar
 from reflex.vars.number import ternary_operation
 from reflex.vars.object import ObjectVar
 from reflex.vars.sequence import LiteralArrayVar
@@ -1900,7 +1901,44 @@ def _get_all_custom_components(
 
         return custom_components
 
-    def get_prop_vars(self) -> List[Var]:
+    @staticmethod
+    def _get_event_spec_from_args_spec(name: str, event: EventChain) -> Callable:
+        """Get the event spec from the args spec.
+
+        Args:
+            name: The name of the event
+            event: The args spec.
+
+        Returns:
+            The event spec.
+        """
+
+        def fn(*args):
+            return run_script(Var(name).to(FunctionVar).call(*args))
+
+        if event.args_spec:
+            arg_spec = (
+                event.args_spec
+                if not isinstance(event.args_spec, Sequence)
+                else event.args_spec[0]
+            )
+            names = inspect.getfullargspec(arg_spec).args
+            fn.__signature__ = inspect.Signature(  # pyright: ignore[reportFunctionMemberAccess]
+                parameters=[
+                    inspect.Parameter(
+                        name=name,
+                        kind=inspect.Parameter.POSITIONAL_ONLY,
+                        annotation=arg._var_type,
+                    )
+                    for name, arg in zip(
+                        names, parse_args_spec(event.args_spec), strict=True
+                    )
+                ]
+            )
+
+        return fn
+
+    def get_prop_vars(self) -> List[Var | Callable]:
         """Get the prop vars.
 
         Returns:
@@ -1909,16 +1947,10 @@ def get_prop_vars(self) -> List[Var]:
         return [
             Var(
                 _js_expr=name,
-                _var_type=(
-                    prop._var_type
-                    if isinstance(prop, Var)
-                    else (
-                        type(prop)
-                        if not isinstance(prop, EventActionsMixin)
-                        else EventChain
-                    )
-                ),
+                _var_type=(prop._var_type if isinstance(prop, Var) else type(prop)),
             ).guess_type()
+            if isinstance(prop, Var) or not isinstance(prop, EventChain)
+            else CustomComponent._get_event_spec_from_args_spec(name, prop)
             for name, prop in self.props.items()
         ]
 

reflex/testing.py

@@ -462,6 +462,10 @@ def __enter__(self) -> "AppHarness":
 
     def stop(self) -> None:
         """Stop the frontend and backend servers."""
+        # Quit browsers first to avoid any lingering events being sent during shutdown.
+        for driver in self._frontends:
+            driver.quit()
+
         self._reload_state_module()
 
         if self.backend is not None:
@@ -492,8 +496,6 @@ def stop(self) -> None:
             self.backend_thread.join()
         if self.frontend_output_thread is not None:
             self.frontend_output_thread.join()
-        for driver in self._frontends:
-            driver.quit()
 
         # Cleanup decorated pages added during testing
         for page in self._decorated_pages:

tests/integration/test_memo.py

@@ -26,9 +26,32 @@ def foo_component(t: str):
     def foo_component2(t: str):
         return FooComponent.create(t, rx.Var("foo"))
 
+    class MemoState(rx.State):
+        last_value: str = ""
+
+        @rx.event
+        def set_last_value(self, value: str):
+            self.last_value = value
+
+    @rx.memo
+    def my_memoed_component(
+        some_value: str,
+        event: rx.EventHandler[rx.event.passthrough_event_spec(str)],
+    ) -> rx.Component:
+        return rx.vstack(
+            rx.button(some_value, id="memo-button", on_click=event(some_value)),
+            rx.input(id="memo-input", on_change=event),
+        )
+
     def index() -> rx.Component:
         return rx.vstack(
-            foo_component(t="foo"), foo_component2(t="bar"), id="memo-custom-code"
+            rx.vstack(
+                foo_component(t="foo"), foo_component2(t="bar"), id="memo-custom-code"
+            ),
+            rx.text(MemoState.last_value, id="memo-last-value"),
+            my_memoed_component(
+                some_value="memod_some_value", event=MemoState.set_last_value
+            ),
         )
 
     app = rx.App()
@@ -64,4 +87,21 @@ async def test_memo_app(memo_app: AppHarness):
 
     # check that the output matches
     memo_custom_code_stack = driver.find_element(By.ID, "memo-custom-code")
+    assert (
+        memo_app.poll_for_content(memo_custom_code_stack, exp_not_equal="")
+        == "foobarbarbar"
+    )
     assert memo_custom_code_stack.text == "foobarbarbar"
+
+    # click the button to trigger partial event application
+    button = driver.find_element(By.ID, "memo-button")
+    button.click()
+    last_value = driver.find_element(By.ID, "memo-last-value")
+    assert memo_app.poll_for_content(last_value, exp_not_equal="") == "memod_some_value"
+
+    # enter text to trigger passed argument to event handler
+    textbox = driver.find_element(By.ID, "memo-input")
+    textbox.send_keys("new_value")
+    assert memo_app._poll_for(
+        lambda: memo_app.poll_for_content(last_value) == "new_value"
+    )