[REF-889] useContext per substate (#2149)

integration/test_var_operations.py

@@ -28,6 +28,7 @@ class VarOperationState(rx.State):
         str_var4: str = "a long string"
         dict1: dict = {1: 2}
         dict2: dict = {3: 4}
+        html_str: str = "<div>hello</div>"
 
     app = rx.App(state=VarOperationState)
 
@@ -522,6 +523,19 @@ def index():
             rx.text(VarOperationState.str_var4.split(" ").to_string(), id="str_split"),
             rx.text(VarOperationState.list3.join(""), id="list_join"),
             rx.text(VarOperationState.list3.join(","), id="list_join_comma"),
+            # Index from an op var
+            rx.text(
+                VarOperationState.list3[VarOperationState.int_var1 % 3],
+                id="list_index_mod",
+            ),
+            rx.html(
+                VarOperationState.html_str,
+                id="html_str",
+            ),
+            rx.highlight(
+                "second",
+                query=[VarOperationState.str_var2],
+            ),
             rx.text(rx.Var.range(2, 5).join(","), id="list_join_range1"),
             rx.text(rx.Var.range(2, 10, 2).join(","), id="list_join_range2"),
             rx.text(rx.Var.range(5, 0, -1).join(","), id="list_join_range3"),
@@ -713,7 +727,14 @@ def test_var_operations(driver, var_operations: AppHarness):
         ("dict_eq_dict", "false"),
         ("dict_neq_dict", "true"),
         ("dict_contains", "true"),
+        # index from an op var
+        ("list_index_mod", "second"),
+        # html component with var
+        ("html_str", "hello"),
     ]
 
     for tag, expected in tests:
         assert driver.find_element(By.ID, tag).text == expected
+
+    # Highlight component with var query (does not plumb ID)
+    assert driver.find_element(By.TAG_NAME, "mark").text == "second"

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

@@ -1,7 +1,7 @@
 {% extends "web/pages/base_page.js.jinja2" %}
 
 {% block declaration %}
-import { EventLoopProvider } from "/utils/context.js";
+import { EventLoopProvider, StateProvider } from "/utils/context.js";
 import { ThemeProvider } from 'next-themes'
 
 {% for custom_code in custom_codes %}
@@ -25,12 +25,14 @@ export default function MyApp({ Component, pageProps }) {
   return (
     <ThemeProvider defaultTheme="light" storageKey="chakra-ui-color-mode" attribute="class">
       <AppWrap>
-        <EventLoopProvider>
-          <Component {...pageProps} />
-        </EventLoopProvider>
+        <StateProvider>
+          <EventLoopProvider>
+            <Component {...pageProps} />
+          </EventLoopProvider>
+        </StateProvider>
       </AppWrap>
     </ThemeProvider>
   );
 }
 
-{% endblock %}
+{% endblock %}

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

