Skip to content

Commit b5e0f9d

Browse files
committed
feat: propagate VarData.module_code into page module code
Adds module_code on VarData so Vars can contribute top-of-file JS helpers/constants. The default collector and the legacy _get_all_custom_code path both pick it up, ensuring snippets carried by Vars on memoized stateful components aren't dropped from the memo file.
1 parent a84e29b commit b5e0f9d

8 files changed

Lines changed: 265 additions & 0 deletions

File tree

packages/reflex-base/src/reflex_base/components/component.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1642,6 +1642,30 @@ def _get_custom_code(self) -> str | None:
16421642
"""
16431643
return None
16441644

1645+
def _iter_var_module_code(self) -> Iterator[str]:
1646+
"""Yield module_code carried by Vars and hook-VarData on this component.
1647+
1648+
Per-component only — does not recurse into children or prop subtrees.
1649+
Callers that need a subtree walk (e.g. :meth:`_get_all_custom_code`)
1650+
recurse externally.
1651+
1652+
Yields:
1653+
module_code snippets contributed by this component's Vars.
1654+
"""
1655+
for var in self._get_vars():
1656+
var_data = var._get_all_var_data()
1657+
if var_data is None:
1658+
continue
1659+
yield from var_data.module_code
1660+
for hook_var_data in self._get_hooks_internal().values():
1661+
if hook_var_data is None:
1662+
continue
1663+
yield from hook_var_data.module_code
1664+
for hook_var_data in self._get_added_hooks().values():
1665+
if hook_var_data is None:
1666+
continue
1667+
yield from hook_var_data.module_code
1668+
16451669
def _get_all_custom_code(self) -> dict[str, None]:
16461670
"""Get custom code for the component and its children.
16471671
@@ -1664,6 +1688,9 @@ def _get_all_custom_code(self) -> dict[str, None]:
16641688
for item in clz.add_custom_code(self):
16651689
code[item] = None
16661690

1691+
for snippet in self._iter_var_module_code():
1692+
code.setdefault(snippet, None)
1693+
16671694
# Add the custom code for the children.
16681695
for child in self.children:
16691696
code |= child._get_all_custom_code()

packages/reflex-base/src/reflex_base/vars/base.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,9 @@ class VarData:
141141
# Components that are part of this var
142142
components: tuple[BaseComponent, ...] = dataclasses.field(default_factory=tuple)
143143

