Move initial state to separate file (#1599)

integration/test_dynamic_routes.py

@@ -145,11 +145,8 @@ def test_on_load_navigate(dynamic_route: AppHarness, driver):
     # look up the backend state and assert that `on_load` was called for all
     # navigation events
     backend_state = dynamic_route.app_instance.state_manager.states[token]
-    # TODO: navigating to dynamic page initially fires hydrate twice
-    # because the new page re-initializes `useEventLoop`, with the same hydrate event
-    # but routeChangeComplete also still fires.
     time.sleep(0.2)
-    assert backend_state.order[-10:] == [str(ix) for ix in range(10)]
+    assert backend_state.order == [str(ix) for ix in range(10)]
 
 
 def test_on_load_navigate_non_dynamic(dynamic_route: AppHarness, driver):

reflex/.templates/jinja/web/pages/index.js.jinja2

@@ -8,6 +8,7 @@
 
 {% block export %}
 export default function Component() {
+  const {{state_name}} = useContext(StateContext)
   const {{const.router}} = useRouter()
   const { {{const.color_mode}}, {{const.toggle_color_mode}} } = {{const.use_color_mode}}()
   const focusRef = useRef();
@@ -19,10 +20,7 @@ export default function Component() {
   }))
 
   // Main event loop.
-  const [{{state_name}}, Event, notConnected] = useEventLoop(
-    {{initial_state|json_dumps}},
-    [E('{{state_name}}.{{const.hydrate}}', {})],
-  )
+  const [Event, notConnected] = useContext(EventLoopContext)
 
   // Set focus to the specified element.
   useEffect(() => {
@@ -31,7 +29,6 @@ export default function Component() {
     }
   })
 
-  {% if is_dynamic %}
   // Route after the initial page hydration.
   useEffect(() => {
     const change_complete = () => Event([E('{{state_name}}.{{const.hydrate}}', {})])
@@ -40,7 +37,6 @@ export default function Component() {
       {{const.router}}.events.off('routeChangeComplete', change_complete)
     }
   }, [{{const.router}}])
-  {% endif %}
 
   {% for hook in hooks %}
   {{ hook }}

reflex/.templates/jinja/web/utils/context.js.jinja2

@@ -0,0 +1,7 @@
+import { createContext } from "react"
+import { E } from "/utils/state.js"
+
+export const initialState = {{ initial_state|json_dumps }}
+export const initialEvents = [E('{{state_name}}.{{const.hydrate}}', {})]
+export const StateContext = createContext(null);
+export const EventLoopContext = createContext(null);

reflex/.templates/web/pages/_app.js

@@ -1,6 +1,8 @@
 import { ChakraProvider, extendTheme } from "@chakra-ui/react";
 import { Global, css } from "@emotion/react";
 import theme from "/utils/theme";
+import { initialEvents, initialState, StateContext, EventLoopContext } from "/utils/context.js";
+import { useEventLoop } from "utils/state";
 
 import '../styles/tailwind.css'
 
@@ -12,11 +14,27 @@ const GlobalStyles = css`
   }
 `;
 
+function EventLoopProvider({ children }) {
+  const [state, Event, notConnected] = useEventLoop(
+    initialState,
+    initialEvents,
+  )
+  return (
+    <EventLoopContext.Provider value={[Event, notConnected]}>
+      <StateContext.Provider value={state}>
+        {children}
+      </StateContext.Provider>
+    </EventLoopContext.Provider>
+  )
+}
+
 function MyApp({ Component, pageProps }) {
   return (
     <ChakraProvider theme={extendTheme(theme)}>
       <Global styles={GlobalStyles} />
-      <Component {...pageProps} />
+      <EventLoopProvider>
+        <Component {...pageProps} />
+      </EventLoopProvider>
     </ChakraProvider>
   );
 }

reflex/app.py

@@ -517,6 +517,9 @@ class App(Base):
         # Compile the theme.
         compile_results.append(compiler.compile_theme(self.style))
 