@@ -8,32 +8,6 @@
 
 {% block export %}
 export default function Component() {
-{% if state_name %}
-  const {{state_name}} = useContext(StateContext)
-{% endif %}
-  const {{const.router}} = useRouter()
-  const [ {{const.color_mode}}, {{const.toggle_color_mode}} ] = useContext(ColorModeContext)
-  const focusRef = useRef();
-
-  // Main event loop.
-  const [addEvents, connectError] = useContext(EventLoopContext)
-
-  // Set focus to the specified element.
-  useEffect(() => {
-    if (focusRef.current) {
-      focusRef.current.focus();
-    }
-  })
-
-  // Route after the initial page hydration.
-  useEffect(() => {
-    const change_complete = () => addEvents(initialEvents())
-    {{const.router}}.events.on('routeChangeComplete', change_complete)
-    return () => {
-      {{const.router}}.events.off('routeChangeComplete', change_complete)
-    }
-  }, [{{const.router}}])
-
   {% for hook in hooks %}
   {{ hook }}
   {% endfor %}

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

@@ -1,5 +1,5 @@
-import { createContext, useState } from "react"
-import { Event, hydrateClientStorage, useEventLoop } from "/utils/state.js"
+import { createContext, useContext, useMemo, useReducer, useState } from "react"
+import { applyDelta, Event, hydrateClientStorage, useEventLoop } from "/utils/state.js"
 
 {% if initial_state %}
 export const initialState = {{ initial_state|json_dumps }}
@@ -8,7 +8,12 @@ export const initialState = {}
 {% endif %}
 
 export const ColorModeContext = createContext(null);
-export const StateContext = createContext(null);
+export const DispatchContext = createContext(null);
+export const StateContexts = {
+  {% for state_name in initial_state %}
+  {{state_name|var_name}}: createContext(null),
+  {% endfor %}
+}
 export const EventLoopContext = createContext(null);
 {% if client_storage %}
 export const clientStorage = {{ client_storage|json_dumps }}
@@ -27,16 +32,40 @@ export const initialEvents = () => []
 export const isDevMode = {{ is_dev_mode|json_dumps }}
 
 export function EventLoopProvider({ children }) {
-  const [state, addEvents, connectError] = useEventLoop(
-    initialState,
+  const dispatch = useContext(DispatchContext)
+  const [addEvents, connectError] = useEventLoop(
+    dispatch,
     initialEvents,
     clientStorage,
   )
   return (
     <EventLoopContext.Provider value={[addEvents, connectError]}>
-      <StateContext.Provider value={state}>
-        {children}
-      </StateContext.Provider>
+      {children}
     </EventLoopContext.Provider>
   )
-}
+}
+
+export function StateProvider({ children }) {
+  {% for state_name in initial_state %}
+  const [{{state_name|var_name}}, dispatch_{{state_name|var_name}}] = useReducer(applyDelta, initialState["{{state_name}}"])
+  {% endfor %}
+  const dispatchers = useMemo(() => {
+    return {
+      {% for state_name in initial_state %}
+      "{{state_name}}": dispatch_{{state_name|var_name}},
+      {% endfor %}
+    }
+  }, [])
+
+  return (
+    {% for state_name in initial_state %}
+    <StateContexts.{{state_name|var_name}}.Provider value={ {{state_name|var_name}} }>
+    {% endfor %}
+      <DispatchContext.Provider value={dispatchers}>
+        {children}
+      </DispatchContext.Provider>
+    {% for state_name in initial_state|reverse %}
+    </StateContexts.{{state_name|var_name}}.Provider>
+    {% endfor %}
+  )
+}

reflex/.templates/web/utils/state.js

@@ -6,7 +6,7 @@ import env from "env.json";
 import Cookies from "universal-cookie";
 import { useEffect, useReducer, useRef, useState } from "react";
 import Router, { useRouter } from "next/router";
-import { initialEvents } from "utils/context.js"
+import { initialEvents, initialState } from "utils/context.js"
 
 // Endpoint URLs.
 const EVENTURL = env.EVENT
@@ -100,37 +100,10 @@ export const getEventURL = () => {
  * @param delta The delta to apply.
  */
 export const applyDelta = (state, delta) => {
-  const new_state = { ...state }
-  for (const substate in delta) {
-    let s = new_state;
-    const path = substate.split(".").slice(1);
-    while (path.length > 0) {
-      s = s[path.shift()];
-    }
-    for (const key in delta[substate]) {
-      s[key] = delta[substate][key];
-    }
-  }
-  return new_state
+  return { ...state, ...delta }
 };
 
 
-/**
- * Get all local storage items in a key-value object.
- * @returns object of items in local storage.
- */
-export const getAllLocalStorageItems = () => {
-  var localStorageItems = {};
-
-  for (var i = 0, len = localStorage.length; i < len; i++) {
-    var key = localStorage.key(i);
-    localStorageItems[key] = localStorage.getItem(key);
-  }
-
-  return localStorageItems;
-}
-
-
 /**
  * Handle frontend event or send the event to the backend via Websocket.
  * @param event The event to send.
@@ -346,7 +319,9 @@ export const connect = async (
   // On each received message, queue the updates and events.
   socket.current.on("event", message => {
     const update = JSON5.parse(message)
-    dispatch(update.delta)
+    for (const substate in update.delta) {
+      dispatch[substate](update.delta[substate])
+    }
     applyClientStorageDelta(client_storage, update.delta)
     event_processing = !update.final
     if (update.events) {
@@ -524,23 +499,21 @@ const applyClientStorageDelta = (client_storage, delta) => {
 
 /**
  * Establish websocket event loop for a NextJS page.
- * @param initial_state The initial app state.
- * @param initial_events Function that returns the initial app events.
+ * @param dispatch The reducer dispatch function to update state.
+ * @param initial_events The initial app events.
  * @param client_storage The client storage object from context.js
  *
- * @returns [state, addEvents, connectError] -
- *   state is a reactive dict,
+ * @returns [addEvents, connectError] -
  *   addEvents is used to queue an event, and
  *   connectError is a reactive js error from the websocket connection (or null if connected).
  */
 export const useEventLoop = (
-  initial_state = {},
+  dispatch,
   initial_events = () => [],
   client_storage = {},
 ) => {
   const socket = useRef(null)
   const router = useRouter()
-  const [state, dispatch] = useReducer(applyDelta, initial_state)
   const [connectError, setConnectError] = useState(null)
 
   // Function to add new events to the event queue.
@@ -570,7 +543,7 @@ export const useEventLoop = (
       return;
     }
     // only use websockets if state is present
-    if (Object.keys(state).length > 0) {
+    if (Object.keys(initialState).length > 0) {
       // Initialize the websocket connection.
       if (!socket.current) {
         connect(socket, dispatch, ['websocket', 'polling'], setConnectError, client_storage)
@@ -583,7 +556,17 @@ export const useEventLoop = (
       })()
     }
   })
-  return [state, addEvents, connectError]
+
+  // Route after the initial page hydration.
+  useEffect(() => {
+    const change_complete = () => addEvents(initial_events())
+    router.events.on('routeChangeComplete', change_complete)
+    return () => {
+      router.events.off('routeChangeComplete', change_complete)
+    }
+  }, [router])
+
+  return [addEvents, connectError]
 }
 
 /***

reflex/app.py

@@ -63,7 +63,7 @@
     StateUpdate,
 )
 from reflex.utils import console, format, prerequisites, types
-from reflex.vars import ImportVar
+from reflex.utils.imports import ImportVar
 
 # Define custom types.
 ComponentCallable = Callable[[], Component]

reflex/compiler/compiler.py

@@ -10,40 +10,10 @@
 from reflex.components.component import Component, ComponentStyle, CustomComponent
 from reflex.config import get_config
 from reflex.state import State
-from reflex.utils import imports
-from reflex.vars import ImportVar
+from reflex.utils.imports import ImportDict, ImportVar
 
 # Imports to be included in every Reflex app.
-DEFAULT_IMPORTS: imports.ImportDict = {
-    "react": [
-        ImportVar(tag="Fragment"),
-        ImportVar(tag="useEffect"),
-        ImportVar(tag="useRef"),
-        ImportVar(tag="useState"),
-        ImportVar(tag="useContext"),
-    ],
-    "next/router": [ImportVar(tag="useRouter")],
-    f"/{constants.Dirs.STATE_PATH}": [
-        ImportVar(tag="uploadFiles"),
-        ImportVar(tag="Event"),
-        ImportVar(tag="isTrue"),
-        ImportVar(tag="spreadArraysOrObjects"),
-        ImportVar(tag="preventDefault"),
-        ImportVar(tag="refs"),
-        ImportVar(tag="getRefValue"),
-        ImportVar(tag="getRefValues"),
-        ImportVar(tag="getAllLocalStorageItems"),
-        ImportVar(tag="useEventLoop"),
-    ],
-    "/utils/context.js": [
-        ImportVar(tag="EventLoopContext"),
-        ImportVar(tag="initialEvents"),
-        ImportVar(tag="StateContext"),
-        ImportVar(tag="ColorModeContext"),
-    ],
-    "/utils/helpers/range.js": [
-        ImportVar(tag="range", is_default=True),
-    ],
+DEFAULT_IMPORTS: ImportDict = {
     "": [ImportVar(tag="focus-visible/dist/focus-visible", install=False)],
 }
 

reflex/compiler/templates.py

@@ -3,7 +3,7 @@
 from jinja2 import Environment, FileSystemLoader, Template
 
 from reflex import constants
-from reflex.utils.format import json_dumps
+from reflex.utils.format import format_state_name, json_dumps
 
 
 class ReflexJinjaEnvironment(Environment):
@@ -19,6 +19,7 @@ def __init__(self) -> None:
         )
         self.filters["json_dumps"] = json_dumps
         self.filters["react_setter"] = lambda state: f"set{state.capitalize()}"
+        self.filters["var_name"] = format_state_name
         self.loader = FileSystemLoader(constants.Templates.Dirs.JINJA_TEMPLATE)
         self.globals["const"] = {
             "socket": constants.CompileVars.SOCKET,

reflex/compiler/utils.py

@@ -24,13 +24,12 @@
 from reflex.state import Cookie, LocalStorage, State
 from reflex.style import Style
 from reflex.utils import console, format, imports, path_ops
-from reflex.vars import ImportVar
 
 # To re-export this function.
 merge_imports = imports.merge_imports
 
 
-def compile_import_statement(fields: list[ImportVar]) -> tuple[str, list[str]]:
+def compile_import_statement(fields: list[imports.ImportVar]) -> tuple[str, list[str]]:
     """Compile an import statement.
 
     Args:
@@ -343,7 +342,9 @@ def get_context_path() -> str:
     Returns:
         The path of the context module.
     """
-    return os.path.join(constants.Dirs.WEB_UTILS, "context" + constants.Ext.JS)
+    return os.path.join(
+        constants.Dirs.WEB, constants.Dirs.CONTEXTS_PATH + constants.Ext.JS
+    )
 
 
 def get_components_path() -> str: