Skip to content
Open
27 changes: 27 additions & 0 deletions packages/reflex-base/src/reflex_base/components/component.py
Original file line number Diff line number Diff line change
Expand Up @@ -1642,6 +1642,30 @@ def _get_custom_code(self) -> str | None:
"""
return None

def _iter_var_module_code(self) -> Iterator[str]:
"""Yield module_code carried by Vars and hook-VarData on this component.

Per-component only — does not recurse into children or prop subtrees.
Callers that need a subtree walk (e.g. :meth:`_get_all_custom_code`)
recurse externally.

Yields:
module_code snippets contributed by this component's Vars.
"""
for var in self._get_vars():
var_data = var._get_all_var_data()
if var_data is None:
continue
yield from var_data.module_code
for hook_var_data in self._get_hooks_internal().values():
if hook_var_data is None:
continue
yield from hook_var_data.module_code
for hook_var_data in self._get_added_hooks().values():
if hook_var_data is None:
continue
yield from hook_var_data.module_code

def _get_all_custom_code(self) -> dict[str, None]:
"""Get custom code for the component and its children.

Expand All @@ -1664,6 +1688,9 @@ def _get_all_custom_code(self) -> dict[str, None]:
for item in clz.add_custom_code(self):
code[item] = None

for snippet in self._iter_var_module_code():
code.setdefault(snippet, None)

# Add the custom code for the children.
for child in self.children:
code |= child._get_all_custom_code()
Expand Down
17 changes: 17 additions & 0 deletions packages/reflex-base/src/reflex_base/vars/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,9 @@ class VarData:
# Components that are part of this var
components: tuple[BaseComponent, ...] = dataclasses.field(default_factory=tuple)

# Module-level JS snippets this var contributes to the page (top-of-file helpers/constants)
module_code: tuple[str, ...] = dataclasses.field(default_factory=tuple)

