fix context template

Author: adhami3310

reflex/compiler/compiler.py

@@ -130,17 +130,17 @@ def _compile_contexts(state: type[BaseState] | None, theme: Component | None) ->
         appearance = LiteralVar.create(SYSTEM_COLOR_MODE)
 
     return (
-        templates.CONTEXT.render(
+        templates.context_template(
             initial_state=utils.compile_state(state),
             state_name=state.get_name(),
             client_storage=utils.compile_client_storage(state),
             is_dev_mode=not is_prod_mode(),
-            default_color_mode=appearance,
+            default_color_mode=str(appearance),
         )
         if state
-        else templates.CONTEXT.render(
+        else templates.context_template(
             is_dev_mode=not is_prod_mode(),
-            default_color_mode=appearance,
+            default_color_mode=str(appearance),
         )
     )
 

reflex/compiler/templates.py

@@ -351,38 +351,161 @@ def theme_template(theme: str):
     return f"""export default {theme}"""
 
 
-def _context_template(**kwargs):
+def context_template(
+    *,
+    is_dev_mode: bool,
+    default_color_mode: str,
+    initial_state: dict[str, Any] | None = None,
+    state_name: str | None = None,
+    client_storage: dict[str, dict[str, dict[str, Any]]] | None = None,
+):
     """Template for the context file.
 
     Args:
-        **kwargs: Template context variables including initial_state, state_name, etc.
+        initial_state: The initial state for the context.
+        state_name: The name of the state.
+        client_storage: The client storage for the context.
+        is_dev_mode: Whether the app is in development mode.
+        default_color_mode: The default color mode for the context.
 
     Returns:
         Rendered context file content as string.
     """
-    initial_state = kwargs.get("initial_state", "null")
-    state_name = kwargs.get("state_name", "")
-    client_storage = kwargs.get("client_storage", "")
-    is_dev_mode = kwargs.get("is_dev_mode", False)
-    default_color_mode = kwargs.get("default_color_mode", "")
+    initial_state = initial_state or {}
+    state_contexts_str = "".join(
+        [f"{state_name}: createContext(null)," for state_name in initial_state]
+    )
+
+    state_str = (
+        rf"""
+export const state_name = "{state_name}"
+
+export const exception_state_name = "{constants.CompileVars.FRONTEND_EXCEPTION_STATE_FULL}"
+
+// These events are triggered on initial load and each page navigation.
+export const onLoadInternalEvent = () => {{
+    const internal_events = [];
+
+    // Get tracked cookie and local storage vars to send to the backend.
+    const client_storage_vars = hydrateClientStorage(clientStorage);
+    // But only send the vars if any are actually set in the browser.
+    if (client_storage_vars && Object.keys(client_storage_vars).length !== 0) {{
+        internal_events.push(
+            Event(
+                '{state_name}.{constants.CompileVars.UPDATE_VARS_INTERNAL}',
+                {{vars: client_storage_vars}},
+            ),
+        );
+    }}
+
+    // `on_load_internal` triggers the correct on_load event(s) for the current page.
+    // If the page does not define any on_load event, this will just set `is_hydrated = true`.
+    internal_events.push(Event('{state_name}.{constants.CompileVars.ON_LOAD_INTERNAL}'));
+
+    return internal_events;
+}}
+
+// The following events are sent when the websocket connects or reconnects.
+export const initialEvents = () => [
+    Event('{state_name}.{constants.CompileVars.HYDRATE}'),
+    ...onLoadInternalEvent()
+]
+    """
+        if state_name
+        else """
+export const state_name = undefined
 
-    return f"""import {{ createContext, useContext, useMemo, useReducer, useState }} from "react"
+export const exception_state_name = undefined
 
-export const StateContexts = {{
-  {state_name or ""}: createContext(null),
+export const onLoadInternalEvent = () => []
+
+export const initialEvents = () => []
+"""
+    )
+
+    state_reducer_str = "\n".join(
+        rf'const [{state_name}, dispatch_{state_name}] = useReducer(applyDelta, initialState["{state_name}"])'
+        for state_name in initial_state
+    )
+
+    create_state_contexts_str = "\n".join(
+        rf"createElement(StateContexts.{state_name},{{value: {state_name}}},"
+        for state_name in initial_state
+    )
+
+    dispatchers_str = "\n".join(
+        f'"{state_name}": dispatch_{state_name},' for state_name in initial_state
+    )
+
+    return rf"""import {{ createContext, useContext, useMemo, useReducer, useState, createElement, useEffect }} from "react"
+import {{ applyDelta, Event, hydrateClientStorage, useEventLoop, refs }} from "$/utils/state"
+import {{ jsx }} from "@emotion/react";
+
+export const initialState = {"{}" if initial_state else json_dumps(initial_state)}
+
+export const defaultColorMode = { default_color_mode }
+export const ColorModeContext = createContext(null);
+export const UploadFilesContext = createContext(null);
+export const DispatchContext = createContext(null);
+export const StateContexts = {{{state_contexts_str}}};
+export const EventLoopContext = createContext(null);
+export const clientStorage = {"{}" if client_storage is None else json_dumps(client_storage)}
+
+{state_str}
+
+export const isDevMode = {json.dumps(is_dev_mode)};
+
+export function UploadFilesProvider({{ children }}) {{
+  const [filesById, setFilesById] = useState({{}})
+  refs["__clear_selected_files"] = (id) => setFilesById(filesById => {{
+    const newFilesById = {{...filesById}}
+    delete newFilesById[id]
+    return newFilesById
+  }})
+  return createElement(
+    UploadFilesContext.Provider,
+    {{ value: [filesById, setFilesById] }},
+    children
+  );
 }}
 
-export const EventContexts = {{
-  {state_name or ""}: createContext(null),
+export function ClientSide(component) {{
+  return ({{ children, ...props }}) => {{
+    const [Component, setComponent] = useState(null);
+    useEffect(() => {{
+      setComponent(component);
+    }}, []);
+    return Component ? jsx(Component, props, children) : null;
+  }};
 }}
 
-export const initialState = {initial_state}
-export const clientStorage = {client_storage}
-export const isDevMode = {str(is_dev_mode).lower()}
-export const colorModeManager = {{
-  get: () => {default_color_mode},
-  set: () => {{}},
-  type: "localStorage"
+export function EventLoopProvider({{ children }}) {{
+  const dispatch = useContext(DispatchContext)
+  const [addEvents, connectErrors] = useEventLoop(
+    dispatch,
+    initialEvents,
+    clientStorage,
+  )
+  return createElement(
+    EventLoopContext.Provider,
+    {{ value: [addEvents, connectErrors] }},
+    children
+  );
+}}
+
+export function StateProvider({{ children }}) {{
+  {state_reducer_str}
+  const dispatchers = useMemo(() => {{
+    return {{
+      {dispatchers_str}
+    }}
+  }}, [])
+
+  return (
+    {create_state_contexts_str}
+    createElement(DispatchContext, {{value: dispatchers}}, children)
+    {")" * len(initial_state)}
+  )
 }}"""
 
 
@@ -698,9 +821,6 @@ def render(self, **kwargs) -> str:
         return self.func(**kwargs)
 
 
-# Template for the context file.
-CONTEXT = TemplateFunction(_context_template)
-
 # Template to render a component tag.
 COMPONENT = TemplateFunction(_component_template)
 

reflex/compiler/utils.py

@@ -245,30 +245,29 @@ def _compile_client_storage_field(
 
 def _compile_client_storage_recursive(
     state: type[BaseState],
-) -> tuple[dict[str, dict], dict[str, dict], dict[str, dict]]:
+) -> tuple[
+    dict[str, dict[str, Any]], dict[str, dict[str, Any]], dict[str, dict[str, Any]]
+]:
     """Compile the client-side storage for the given state recursively.
 
     Args:
         state: The app state object.
 
     Returns:
-        A tuple of the compiled client-side storage info:
-            (
-                cookies: dict[str, dict],
-                local_storage: dict[str, dict[str, str]]
-                session_storage: dict[str, dict[str, str]]
-            ).
+        A tuple of the compiled client-side storage info: (cookies, local_storage, session_storage).
     """
-    cookies = {}
-    local_storage = {}
-    session_storage = {}
+    cookies: dict[str, dict[str, Any]] = {}
+    local_storage: dict[str, dict[str, Any]] = {}
+    session_storage: dict[str, dict[str, Any]] = {}
     state_name = state.get_full_name()
     for name, field in state.__fields__.items():
         if name in state.inherited_vars:
             # only include vars defined in this state
             continue
         state_key = f"{state_name}.{name}" + FIELD_MARKER
         field_type, options = _compile_client_storage_field(field)
+        if field_type is None or options is None:
+            continue
         if field_type is Cookie:
             cookies[state_key] = options
         elif field_type is LocalStorage:
@@ -289,7 +288,9 @@ def _compile_client_storage_recursive(
     return cookies, local_storage, session_storage
 
 
-def compile_client_storage(state: type[BaseState]) -> dict[str, dict]:
+def compile_client_storage(
+    state: type[BaseState],
+) -> dict[str, dict[str, dict[str, Any]]]:
     """Compile the client-side storage for the given state.
 
     Args: