From 387bacff734078267183e72d05496fc3b271efc8 Mon Sep 17 00:00:00 2001
From: Lucas <lucasrhofer@gmail.com>
Date: Sat, 25 Feb 2023 18:13:02 +0000
Subject: [PATCH] Add multiple load events for a page (#596)

---
 pynecone/app.py                           |  6 +++---
 pynecone/middleware/hydrate_middleware.py | 22 ++++++++++++++++++----
 pynecone/route.py                         |  6 +++---
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/pynecone/app.py b/pynecone/app.py
index f831091c1..d2d5cc22d 100644
--- a/pynecone/app.py
+++ b/pynecone/app.py
@@ -53,7 +53,7 @@ class App(Base):
     middleware: List[Middleware] = []
 
     # Event handlers to trigger when a page loads.
-    load_events: Dict[str, EventHandler] = {}
+    load_events: Dict[str, Union[EventHandler, List[EventHandler]]] = {}
 
     def __init__(self, *args, **kwargs):
         """Initialize the app.
@@ -197,7 +197,7 @@ class App(Base):
         title: str = constants.DEFAULT_TITLE,
         description: str = constants.DEFAULT_DESCRIPTION,
         image=constants.DEFAULT_IMAGE,
-        on_load: Optional[EventHandler] = None,
+        on_load: Optional[Union[EventHandler, List[EventHandler]]] = None,
         path: Optional[str] = None,
         meta: List[Dict] = constants.DEFAULT_META_LIST,
     ):
@@ -213,7 +213,7 @@ class App(Base):
             title: The title of the page.
             description: The description of the page.
             image: The image to display on the page.
-            on_load: The event handler that will be called each time the page load.
+            on_load: The event handler(s) that will be called each time the page load.
             meta: The metadata of the page.
         """
         if path is not None:
diff --git a/pynecone/middleware/hydrate_middleware.py b/pynecone/middleware/hydrate_middleware.py
index 5d933e5c0..a02da7416 100644
--- a/pynecone/middleware/hydrate_middleware.py
+++ b/pynecone/middleware/hydrate_middleware.py
@@ -4,7 +4,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Optional
 
 from pynecone import constants, utils
-from pynecone.event import Event
+from pynecone.event import Event, EventHandler
 from pynecone.middleware.middleware import Middleware
 from pynecone.state import Delta, State
 
@@ -34,8 +34,22 @@ class HydrateMiddleware(Middleware):
                 load_event = app.load_events.get(route.lstrip("/"))
             else:
                 load_event = None
+
             if load_event:
-                substate_path = utils.format_event_handler(load_event).split(".")
-                ex_state = state.get_substate(substate_path[:-1])
-                load_event.fn(ex_state)
+                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)
             return utils.format_state({state.get_name(): state.dict()})
+
+    def execute_load_event(self, state: State, load_event: EventHandler) -> None:
+        """Execute single load event.
+
+        Args:
+            state: The client state.
+            load_event: A single load event to execute.
+        """
+        substate_path = utils.format_event_handler(load_event).split(".")
+        ex_state = state.get_substate(substate_path[:-1])
+        load_event.fn(ex_state)
diff --git a/pynecone/route.py b/pynecone/route.py
index f6e43b934..cfed25a87 100644
--- a/pynecone/route.py
+++ b/pynecone/route.py
@@ -1,6 +1,6 @@
 """The route decorator and associated variables."""
 
-from typing import Optional
+from typing import List, Optional, Union
 
 from pynecone.event import EventHandler
 
@@ -12,7 +12,7 @@ def route(
     title: Optional[str] = None,
     image: Optional[str] = None,
     description: Optional[str] = None,
-    on_load: Optional[EventHandler] = None,
+    on_load: Optional[Union[EventHandler, List[EventHandler]]] = None,
 ):
     """Decorate a function as a page.
 
@@ -28,7 +28,7 @@ def route(
         title: The title of the page.
         image: The favicon of the page.
         description: The description of the page
-        on_load: The event handler called when the page load.
+        on_load: The event handler(s) called when the page load.
 
     Returns:
         The decorated function.