144+
# Module-level JS snippets this var contributes to the page (top-of-file helpers/constants)
145+
module_code: tuple[str, ...] = dataclasses.field(default_factory=tuple)
146+
144147
def __init__(
145148
self,
146149
state: str = "",
@@ -150,6 +153,7 @@ def __init__(
150153
deps: list[Var] | None = None,
151154
position: Hooks.HookPosition | None = None,
152155
components: Iterable[BaseComponent] | None = None,
156+
module_code: Iterable[str] | None = None,
153157
):
154158
"""Initialize the var data.
155159
@@ -161,6 +165,7 @@ def __init__(
161165
deps: Dependencies of the var for useCallback.
162166
position: Position of the hook in the component.
163167
components: Components that are part of this var.
168+
module_code: Module-level JS snippets this var contributes to the page.
164169
"""
165170
if isinstance(hooks, str):
166171
hooks = [hooks]
@@ -176,6 +181,7 @@ def __init__(
176181
object.__setattr__(self, "deps", tuple(deps or []))
177182
object.__setattr__(self, "position", position or None)
178183
object.__setattr__(self, "components", tuple(components or []))
184+
object.__setattr__(self, "module_code", tuple(module_code or []))
179185

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

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

269+
module_code = tuple(
270+
dict.fromkeys(
271+
snippet
272+
for var_data in all_var_datas
273+
for snippet in var_data.module_code
274+
)
275+
)
276+
262277
return VarData(
263278
state=state,
264279
field_name=field_name,
@@ -267,6 +282,7 @@ def merge(*all: VarData | None) -> VarData | None:
267282
deps=deps,
268283
position=position,
269284
components=components,
285+
module_code=module_code,
270286
)
271287

272288
def __bool__(self) -> bool:
@@ -283,6 +299,7 @@ def __bool__(self) -> bool:
283299
or self.deps
284300
or self.position
285301
or self.components
302+
or self.module_code
286303
)
287304

288305
@classmethod

reflex/compiler/plugins/builtin.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,7 @@ def leave_component(
191191
self._extend_imports(page_context.frontend_imports, imports)
192192

193193
self._collect_component_custom_code(page_context.module_code, comp)
194+
self._collect_var_module_code(page_context.module_code, comp)
194195

195196
if not in_prop_tree:
196197
self._collect_component_hooks(page_context.hooks, comp)
@@ -252,6 +253,7 @@ def _compiler_bind_leave_component(
252253
extend_imports = self._extend_imports
253254
collect_component_hooks = self._collect_component_hooks
254255
collect_component_custom_code = self._collect_component_custom_code
256+
collect_var_module_code = self._collect_var_module_code
255257
collect_app_wrap_components = self._collect_app_wrap_components
256258
base_get_app_wrap_components = Component._get_app_wrap_components
257259
seen_app_wrap_methods: set[object] = set()
@@ -269,6 +271,7 @@ def leave_component(
269271
extend_imports(frontend_imports, imports_for_component)
270272

271273
collect_component_custom_code(module_code, comp)
274+
collect_var_module_code(module_code, comp)
272275

273276
if not in_prop_tree:
274277
collect_component_hooks(hooks, comp)
@@ -329,6 +332,20 @@ def _collect_component_custom_code(
329332
for item in clz.add_custom_code(component):
330333
module_code[item] = None
331334

335+
@staticmethod
336+
def _collect_var_module_code(
337+
module_code: dict[str, None],
338+
component: Component,
339+
) -> None:
340+
"""Collect module_code from VarData attached to this component's Vars.
341+
342+
Per-component contract — the walker re-enters each prop subtree with
343+
``in_prop_tree=True`` so this helper does not recurse, mirroring
344+
:meth:`_collect_component_custom_code`.
345+
"""
346+
for snippet in component._iter_var_module_code():
347+
module_code.setdefault(snippet, None)
348+
332349
def _collect_app_wrap_components(
333350
self,
334351
page_app_wrap_components: dict[tuple[int, str], Component],

reflex/compiler/utils.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -544,6 +544,8 @@ def _root_only_custom_code(component: Component) -> dict[str, None]:
544544
for clz in component._iter_parent_classes_with_method("add_custom_code"):
545545
for item in clz.add_custom_code(component):
546546
code[item] = None
547+
for snippet in component._iter_var_module_code():
548+
code.setdefault(snippet, None)
547549
return code
548550

549551

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
"""Integration test for ``VarData.module_code``.
2+
3+
A Var can declare module-level JS that the page compiler emits at the top of
4+
the page module — alongside ``custom_code`` from Components. When the Var's
5+
``_js_expr`` references that helper, it must be defined for the rendered
6+
output to be correct.
7+
8+
This exercises three facets in one app:
9+
10+
- A Var carrying ``module_code`` directly, used twice on the same page
11+
(deduplication doesn't break correctness).
12+
- Two distinct Vars with different helpers coexisting on a single page
13+
(merge preserves both snippets).
14+
- A Var whose ``module_code`` rides on a *hook's* VarData (the ``__init__``
15+
hook-merge path on ``VarData`` propagates ``module_code`` up).
16+
"""
17+
18+
from collections.abc import Generator
19+
20+
import pytest
21+
from playwright.sync_api import Page, expect
22+
23+
from reflex.testing import AppHarness
24+
25+
26+
def VarModuleCodeApp():
27+
"""App where Vars contribute module-level JS helpers."""
28+
import reflex as rx
29+
from reflex.vars.base import Var, VarData
30+
31+
greet_helper = "const greet = (name) => `Hello, ${name}!`;"
32+
pi_helper = "const PI_APPROX = 3.14;"
33+
counter_helper = "const fmtCount = (n) => `count=${n}`;"
34+
35+
greeting = Var(
36+
_js_expr="greet('World')",
37+
_var_type=str,
38+
_var_data=VarData(module_code=(greet_helper,)),
39+
)
40+
pi = Var(
41+
_js_expr="PI_APPROX",
42+
_var_type=str,
43+
_var_data=VarData(module_code=(pi_helper,)),
44+
)
45+
counter = Var(
46+
_js_expr="fmtCount(0)",
47+
_var_type=str,
48+
_var_data=VarData(
49+
hooks={
50+
"const _unused_counter = 0": VarData(module_code=(counter_helper,)),
51+
},
52+
),
53+
)
54+
55+
def basic():
56+
return rx.box(
57+
rx.text(greeting, id="greeting"),
58+
rx.text(greeting, id="greeting-2"),
59+
)
60+
61+
def multi():
62+
return rx.box(
63+
rx.text(greeting, id="greeting"),
64+
rx.text(pi, id="pi"),
65+
)
66+
67+
def hook():
68+
return rx.box(rx.text(counter, id="counter"))
69+
70+
app = rx.App()
71+
app.add_page(basic, route="/")
72+
app.add_page(multi, route="/multi")
73+
app.add_page(hook, route="/hook")
74+
75+
76+
@pytest.fixture(scope="module")
77+
def var_module_code_app(
78+
tmp_path_factory: pytest.TempPathFactory,
79+
) -> Generator[AppHarness, None, None]:
80+
"""Run the var-module-code app under an AppHarness.
81+
82+
Args:
83+
tmp_path_factory: Pytest fixture for creating temporary directories.
84+
85+
Yields:
86+
The running harness.
87+
"""
88+
with AppHarness.create(
89+
root=tmp_path_factory.mktemp("var_module_code"),
90+
app_source=VarModuleCodeApp,
91+
) as harness:
92+
yield harness
93+
94+
95+
def test_var_module_code_renders_helper_output(
96+
var_module_code_app: AppHarness, page: Page
97+
) -> None:
98+
"""A Var whose ``_js_expr`` calls a ``module_code`` helper renders correctly.
99+
100+
Two usages of the same Var on one page must both resolve — proving the
101+
helper is emitted at module level and that deduplication does not drop it.
102+
103+
Args:
104+
var_module_code_app: Running app harness.
105+
page: Playwright page.
106+
"""
107+
assert var_module_code_app.frontend_url is not None
108+
page.goto(var_module_code_app.frontend_url)
109+
110+
expect(page.locator("#greeting")).to_have_text("Hello, World!")
111+
expect(page.locator("#greeting-2")).to_have_text("Hello, World!")
112+
113+
114+
def test_var_module_code_multiple_distinct_helpers(
115+
var_module_code_app: AppHarness, page: Page
116+
) -> None:
117+
"""Two distinct ``module_code`` Vars on one page each resolve their helper.
118+
119+
Args:
120+
var_module_code_app: Running app harness.
121+
page: Playwright page.
122+
"""
123+
assert var_module_code_app.frontend_url is not None
124+
page.goto(var_module_code_app.frontend_url + "multi")
125+
126+
expect(page.locator("#greeting")).to_have_text("Hello, World!")
127+
expect(page.locator("#pi")).to_have_text("3.14")
128+
129+
130+
def test_var_module_code_via_hook_var_data(
131+
var_module_code_app: AppHarness, page: Page
132+
) -> None:
133+
"""``module_code`` carried on a hook's VarData propagates to the page.
134+
135+
Constructing the outer ``VarData`` triggers the hook-merge fast-forward in
136+
``VarData.__init__``, which must surface the inner ``module_code`` so the
137+
helper is emitted alongside the hook itself.
138+
139+
Args:
140+
var_module_code_app: Running app harness.
141+
page: Playwright page.
142+
"""
143+
assert var_module_code_app.frontend_url is not None
144+
page.goto(var_module_code_app.frontend_url + "hook")
145+
146+
expect(page.locator("#counter")).to_have_text("count=0")

tests/units/compiler/test_plugins.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -788,6 +788,20 @@ def test_default_collector_collects_nested_prop_tree_custom_code_without_recursi
788788
assert "const childCustomCode = 1;" in page_ctx.module_code
789789

790790

791+
def test_default_collector_collects_var_module_code() -> None:
792+
var_with_module_code = LiteralVar.create("v")._replace(
793+
merge_var_data=VarData(module_code=("const fromVar = 42;",))
794+
)
795+
component = ChildComponent.create(id=var_with_module_code)
796+
797+
page_ctx = collect_page_context(
798+
component,
799+
plugins=(DefaultCollectorPlugin(),),
800+
)
801+
802+
assert "const fromVar = 42;" in page_ctx.module_code
803+
804+
791805
def test_default_page_plugins_are_minimal_and_ordered() -> None:
792806
from reflex.compiler.plugins.memoize import MemoizeStatefulPlugin
793807

tests/units/components/test_component.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -571,6 +571,25 @@ def test_get_custom_code(component1: Component, component2: Component):
571571
}
572572

573573

574+
def test_get_all_custom_code_includes_var_module_code(component1: Component):
575+
"""Var-level module_code rides into the legacy _get_all_custom_code path.
576+
577+
This is the entry point used by the memo compile pipeline (see
578+
``compile_experimental_component_memo``); without it, snippets carried by
579+
Vars on a memoized stateful component are silently dropped from the memo
580+
file, and the helper is a runtime ReferenceError.
581+
"""
582+
from reflex.vars.base import Var, VarData
583+
584+
var_with_module_code = Var(
585+
_js_expr="my_helper()",
586+
_var_type=str,
587+
_var_data=VarData(module_code=("const my_helper = () => 1;",)),
588+
)
589+
c = component1.create(id=var_with_module_code)
590+
assert "const my_helper = () => 1;" in c._get_all_custom_code()
591+
592+
574593
def test_get_props(component1, component2):
575594
"""Test that the props are set correctly.
576595

tests/units/test_var.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1909,6 +1909,29 @@ def test_var_data_with_hooks_value():
19091909
assert var_data == VarData(hooks=["whott", "whot", "what"])
19101910

19111911

1912+
def test_var_data_module_code_default_and_truthiness():
1913+
assert VarData().module_code == ()
1914+
assert not bool(VarData())
1915+
assert bool(VarData(module_code=("const A = 1;",)))
1916+
1917+
1918+
def test_var_data_module_code_merge_dedupes_preserving_order():
1919+
merged = VarData.merge(
1920+
VarData(module_code=("a;",)),
1921+
VarData(module_code=("b;",)),
1922+
VarData(module_code=("a;",)),
1923+
)
1924+
assert merged is not None
1925+
assert merged.module_code == ("a;", "b;")
1926+
1927+
1928+
def test_var_data_module_code_propagates_through_nested_hook_var_data():
1929+
var_data = VarData(
1930+
hooks={"useThing": VarData(module_code=("const helper = 1;",))},
1931+
)
1932+
assert var_data.module_code == ("const helper = 1;",)
1933+
1934+
19121935
def test_str_var_in_components(mocker: MockerFixture):
19131936
class StateWithVar(rx.State):
19141937
field: int = 1

0 commit comments

Comments
 (0)