+        # Compile the contexts.
+        compile_results.append(compiler.compile_contexts(self.state))
+
         # Compile the Tailwind config.
         if config.tailwind is not None:
             config.tailwind["content"] = config.tailwind.get(

reflex/compiler/compiler.py

@@ -6,7 +6,6 @@ from typing import List, Set, Tuple, Type
 from reflex import constants
 from reflex.compiler import templates, utils
 from reflex.components.component import Component, ComponentStyle, CustomComponent
-from reflex.route import get_route_args
 from reflex.state import State
 from reflex.utils import imports
 from reflex.vars import ImportVar
@@ -18,6 +17,7 @@ DEFAULT_IMPORTS: imports.ImportDict = {
         ImportVar(tag="useEffect"),
         ImportVar(tag="useRef"),
         ImportVar(tag="useState"),
+        ImportVar(tag="useContext"),
     },
     "next/router": {ImportVar(tag="useRouter")},
     f"/{constants.STATE_PATH}": {
@@ -31,6 +31,10 @@ DEFAULT_IMPORTS: imports.ImportDict = {
         ImportVar(tag="getAllLocalStorageItems"),
         ImportVar(tag="useEventLoop"),
     },
+    "/utils/context.js": {
+        ImportVar(tag="EventLoopContext"),
+        ImportVar(tag="StateContext"),
+    },
     "": {ImportVar(tag="focus-visible/dist/focus-visible")},
     "@chakra-ui/react": {
         ImportVar(tag=constants.USE_COLOR_MODE),
@@ -67,11 +71,25 @@ def _compile_theme(theme: dict) -> str:
     return templates.THEME.render(theme=theme)
 
 
+def _compile_contexts(state: Type[State]) -> str:
+    """Compile the initial state and contexts.
+
+    Args:
+        state: The app state.
+
+    Returns:
+        The compiled context file.
+    """
+    return templates.CONTEXT.render(
+        initial_state=utils.compile_state(state),
+        state_name=state.get_name(),
+    )
+
+
 def _compile_page(
     component: Component,
     state: Type[State],
     connect_error_component,
-    is_dynamic: bool,
 ) -> str:
     """Compile the component given the app state.
 
@@ -79,7 +97,6 @@ def _compile_page(
         component: The component to compile.
         state: The app state.
         connect_error_component: The component to render on sever connection error.
-        is_dynamic: if True, include route change re-hydration logic
 
     Returns:
         The compiled component.
@@ -93,13 +110,11 @@ def _compile_page(
     return templates.PAGE.render(
         imports=imports,
         custom_codes=component.get_custom_code(),
-        initial_state=utils.compile_state(state),
         state_name=state.get_name(),
         hooks=component.get_hooks(),
         render=component.render(),
         transports=constants.Transports.POLLING_WEBSOCKET.get_transports(),
         err_comp=connect_error_component.render() if connect_error_component else None,
-        is_dynamic=is_dynamic,
     )
 
 
@@ -186,6 +201,23 @@ def compile_theme(style: ComponentStyle) -> Tuple[str, str]:
     return output_path, code
 
 
+def compile_contexts(
+    state: Type[State],
+) -> Tuple[str, str]:
+    """Compile the initial state / context.
+
+    Args:
+        state: The app state.
+
+    Returns:
+        The path and code of the compiled context.
+    """
+    # Get the path for the output file.
+    output_path = utils.get_context_path()
+
+    return output_path, _compile_contexts(state)
+
+
 def compile_page(
     path: str,
     component: Component,
@@ -208,7 +240,9 @@ def compile_page(
 
     # Add the style to the component.
     code = _compile_page(
-        component, state, connect_error_component, is_dynamic=bool(get_route_args(path))
+        component,
+        state,
+        connect_error_component,
     )
     return output_path, code
 

reflex/compiler/templates.py

@@ -65,6 +65,9 @@ DOCUMENT_ROOT = get_template("web/pages/_document.js.jinja2")
 # Template for the theme file.
 THEME = get_template("web/utils/theme.js.jinja2")
 
+# Template for the context file.
+CONTEXT = get_template("web/utils/context.js.jinja2")
+
 # Template for Tailwind config.
 TAILWIND_CONFIG = get_template("web/tailwind.config.js.jinja2")
 

reflex/compiler/utils.py

@@ -236,6 +236,15 @@ def get_theme_path() -> str:
     return os.path.join(constants.WEB_UTILS_DIR, constants.THEME + constants.JS_EXT)
 
 
+def get_context_path() -> str:
+    """Get the path of the context / initial state file.
+
+    Returns:
+        The path of the context module.
+    """
+    return os.path.join(constants.WEB_UTILS_DIR, "context" + constants.JS_EXT)
+
+
 def get_components_path() -> str:
     """Get the path of the compiled components.