Repository: reflex-dev/reflex
Pull request: Fix Event chaining in the on_load event handler return not working
This PR fixes a couple of issues with the way on-load events are handled. In this PR, the following should work now:

- event-chaining( handlers that return other handlers)
- list of on-load events (eg. `app.add_page('index', on_load = [State.first_handler, State.second_handler]`)
- Async handlers

fixes #610 , #603


Problem statement:
Event chaining in the on_load event handler return is not working.
**Describe the bug**
A clear and concise description of what the bug is.

Event chaining in the on_load event handler return is not working.

**To Reproduce**
Steps to reproduce the behavior:
 - Code/Link to Repo:

I implemented it like this

```python
class BaseState(pc.State):
    ...


class QueryState(BaseState):
    label: str

    def init(self):
        print("init!")
        return [
            self.reset_label,
        ]

    def reset_label(self):
        print("reset label!")
        self.label = ""

app = pc.App(state=state.BaseState)
app.add_page(query.index, route="/", on_load=QueryState.init)
```

When I ran the code, the result was as follows.

```
init!
```

The `reset_label()` method doesn't seem to have actually been executed.

**Expected behavior**
A clear and concise description of what you expected to happen.

The `reset_label()` method should be executed.
This means that the result of its execution should be

```
init!
reset label!
```

**Screenshots**
If applicable, add screenshots to help explain your problem.

** Specifics (please complete the following information):**
 - Python Version: 3.11.0
 - Pynecone Version: 0.1.18
 - OS: Mac OS Monterey 12.6
 - Browser (Optional): .
 
**Additional context**
Add any other context about the problem here.

Hints from issue discussion:
Currently event chains don't work in `on_load` since `on_load` events are handled in [HydrateMiddleware](https://github.com/pynecone-io/pynecone/blob/387bacff734078267183e72d05496fc3b271efc8/pynecone/middleware/hydrate_middleware.py). 

A different method for handling `on_load` events has been proposed to handle chaining, but is not yet implemented (see [#603](https://github.com/pynecone-io/pynecone/issues/603)).

As a work around in the meantime, you can use a list of event handlers in `on_load`:

```class BaseState(pc.State):
    ...


class QueryState(BaseState):
    label: str

    def init(self):
        print("init!")

    def reset_label(self):
        print("reset label!")
        self.label = ""

app = pc.App(state=state.BaseState)
app.add_page(query.index, route="/", 
                        on_load=[QueryState.init, 
                                 QueryState.reset_label]
                       )
```


I modified it like this

```
app.add_page(query.index, route="/", on_load=[QueryState.init, QueryState.reset_label])
```

However, we still get the following error.

```
AttributeError: 'list' object has no attribute 'fn'
```

> As a work around in the meantime, you can use a list of event handlers in on_load:

I'm using the most recent version, 0.1.18, but the list of event handlers for on_load doesn't seem to be implemented yet.
Okay. I'll consider it a bug and wait for it to be fixed.
Thanks for all your hard work.
The bug has been fixed since 0.1.18. 

If you clone the repo and then pip install it, it should work

Commit message: added async tests

Patch to review:
diff --git a/pynecone/app.py b/pynecone/app.py
index 8c359e30da1..0a3ef91d49d 100644
--- a/pynecone/app.py
+++ b/pynecone/app.py
@@ -148,7 +148,7 @@ def add_cors(self):
             allow_origins=["*"],
         )
 
-    def preprocess(self, state: State, event: Event) -> Optional[Delta]:
+    async def preprocess(self, state: State, event: Event) -> Optional[Delta]:
         """Preprocess the event.
 
         This is where middleware can modify the event before it is processed.
@@ -165,7 +165,7 @@ def preprocess(self, state: State, event: Event) -> Optional[Delta]:
             An optional state to return.
         """
         for middleware in self.middleware:
-            out = middleware.preprocess(app=self, state=state, event=event)
+            out = await middleware.preprocess(app=self, state=state, event=event)
             if out is not None:
                 return out
 
@@ -400,7 +400,7 @@ def compile(self, force_compile: bool = False):
 
 async def process(
     app: App, event: Event, sid: str, headers: Dict, client_ip: str
-) -> StateUpdate:
+) -> Union[StateUpdate, List[StateUpdate]]:
     """Process an event.
 
     Args:
@@ -411,7 +411,7 @@ async def process(
         client_ip: The client_ip.
 
     Returns:
-        The state update after processing the event.
+        The state update(s) after processing the event.
     """
     # Get the state for the session.
     state = app.state_manager.get_state(event.token)