def __init__(
self,
state: str = "",
Expand All @@ -150,6 +153,7 @@ def __init__(
deps: list[Var] | None = None,
position: Hooks.HookPosition | None = None,
components: Iterable[BaseComponent] | None = None,
module_code: Iterable[str] | None = None,
):
"""Initialize the var data.

Expand All @@ -161,6 +165,7 @@ def __init__(
deps: Dependencies of the var for useCallback.
position: Position of the hook in the component.
components: Components that are part of this var.
module_code: Module-level JS snippets this var contributes to the page.
"""
if isinstance(hooks, str):
hooks = [hooks]
Expand All @@ -176,6 +181,7 @@ def __init__(
object.__setattr__(self, "deps", tuple(deps or []))
object.__setattr__(self, "position", position or None)
object.__setattr__(self, "components", tuple(components or []))
object.__setattr__(self, "module_code", tuple(module_code or []))

if hooks and any(hooks.values()):
# Merge our dependencies first, so they can be referenced.
Expand All @@ -188,6 +194,7 @@ def __init__(
object.__setattr__(self, "deps", merged_var_data.deps)
object.__setattr__(self, "position", merged_var_data.position)
object.__setattr__(self, "components", merged_var_data.components)
object.__setattr__(self, "module_code", merged_var_data.module_code)

def old_school_imports(self) -> ImportDict:
"""Return the imports as a mutable dict.
Expand Down Expand Up @@ -259,6 +266,14 @@ def merge(*all: VarData | None) -> VarData | None:
component for var_data in all_var_datas for component in var_data.components
)

module_code = tuple(
dict.fromkeys(
snippet
for var_data in all_var_datas
for snippet in var_data.module_code
)
)

return VarData(
state=state,
field_name=field_name,
Expand All @@ -267,6 +282,7 @@ def merge(*all: VarData | None) -> VarData | None:
deps=deps,
position=position,
components=components,
module_code=module_code,
)

def __bool__(self) -> bool:
Expand All @@ -283,6 +299,7 @@ def __bool__(self) -> bool:
or self.deps
or self.position
or self.components
or self.module_code
)

@classmethod
Expand Down
17 changes: 17 additions & 0 deletions reflex/compiler/plugins/builtin.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,7 @@ def leave_component(
self._extend_imports(page_context.frontend_imports, imports)

self._collect_component_custom_code(page_context.module_code, comp)
self._collect_var_module_code(page_context.module_code, comp)

if not in_prop_tree:
self._collect_component_hooks(page_context.hooks, comp)
Expand Down Expand Up @@ -252,6 +253,7 @@ def _compiler_bind_leave_component(
extend_imports = self._extend_imports
collect_component_hooks = self._collect_component_hooks
collect_component_custom_code = self._collect_component_custom_code
collect_var_module_code = self._collect_var_module_code
collect_app_wrap_components = self._collect_app_wrap_components
base_get_app_wrap_components = Component._get_app_wrap_components
seen_app_wrap_methods: set[object] = set()
Expand All @@ -269,6 +271,7 @@ def leave_component(
extend_imports(frontend_imports, imports_for_component)

collect_component_custom_code(module_code, comp)
collect_var_module_code(module_code, comp)

if not in_prop_tree:
collect_component_hooks(hooks, comp)
Expand Down Expand Up @@ -329,6 +332,20 @@ def _collect_component_custom_code(
for item in clz.add_custom_code(component):
module_code[item] = None

@staticmethod
def _collect_var_module_code(
module_code: dict[str, None],
component: Component,
) -> None:
"""Collect module_code from VarData attached to this component's Vars.

Per-component contract — the walker re-enters each prop subtree with
``in_prop_tree=True`` so this helper does not recurse, mirroring
:meth:`_collect_component_custom_code`.
"""
for snippet in component._iter_var_module_code():
module_code.setdefault(snippet, None)

def _collect_app_wrap_components(
self,
page_app_wrap_components: dict[tuple[int, str], Component],
Expand Down
2 changes: 2 additions & 0 deletions reflex/compiler/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -544,6 +544,8 @@ def _root_only_custom_code(component: Component) -> dict[str, None]:
for clz in component._iter_parent_classes_with_method("add_custom_code"):
for item in clz.add_custom_code(component):
code[item] = None
for snippet in component._iter_var_module_code():
code.setdefault(snippet, None)
return code


Expand Down
146 changes: 146 additions & 0 deletions tests/integration/tests_playwright/test_var_module_code.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,146 @@
"""Integration test for ``VarData.module_code``.

A Var can declare module-level JS that the page compiler emits at the top of
the page module — alongside ``custom_code`` from Components. When the Var's
``_js_expr`` references that helper, it must be defined for the rendered
output to be correct.

This exercises three facets in one app:

- A Var carrying ``module_code`` directly, used twice on the same page
(deduplication doesn't break correctness).
- Two distinct Vars with different helpers coexisting on a single page
(merge preserves both snippets).
- A Var whose ``module_code`` rides on a *hook's* VarData (the ``__init__``
hook-merge path on ``VarData`` propagates ``module_code`` up).
"""

from collections.abc import Generator

import pytest
from playwright.sync_api import Page, expect

from reflex.testing import AppHarness


def VarModuleCodeApp():
"""App where Vars contribute module-level JS helpers."""
import reflex as rx
from reflex.vars.base import Var, VarData

greet_helper = "const greet = (name) => `Hello, ${name}!`;"
pi_helper = "const PI_APPROX = 3.14;"
counter_helper = "const fmtCount = (n) => `count=${n}`;"

greeting = Var(
_js_expr="greet('World')",
_var_type=str,
_var_data=VarData(module_code=(greet_helper,)),
)
pi = Var(
_js_expr="PI_APPROX",
_var_type=str,
_var_data=VarData(module_code=(pi_helper,)),
)
counter = Var(
_js_expr="fmtCount(0)",
_var_type=str,
_var_data=VarData(
hooks={
"const _unused_counter = 0": VarData(module_code=(counter_helper,)),
},
),
)

def basic():
return rx.box(
rx.text(greeting, id="greeting"),
rx.text(greeting, id="greeting-2"),
)

def multi():
return rx.box(
rx.text(greeting, id="greeting"),
rx.text(pi, id="pi"),
)

def hook():
return rx.box(rx.text(counter, id="counter"))

app = rx.App()
app.add_page(basic, route="/")
app.add_page(multi, route="/multi")
app.add_page(hook, route="/hook")


@pytest.fixture(scope="module")
def var_module_code_app(
tmp_path_factory: pytest.TempPathFactory,
) -> Generator[AppHarness, None, None]:
"""Run the var-module-code app under an AppHarness.

Args:
tmp_path_factory: Pytest fixture for creating temporary directories.

Yields:
The running harness.
"""
with AppHarness.create(
root=tmp_path_factory.mktemp("var_module_code"),
app_source=VarModuleCodeApp,
) as harness:
yield harness


def test_var_module_code_renders_helper_output(
var_module_code_app: AppHarness, page: Page
) -> None:
"""A Var whose ``_js_expr`` calls a ``module_code`` helper renders correctly.

Two usages of the same Var on one page must both resolve — proving the
helper is emitted at module level and that deduplication does not drop it.

Args:
var_module_code_app: Running app harness.
page: Playwright page.
"""
assert var_module_code_app.frontend_url is not None
page.goto(var_module_code_app.frontend_url)

expect(page.locator("#greeting")).to_have_text("Hello, World!")
expect(page.locator("#greeting-2")).to_have_text("Hello, World!")


def test_var_module_code_multiple_distinct_helpers(
var_module_code_app: AppHarness, page: Page
) -> None:
"""Two distinct ``module_code`` Vars on one page each resolve their helper.

Args:
var_module_code_app: Running app harness.
page: Playwright page.
"""
assert var_module_code_app.frontend_url is not None
page.goto(var_module_code_app.frontend_url + "multi")

expect(page.locator("#greeting")).to_have_text("Hello, World!")
expect(page.locator("#pi")).to_have_text("3.14")
Comment thread
FarhanAliRaza marked this conversation as resolved.


def test_var_module_code_via_hook_var_data(
var_module_code_app: AppHarness, page: Page
) -> None:
"""``module_code`` carried on a hook's VarData propagates to the page.

Constructing the outer ``VarData`` triggers the hook-merge fast-forward in
``VarData.__init__``, which must surface the inner ``module_code`` so the
helper is emitted alongside the hook itself.

Args:
var_module_code_app: Running app harness.
page: Playwright page.
"""
assert var_module_code_app.frontend_url is not None
page.goto(var_module_code_app.frontend_url + "hook")

expect(page.locator("#counter")).to_have_text("count=0")
Comment thread
FarhanAliRaza marked this conversation as resolved.
14 changes: 14 additions & 0 deletions tests/units/compiler/test_plugins.py
Original file line number Diff line number Diff line change
Expand Up @@ -788,6 +788,20 @@ def test_default_collector_collects_nested_prop_tree_custom_code_without_recursi
assert "const childCustomCode = 1;" in page_ctx.module_code


def test_default_collector_collects_var_module_code() -> None:
var_with_module_code = LiteralVar.create("v")._replace(
merge_var_data=VarData(module_code=("const fromVar = 42;",))
)
component = ChildComponent.create(id=var_with_module_code)

page_ctx = collect_page_context(
component,
plugins=(DefaultCollectorPlugin(),),
)

assert "const fromVar = 42;" in page_ctx.module_code


def test_default_page_plugins_are_minimal_and_ordered() -> None:
from reflex.compiler.plugins.memoize import MemoizeStatefulPlugin

Expand Down
19 changes: 19 additions & 0 deletions tests/units/components/test_component.py
Original file line number Diff line number Diff line change
Expand Up @@ -571,6 +571,25 @@ def test_get_custom_code(component1: Component, component2: Component):
}


def test_get_all_custom_code_includes_var_module_code(component1: Component):
"""Var-level module_code rides into the legacy _get_all_custom_code path.

This is the entry point used by the memo compile pipeline (see
``compile_experimental_component_memo``); without it, snippets carried by
Vars on a memoized stateful component are silently dropped from the memo
file, and the helper is a runtime ReferenceError.
"""
from reflex.vars.base import Var, VarData

var_with_module_code = Var(
_js_expr="my_helper()",
_var_type=str,
_var_data=VarData(module_code=("const my_helper = () => 1;",)),
)
c = component1.create(id=var_with_module_code)
assert "const my_helper = () => 1;" in c._get_all_custom_code()


def test_get_props(component1, component2):
"""Test that the props are set correctly.

Expand Down
23 changes: 23 additions & 0 deletions tests/units/test_var.py
Original file line number Diff line number Diff line change
Expand Up @@ -1909,6 +1909,29 @@ def test_var_data_with_hooks_value():
assert var_data == VarData(hooks=["whott", "whot", "what"])


def test_var_data_module_code_default_and_truthiness():
assert VarData().module_code == ()
assert not bool(VarData())
assert bool(VarData(module_code=("const A = 1;",)))


def test_var_data_module_code_merge_dedupes_preserving_order():
merged = VarData.merge(
VarData(module_code=("a;",)),
VarData(module_code=("b;",)),
VarData(module_code=("a;",)),
)
assert merged is not None
assert merged.module_code == ("a;", "b;")


def test_var_data_module_code_propagates_through_nested_hook_var_data():
var_data = VarData(
hooks={"useThing": VarData(module_code=("const helper = 1;",))},
)
assert var_data.module_code == ("const helper = 1;",)


def test_str_var_in_components(mocker: MockerFixture):
class StateWithVar(rx.State):
field: int = 1
Expand Down
Loading