[0.6.0 blocker] state: update inherited_vars and tracking dicts when adding vars (#2822)

* state: update inherited_vars and tracking dicts when adding vars

Ensure that dynamically added vars are accounted for in dependency and
inheritence tree to avoid unrenderable or stale data.

* Regression test for dynamic route args and inherited_vars

* [flexgen] Initialize app from refactored code

Use the new /api/gen/{hash}/refactored endpoint to get refactored reflex code.

* Use _js_expr instead of _var_name

integration/test_dynamic_routes.py

@@ -66,6 +66,64 @@ def DynamicRoute():
             ),
         )
 
+    class ArgState(rx.State):
+        """The app state."""
+
+        @rx.var
+        def arg(self) -> int:
+            return int(self.arg_str or 0)
+
+    class ArgSubState(ArgState):
+        @rx.var(cache=True)
+        def cached_arg(self) -> int:
+            return self.arg
+
+        @rx.var(cache=True)
+        def cached_arg_str(self) -> str:
+            return self.arg_str
+
+    @rx.page(route="/arg/[arg_str]")
+    def arg() -> rx.Component:
+        return rx.vstack(
+            rx.data_list.root(
+                rx.data_list.item(
+                    rx.data_list.label("rx.State.arg_str (dynamic)"),
+                    rx.data_list.value(rx.State.arg_str, id="state-arg_str"),  # type: ignore
+                ),
+                rx.data_list.item(
+                    rx.data_list.label("ArgState.arg_str (dynamic) (inherited)"),
+                    rx.data_list.value(ArgState.arg_str, id="argstate-arg_str"),  # type: ignore
+                ),
+                rx.data_list.item(
+                    rx.data_list.label("ArgState.arg"),
+                    rx.data_list.value(ArgState.arg, id="argstate-arg"),
+                ),
+                rx.data_list.item(
+                    rx.data_list.label("ArgSubState.arg_str (dynamic) (inherited)"),
+                    rx.data_list.value(ArgSubState.arg_str, id="argsubstate-arg_str"),  # type: ignore
+                ),
+                rx.data_list.item(
+                    rx.data_list.label("ArgSubState.arg (inherited)"),
+                    rx.data_list.value(ArgSubState.arg, id="argsubstate-arg"),
+                ),
+                rx.data_list.item(
+                    rx.data_list.label("ArgSubState.cached_arg"),
+                    rx.data_list.value(
+                        ArgSubState.cached_arg, id="argsubstate-cached_arg"
+                    ),
+                ),
+                rx.data_list.item(
+                    rx.data_list.label("ArgSubState.cached_arg_str"),
+                    rx.data_list.value(
+                        ArgSubState.cached_arg_str, id="argsubstate-cached_arg_str"
+                    ),
+                ),
+            ),
+            rx.link("+", href=f"/arg/{ArgState.arg + 1}", id="next-page"),
+            align="center",
+            height="100vh",
+        )
+
     @rx.page(route="/redirect-page/[page_id]", on_load=DynamicState.on_load_redir)  # type: ignore
     def redirect_page():
         return rx.fragment(rx.text("redirecting..."))
@@ -302,3 +360,50 @@ async def test_on_load_navigate_non_dynamic(
         link.click()
     assert urlsplit(driver.current_url).path == "/static/x/"
     await poll_for_order(["/static/x-no page id", "/static/x-no page id"])
+
+
+@pytest.mark.asyncio
+async def test_render_dynamic_arg(
+    dynamic_route: AppHarness,
+    driver: WebDriver,
+):
+    """Assert that dynamic arg var is rendered correctly in different contexts.
+
+    Args:
+        dynamic_route: harness for DynamicRoute app.
+        driver: WebDriver instance.
+    """
+    assert dynamic_route.app_instance is not None
+    with poll_for_navigation(driver):
+        driver.get(f"{dynamic_route.frontend_url}/arg/0")
+
+    def assert_content(expected: str, expect_not: str):
+        ids = [
+            "state-arg_str",
+            "argstate-arg",
+            "argstate-arg_str",
+            "argsubstate-arg_str",
+            "argsubstate-arg",
+            "argsubstate-cached_arg",
+            "argsubstate-cached_arg_str",
+        ]
+        for id in ids:
+            el = driver.find_element(By.ID, id)
+            assert el
+            assert (
+                dynamic_route.poll_for_content(el, exp_not_equal=expect_not) == expected
+            )
+
+    assert_content("0", "")
+    next_page_link = driver.find_element(By.ID, "next-page")
+    assert next_page_link
+    with poll_for_navigation(driver):
+        next_page_link.click()
+    assert driver.current_url == f"{dynamic_route.frontend_url}/arg/1/"
+    assert_content("1", "0")
+    next_page_link = driver.find_element(By.ID, "next-page")
+    assert next_page_link
+    with poll_for_navigation(driver):
+        next_page_link.click()
+    assert driver.current_url == f"{dynamic_route.frontend_url}/arg/2/"
+    assert_content("2", "1")

reflex/state.py

@@ -1047,6 +1047,27 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             if not func[0].startswith("__")
         }
 