@@ -430,21 +430,27 @@ async def process(
     state.router_data[constants.RouteVar.CLIENT_IP] = client_ip
 
     # Preprocess the event.
-    pre = app.preprocess(state, event)
-    if pre is not None:
+    pre = await app.preprocess(state, event)
+    if pre is not None and not isinstance(pre, List):
         return StateUpdate(delta=pre)
 
     # Apply the event to the state.
-    update = await state.process(event)
+    updates = pre if pre else await state.process(event)
     app.state_manager.set_state(event.token, state)
 
+    updates = updates if isinstance(updates, List) else [updates]
+
     # Postprocess the event.
-    post = app.postprocess(state, event, update.delta)
-    if post is not None:
-        return StateUpdate(delta=post)
+    post_list = []
+    for update in updates:
+        post = app.postprocess(state, event, update.delta)  # type: ignore
+        post_list.append(post) if post else None
+
+    if post_list:
+        return [StateUpdate(delta=post) for post in post_list]
 
     # Return the update.
-    return update
+    return updates
 
 
 async def ping() -> str:
@@ -582,7 +588,13 @@ async def on_event(self, sid, data):
         update = await process(self.app, event, sid, headers, client_ip)
 
         # Emit the event.
-        await self.emit(str(constants.SocketEvent.EVENT), update.json(), to=sid)
+        if isinstance(update, List):
+            for single_update in update:
+                await self.emit(
+                    str(constants.SocketEvent.EVENT), single_update.json(), to=sid
+                )
+        else:
+            await self.emit(str(constants.SocketEvent.EVENT), update.json(), to=sid)
 
     async def on_ping(self, sid):
         """Event for testing the API endpoint.
diff --git a/pynecone/components/typography/markdown.py b/pynecone/components/typography/markdown.py
index a2fa771fa08..8d50912d1fc 100644
--- a/pynecone/components/typography/markdown.py
+++ b/pynecone/components/typography/markdown.py
@@ -69,7 +69,7 @@ def _render(self):
                     "li": "{ListItem}",
                     "p": "{Text}",
                     "a": "{Link}",
-                    "code": """{({node, inline, className, children, ...props}) => 
+                    "code": """{({node, inline, className, children, ...props}) =>
                     {
         const match = (className || '').match(/language-(?<lang>.*)/);
         return !inline ? (
diff --git a/pynecone/middleware/hydrate_middleware.py b/pynecone/middleware/hydrate_middleware.py
index 4c2b254f17b..c7b101bad70 100644
--- a/pynecone/middleware/hydrate_middleware.py
+++ b/pynecone/middleware/hydrate_middleware.py
@@ -1,12 +1,12 @@
 """Middleware to hydrate the state."""
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 from pynecone import constants
 from pynecone.event import Event, EventHandler, get_hydrate_event
 from pynecone.middleware.middleware import Middleware
-from pynecone.state import Delta, State
+from pynecone.state import Delta, State, StateUpdate
 from pynecone.utils import format
 
 if TYPE_CHECKING:
@@ -16,7 +16,9 @@
 class HydrateMiddleware(Middleware):
     """Middleware to handle initial app hydration."""
 
-    def preprocess(self, app: App, state: State, event: Event) -> Optional[Delta]:
+    async def preprocess(
+        self, app: App, state: State, event: Event
+    ) -> Union[Optional[Delta], List[StateUpdate]]:
         """Preprocess the event.
 
         Args:
@@ -25,7 +27,7 @@ def preprocess(self, app: App, state: State, event: Event) -> Optional[Delta]:
             event: The event to preprocess.
 
         Returns:
-            An optional state to return.
+            An optional delta or list of state updates to return.
         """
         if event.name == get_hydrate_event(state):
             route = event.router_data.get(constants.RouteVar.PATH, "")
@@ -37,20 +39,42 @@ def preprocess(self, app: App, state: State, event: Event) -> Optional[Delta]:
                 load_event = None
 
             if load_event:
-                if isinstance(load_event, list):
-                    for single_event in load_event:
-                        self.execute_load_event(state, single_event)
-                else:
-                    self.execute_load_event(state, load_event)
+                if not isinstance(load_event, List):
+                    load_event = [load_event]
+                updates = []
+                for single_event in load_event:
+                    updates.append(
+                        await self.execute_load_event(
+                            state, single_event, event.token, event.payload
+                        )
+                    )
+                return updates
             return format.format_state({state.get_name(): state.dict()})
 
-    def execute_load_event(self, state: State, load_event: EventHandler) -> None:
+    async def execute_load_event(
+        self, state: State, load_event: EventHandler, token: str, payload: Dict
+    ) -> StateUpdate:
         """Execute single load event.
 
         Args:
             state: The client state.
             load_event: A single load event to execute.
+            token: client token
+            payload: the event payload
+
+        Returns:
+            A state Update.
+
+        Raises:
+            ValueError: If the state value is None.
         """
         substate_path = format.format_event_handler(load_event).split(".")
         ex_state = state.get_substate(substate_path[:-1])
-        load_event.fn(ex_state)
+        if ex_state:
+            return await state.process_event(
+                handler=load_event, state=ex_state, payload=payload, token=token
+            )
+        else:
+            raise ValueError(
+                "The value of state cannot be None when processing an on-load event."
+            )
diff --git a/pynecone/middleware/logging_middleware.py b/pynecone/middleware/logging_middleware.py
index e5e95399dab..480c41efac4 100644
--- a/pynecone/middleware/logging_middleware.py
+++ b/pynecone/middleware/logging_middleware.py
@@ -14,7 +14,7 @@
 class LoggingMiddleware(Middleware):
     """Middleware to log requests and responses."""
 
-    def preprocess(self, app: App, state: State, event: Event):
+    async def preprocess(self, app: App, state: State, event: Event):
         """Preprocess the event.
 
         Args:
diff --git a/pynecone/middleware/middleware.py b/pynecone/middleware/middleware.py
index 94fbc0cd4a3..873db1db4e8 100644
--- a/pynecone/middleware/middleware.py
+++ b/pynecone/middleware/middleware.py
@@ -15,7 +15,7 @@
 class Middleware(Base, ABC):
     """Middleware to preprocess and postprocess requests."""
 
-    def preprocess(self, app: App, state: State, event: Event) -> Optional[Delta]:
+    async def preprocess(self, app: App, state: State, event: Event) -> Optional[Delta]:
         """Preprocess the event.
 
         Args:
diff --git a/pynecone/pc.py b/pynecone/pc.py
index 7edc6bff911..ab3c19e2320 100644
--- a/pynecone/pc.py
+++ b/pynecone/pc.py
@@ -203,12 +203,12 @@ def export(
 
     if zipping:
         console.rule(
-            """Backend & Frontend compiled. See [green bold]backend.zip[/green bold] 
+            """Backend & Frontend compiled. See [green bold]backend.zip[/green bold]
             and [green bold]frontend.zip[/green bold]."""
         )
     else:
         console.rule(
-            """Backend & Frontend compiled. See [green bold]app[/green bold] 
+            """Backend & Frontend compiled. See [green bold]app[/green bold]
             and [green bold].web/_static[/green bold] directories."""
         )
 
diff --git a/pynecone/state.py b/pynecone/state.py
index 8dff76cfafb..2e5ba2b364a 100644
--- a/pynecone/state.py
+++ b/pynecone/state.py
@@ -546,13 +546,16 @@ def get_substate(self, path: Sequence[str]) -> Optional[State]:
         return self.substates[path[0]].get_substate(path[1:])
 
     async def process(self, event: Event) -> StateUpdate:
-        """Process an event.
+        """Obtain event info and process event.
 
         Args:
             event: The event to process.
 
         Returns:
             The state update after processing the event.
+
+        Raises:
+            ValueError: If the state value is None.
         """
         # Get the event handler.
         path = event.name.split(".")
@@ -560,23 +563,48 @@ async def process(self, event: Event) -> StateUpdate:
         substate = self.get_substate(path)
         handler = substate.event_handlers[name]  # type: ignore
 
-        # Process the event.
-        fn = functools.partial(handler.fn, substate)
+        if substate:
+            return await self.process_event(
+                handler=handler,
+                state=substate,
+                payload=event.payload,
+                token=event.token,
+            )
+        else:
+            raise ValueError(
+                "The value of state cannot be None when processing an event."
+            )
+
+    async def process_event(
+        self, handler: EventHandler, state: State, payload: Dict, token: str
+    ) -> StateUpdate:
+        """Process event.
+
+        Args:
+            handler: Eventhandler to process.
+            state: state to process the handler.
+            payload: the event payload.
+            token: client token.
+
+        Returns:
+            The state update after processing the event.
+        """
+        fn = functools.partial(handler.fn, state)
         try:
             if asyncio.iscoroutinefunction(fn.func):
-                events = await fn(**event.payload)
+                events = await fn(**payload)
             else:
-                events = fn(**event.payload)
+                events = fn(**payload)
         except Exception:
             error = traceback.format_exc()
             print(error)
             events = fix_events(
-                [window_alert("An error occurred. See logs for details.")], event.token
+                [window_alert("An error occurred. See logs for details.")], token
             )
             return StateUpdate(events=events)
 
         # Fix the returned events.
-        events = fix_events(events, event.token)
+        events = fix_events(events, token)
 
         # Get the delta after processing the event.
         delta = self.get_delta()
diff --git a/tests/components/datadisplay/test_datatable.py b/tests/components/datadisplay/test_datatable.py
index b2ed7f2a9cd..f733cc63c58 100644
--- a/tests/components/datadisplay/test_datatable.py
+++ b/tests/components/datadisplay/test_datatable.py
@@ -94,7 +94,6 @@ def test_computed_var_without_annotation(fixture, request, err_msg, is_data_fram
         is_data_frame: whether data field is a pandas dataframe.
     """
     with pytest.raises(ValueError) as err:
-
         if is_data_frame:
             data_table(data=request.getfixturevalue(fixture).data)
         else:
diff --git a/tests/middleware/__init__.py b/tests/middleware/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/middleware/conftest.py b/tests/middleware/conftest.py
new file mode 100644
index 00000000000..b0640693008
--- /dev/null
+++ b/tests/middleware/conftest.py
@@ -0,0 +1,34 @@
+import pytest
+
+from pynecone.event import Event
+
+
+def create_event(name):
+    return Event(
+        token="<token>",
+        name=name,
+        router_data={
+            "pathname": "/",
+            "query": {},
+            "token": "<token>",
+            "sid": "<sid>",
+            "headers": {},
+            "ip": "127.0.0.1",
+        },
+        payload={},
+    )
+
+
+@pytest.fixture
+def event1():
+    return create_event("test_state.hydrate")
+
+
+@pytest.fixture
+def event2():
+    return create_event("test_state2.hydrate")
+
+
+@pytest.fixture
+def event3():
+    return create_event("test_state3.hydrate")
diff --git a/tests/middleware/test_hydrate_middleware.py b/tests/middleware/test_hydrate_middleware.py
new file mode 100644
index 00000000000..948b1f06a90
--- /dev/null
+++ b/tests/middleware/test_hydrate_middleware.py
@@ -0,0 +1,96 @@
+from typing import List
+
+import pytest
+
+from pynecone.app import App
+from pynecone.middleware.hydrate_middleware import HydrateMiddleware
+from pynecone.state import State
+
+
+class TestState(State):
+    """A test state with no return in handler."""
+
+    num: int = 0
+
+    def test_handler(self):
+        """Test handler."""
+        self.num += 1
+
+
+class TestState2(State):
+    """A test state with return in handler."""
+
+    num: int = 0
+    name: str
+
+    def test_handler(self):
+        """Test handler that calls another handler.
+
+        Returns:
+            Chain of EventHandlers
+        """
+        self.num += 1
+        return [self.change_name()]
+
+    def change_name(self):
+        """Test handler to change name."""
+        self.name = "random"
+
+
+class TestState3(State):
+    """A test state with async handler."""
+
+    num: int = 0
+
+    async def test_handler(self):
+        """Test handler."""
+        self.num += 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "state, expected, event_fixture",
+    [
+        (TestState, {"test_state": {"num": 1}}, "event1"),
+        (TestState2, {"test_state2": {"num": 1}}, "event2"),
+        (TestState3, {"test_state3": {"num": 1}}, "event3"),
+    ],
+)
+async def test_preprocess(state, request, event_fixture, expected):
+    """Test that a state hydrate event is processed correctly.
+
+    Args:
+        state: state to process event
+        request: pytest fixture request
+        event_fixture: The event fixture(an Event)
+        expected: expected delta
+    """
+    app = App(state=state, load_events={"index": state.test_handler})
+
+    hydrate_middleware = HydrateMiddleware()
+    result = await hydrate_middleware.preprocess(
+        app=app, event=request.getfixturevalue(event_fixture), state=state()
+    )
+    assert isinstance(result, List)
+    assert result[0].delta == expected
+
+
+@pytest.mark.asyncio
+async def test_preprocess_multiple_load_events(event1):
+    """Test that a state hydrate event for multiple on-load events is processed correctly.
+
+    Args:
+        event1: an Event.
+    """
+    app = App(
+        state=TestState,
+        load_events={"index": [TestState.test_handler, TestState.test_handler]},
+    )
+
+    hydrate_middleware = HydrateMiddleware()
+    result = await hydrate_middleware.preprocess(
+        app=app, event=event1, state=TestState()
+    )
+    assert isinstance(result, List)
+    assert result[0].delta == {"test_state": {"num": 1}}
+    assert result[1].delta == {"test_state": {"num": 2}}


Code review comments:
diff_hunk: @@ -37,20 +39,42 @@ def preprocess(self, app: App, state: State, event: Event) -> Optional[Delta]:
                 load_event = None
 
             if load_event:
-                if isinstance(load_event, list):
-                    for single_event in load_event:
-                        self.execute_load_event(state, single_event)
-                else:
-                    self.execute_load_event(state, load_event)
+                if not isinstance(load_event, List):
+                    load_event = [load_event]
+                updates = []
+                for single_event in load_event:
+                    updates.append(
+                        await self.execute_load_event(
+                            state, single_event, event.token, event.payload
+                        )
+                    )
+                return updates
             return format.format_state({state.get_name(): state.dict()})
 
-    def execute_load_event(self, state: State, load_event: EventHandler) -> None:
+    async def execute_load_event(
+        self, state: State, load_event: EventHandler, token: str, payload: Dict
+    ) -> StateUpdate:
         """Execute single load event.
 
         Args:
             state: The client state.
             load_event: A single load event to execute.
+            token: client token
+            payload: the event payload
+
+        Returns:
+            A state Update.
+
+        Raises:
+            ValueError: If the state value is None.
         """
         substate_path = format.format_event_handler(load_event).split(".")
         ex_state = state.get_substate(substate_path[:-1])
-        load_event.fn(ex_state)
+        if ex_state:
+            return await state.process_event(
+                handler=load_event, state=ex_state, payload=payload, token=token
+            )
+        else:
path: pynecone/middleware/hydrate_middleware.py
text: @user1:
Don't need the `else` since the first case returns

diff_hunk: @@ -546,37 +546,65 @@ def get_substate(self, path: Sequence[str]) -> Optional[State]:
         return self.substates[path[0]].get_substate(path[1:])
 
     async def process(self, event: Event) -> StateUpdate:
-        """Process an event.
+        """Obtain event info and process event.
 
         Args:
             event: The event to process.
 
         Returns:
             The state update after processing the event.
+
+        Raises:
+            ValueError: If the state value is None.
         """
         # Get the event handler.
         path = event.name.split(".")
         path, name = path[:-1], path[-1]
         substate = self.get_substate(path)
         handler = substate.event_handlers[name]  # type: ignore
 
-        # Process the event.
-        fn = functools.partial(handler.fn, substate)
+        if substate:
path: pynecone/state.py
text: @user1:
We can change this to

```python
if not substate:
    raise ValueError(...)

return ...

```

diff_hunk: @@ -37,20 +39,42 @@ def preprocess(self, app: App, state: State, event: Event) -> Optional[Delta]:
                 load_event = None
 
             if load_event:
-                if isinstance(load_event, list):
-                    for single_event in load_event:
-                        self.execute_load_event(state, single_event)
-                else:
-                    self.execute_load_event(state, load_event)
+                if not isinstance(load_event, List):
+                    load_event = [load_event]
+                updates = []
+                for single_event in load_event:
+                    updates.append(
+                        await self.execute_load_event(
+                            state, single_event, event.token, event.payload
+                        )
+                    )
+                return updates
             return format.format_state({state.get_name(): state.dict()})
path: pynecone/middleware/hydrate_middleware.py
text: @user1:
We can maybe just make this a `StateUpdate` object to simplify the API

diff_hunk: @@ -582,7 +588,13 @@ async def on_event(self, sid, data):
         update = await process(self.app, event, sid, headers, client_ip)
 
         # Emit the event.
-        await self.emit(str(constants.SocketEvent.EVENT), update.json(), to=sid)
+        if isinstance(update, List):
path: pynecone/app.py
text: @user1:
We can make sure update is always a List so we don't need this check

diff_hunk: @@ -430,21 +430,27 @@ async def process(
     state.router_data[constants.RouteVar.CLIENT_IP] = client_ip
 
     # Preprocess the event.
-    pre = app.preprocess(state, event)
-    if pre is not None:
+    pre = await app.preprocess(state, event)
+    if pre is not None and not isinstance(pre, List):
         return StateUpdate(delta=pre)
path: pynecone/app.py
text: @user1:
We can listify this here

diff_hunk: @@ -148,7 +148,7 @@ def add_cors(self):
             allow_origins=["*"],
         )
 
-    def preprocess(self, state: State, event: Event) -> Optional[Delta]:
+    async def preprocess(self, state: State, event: Event) -> Optional[Delta]:
path: pynecone/app.py
text: @user1:
Should we make `postprocess` async also for consistency?

@author:
yep, I agree