From 1029d3483adb7174db532005ea8b4f24f76f0c54 Mon Sep 17 00:00:00 2001
From: Nikhil Rao <nikhil@pynecone.io>
Date: Mon, 24 Apr 2023 17:54:16 -0700
Subject: [PATCH] Fix substate event handler conversions (#861)

---
 pynecone/app.py                             |  4 ----
 pynecone/state.py                           | 25 ++++++---------------
 tests/middleware/test_hydrate_middleware.py |  2 +-
 3 files changed, 8 insertions(+), 23 deletions(-)

diff --git a/pynecone/app.py b/pynecone/app.py
index 67b86ee94..94942e035 100644
--- a/pynecone/app.py
+++ b/pynecone/app.py
@@ -397,7 +397,6 @@ class App(Base):
         compiler.compile_theme(self.style)
 
         # Compile the pages.
-        self.state.set_handlers()
         custom_components = set()
         for route, component in self.pages.items():
             component.add_style(self.style)
@@ -409,9 +408,6 @@ class App(Base):
         # Compile the custom components.
         compiler.compile_components(custom_components)
 
-        # To support calling event handlers from other handlers.
-        self.state.convert_handlers_to_fns()
-
 
 async def process(
     app: App, event: Event, sid: str, headers: Dict, client_ip: str
diff --git a/pynecone/state.py b/pynecone/state.py
index a0b3c379d..da67fb802 100644
--- a/pynecone/state.py
+++ b/pynecone/state.py
@@ -21,6 +21,7 @@ from typing import (
 )
 
 import cloudpickle
+import pydantic
 from redis import Redis
 
 from pynecone import constants
@@ -32,7 +33,7 @@ from pynecone.var import BaseVar, ComputedVar, PCDict, PCList, Var
 Delta = Dict[str, Any]
 
 
-class State(Base, ABC):
+class State(Base, ABC, extra=pydantic.Extra.allow):
     """The state of the app."""
 
     # A map from the var name to the var.
@@ -87,6 +88,11 @@ class State(Base, ABC):
         for substate in self.get_substates():
             self.substates[substate.get_name()] = substate().set(parent_state=self)
 
+        # Convert the event handlers to functions.
+        for name, event_handler in self.event_handlers.items():
+            setattr(self, name, event_handler.fn)
+
+        # Initialize the mutable fields.
         self._init_mutable_fields()
 
     def _init_mutable_fields(self):
@@ -189,23 +195,6 @@ class State(Base, ABC):
             if not name.startswith("_") and isinstance(fn, Callable)
         }
         cls.event_handlers = {name: EventHandler(fn=fn) for name, fn in events.items()}
-
-        cls.set_handlers()
-
-    @classmethod
-    def convert_handlers_to_fns(cls):
-        """Convert the event handlers to functions.
-
-        This is done so the state functions can be called as normal functions during runtime.
-        """
-        for name, event_handler in cls.event_handlers.items():
-            setattr(cls, name, event_handler.fn)
-        for substate in cls.get_substates():
-            substate.convert_handlers_to_fns()
-
-    @classmethod
-    def set_handlers(cls):
-        """Set the state class handlers."""
         for name, event_handler in cls.event_handlers.items():
             setattr(cls, name, event_handler)
 
diff --git a/tests/middleware/test_hydrate_middleware.py b/tests/middleware/test_hydrate_middleware.py
index 948b1f06a..4ed0a1082 100644
--- a/tests/middleware/test_hydrate_middleware.py
+++ b/tests/middleware/test_hydrate_middleware.py
@@ -30,7 +30,7 @@ class TestState2(State):
             Chain of EventHandlers
         """
         self.num += 1
-        return [self.change_name()]
+        return self.change_name
 
     def change_name(self):
         """Test handler to change name."""