+    @classmethod
+    def _update_substate_inherited_vars(cls, vars_to_add: dict[str, Var]):
+        """Update the inherited vars of substates recursively when new vars are added.
+
+        Also updates the var dependency tracking dicts after adding vars.
+
+        Args:
+            vars_to_add: names to Var instances to add to substates
+        """
+        for substate_class in cls.class_subclasses:
+            for name, var in vars_to_add.items():
+                if types.is_backend_base_variable(name, cls):
+                    substate_class.backend_vars.setdefault(name, var)
+                    substate_class.inherited_backend_vars.setdefault(name, var)
+                else:
+                    substate_class.vars.setdefault(name, var)
+                    substate_class.inherited_vars.setdefault(name, var)
+                substate_class._update_substate_inherited_vars(vars_to_add)
+        # Reinitialize dependency tracking dicts.
+        cls._init_var_dependency_dicts()
+
     @classmethod
     def setup_dynamic_args(cls, args: dict[str, str]):
         """Set up args for easy access in renderer.
@@ -1063,14 +1084,15 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             def inner_func(self) -> str:
                 return self.router.page.params.get(param, "")
 
-            return DynamicRouteVar(fget=inner_func, cache=True)
+            return inner_func
 
         def arglist_factory(param):
             def inner_func(self) -> List[str]:
                 return self.router.page.params.get(param, [])
 
-            return DynamicRouteVar(fget=inner_func, cache=True)
+            return inner_func
 
+        dynamic_vars = {}
         for param, value in args.items():
             if value == constants.RouteArgType.SINGLE:
                 func = argsingle_factory(param)
@@ -1078,16 +1100,18 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
                 func = arglist_factory(param)
             else:
                 continue
-            # to allow passing as a prop, evade python frozen rules (bad practice)
-            object.__setattr__(func, "_js_expr", param)
-            # cls.vars[param] = cls.computed_vars[param] = func._var_set_state(cls)  # type: ignore
-            cls.vars[param] = cls.computed_vars[param] = func._replace(
-                _var_data=VarData.from_state(cls)
+            dynamic_vars[param] = DynamicRouteVar(
+                fget=func,
+                cache=True,
+                _js_expr=param,
+                _var_data=VarData.from_state(cls),
             )
-            setattr(cls, param, func)
+            setattr(cls, param, dynamic_vars[param])
 
-        # Reinitialize dependency tracking dicts.
-        cls._init_var_dependency_dicts()
+        # Update tracking dicts.
+        cls.computed_vars.update(dynamic_vars)
+        cls.vars.update(dynamic_vars)
+        cls._update_substate_inherited_vars(dynamic_vars)
 
     @classmethod
     def _check_overwritten_dynamic_args(cls, args: list[str]):

reflex/vars/base.py

@@ -1687,7 +1687,7 @@ class ComputedVar(Var[RETURN_TYPE]):
         """
         if instance is None:
             state_where_defined = owner
-            while self.fget.__name__ in state_where_defined.inherited_vars:
+            while self._js_expr in state_where_defined.inherited_vars:
                 state_where_defined = state_where_defined.get_parent_state()
 
             return self._replace(