From 0bb48818b288a86682531f6580980af7897aec49 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 9 Jul 2026 22:45:15 +0000 Subject: [PATCH 1/3] Allow customizing the wrapper function in generated rx.memo (ENG-10089) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit @rx.memo now accepts a wrapper= keyword controlling the JS function the compiled function component is wrapped in. The default is unchanged: React's memo, now expressed as DEFAULT_MEMO_WRAPPER — a FunctionStringVar carrying its own react import. Passing another Var swaps the wrapper (its VarData imports ride along, feeding module imports and frontend package detection), and wrapper=None exports the bare function component. - MemoComponentDefinition grows a wrapper field; the compiler emits it via the render dict and a shared _render_memo_component template helper. - The react memo import moves off _MEMO_BASE_IMPORTS onto the wrapper var, so modules whose memos swap or drop the wrapper no longer import it. - Var-returning memos reject an explicit wrapper= with a TypeError. - @rx.memo() with no arguments now also works, typed via a _MemoDecorator protocol so both component- and var-returning memos check cleanly. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CtJ4TyjTc7nDyDyFn446TH --- docs/library/other/memo.md | 27 +++- .../src/reflex_base/compiler/templates.py | 42 +++--- .../src/reflex_base/components/memo.py | 125 ++++++++++++++--- pyi_hashes.json | 2 +- reflex/compiler/compiler.py | 9 +- reflex/compiler/utils.py | 9 ++ .../integration/tests_playwright/test_memo.py | 27 ++++ tests/units/components/test_memo.py | 126 +++++++++++++++++- 8 files changed, 320 insertions(+), 47 deletions(-) diff --git a/docs/library/other/memo.md b/docs/library/other/memo.md index a1a7ea10e4c..f56b2bbbe6e 100644 --- a/docs/library/other/memo.md +++ b/docs/library/other/memo.md @@ -165,6 +165,29 @@ def index(): The body of a `Var`-returning memo runs at compile time and is restricted to Var operations — no hooks, no Python branching on the Vars. +## Customizing the JavaScript Wrapper + +By default the compiled function component is wrapped in React's [`memo`](https://react.dev/reference/react/memo) helper. Pass `wrapper=` to swap it for a different function — any `Var` whose JavaScript expression is callable, typically an `rx.vars.FunctionStringVar` carrying its own imports — or pass `wrapper=None` to export the bare function component with no wrapper at all. + +```python +observer = rx.vars.FunctionStringVar( + "observer", + _var_data=rx.vars.VarData(imports={"mobx-react-lite": "observer"}), +) + + +@rx.memo(wrapper=observer) +def observed_panel(label: rx.Var[str]) -> rx.Component: + return rx.text(label) + + +@rx.memo(wrapper=None) +def custom_sankey_node(x: rx.Var[int], y: rx.Var[int]) -> rx.Component: + return rx.box(width=x.to_string(), height=y.to_string()) +``` + +The wrapper's imports ride along with the `Var`, so a custom wrapper brings its own import statement and `wrapper=None` pulls in nothing. `wrapper=` only applies to component-returning memos — a `Var`-returning memo compiles to a plain function and rejects the argument. + ## Performance Considerations Reach for `rx.memo` when: @@ -201,10 +224,12 @@ The old `rx._x.memo` alias still resolves to the new memo and prints a one-time ```python rx.memo(component_fn) +rx.memo(wrapper=...)(component_fn) ``` -Wraps a function whose parameters are all `rx.Var[...]` or `rx.RestProp`. Returns a callable that constructs the memoized component (or a `Var` if the function's return annotation is `rx.Var[T]`). +Wraps a function whose parameters are all `rx.Var[...]` or `rx.RestProp`. Returns a callable that constructs the memoized component (or a `Var` if the function's return annotation is `rx.Var[T]`). Called with only keyword arguments, it returns a decorator applying them. | Argument | Type | Description | | --- | --- | --- | | `component_fn` | `Callable[..., rx.Component \| rx.Var]` | The function to memoize. All parameters must be `rx.Var[...]` or `rx.RestProp`. | +| `wrapper` | `rx.Var \| None` | The JS function the compiled function component is wrapped in. Defaults to React's `memo`; pass `None` to omit the wrapper. Only supported on component-returning memos. | diff --git a/packages/reflex-base/src/reflex_base/compiler/templates.py b/packages/reflex-base/src/reflex_base/compiler/templates.py index 8e6e63a2bb3..7d358df0f89 100644 --- a/packages/reflex-base/src/reflex_base/compiler/templates.py +++ b/packages/reflex-base/src/reflex_base/compiler/templates.py @@ -706,6 +706,28 @@ def dynamic_components_module_template( return f"{imports_str}\n{memoized_code}" +def _render_memo_component(component: dict[str, Any]) -> str: + """Render the ``export const`` statement for one memoized component. + + Args: + component: The component render dict (name, signature, render, hooks, + and the optional ``wrapper`` JS expression the function component + is wrapped in). + + Returns: + Rendered component export as string. + """ + function_expr = f"""(({component["signature"]}) => {{ + {_render_hooks(component.get("hooks", {}))} + return( + {_RenderUtils.render(component["render"])} + ) +}})""" + wrapper = component.get("wrapper") + export_expr = f"{wrapper}{function_expr}" if wrapper else function_expr + return f"\nexport const {component['name']} = {export_expr};\n" + + def memo_components_template( imports: list[_ImportDict], components: list[dict[str, Any]], @@ -729,16 +751,7 @@ def memo_components_template( dynamic_imports_str = "\n".join(dynamic_imports) custom_code_str = "\n".join(custom_codes) - components_code = "" - for component in components: - components_code += f""" -export const {component["name"]} = memo(({component["signature"]}) => {{ - {_render_hooks(component.get("hooks", {}))} - return( - {_RenderUtils.render(component["render"])} - ) -}}); -""" + components_code = "".join(map(_render_memo_component, components)) functions_code = "" for function in functions: @@ -779,14 +792,7 @@ def memo_single_component_template( dynamic_imports_str = "\n".join(dynamic_imports) custom_code_str = "\n".join(custom_codes) - component_code = f""" -export const {component["name"]} = memo(({component["signature"]}) => {{ - {_render_hooks(component.get("hooks", {}))} - return( - {_RenderUtils.render(component["render"])} - ) -}}); -""" + component_code = _render_memo_component(component) return f""" {imports_str} diff --git a/packages/reflex-base/src/reflex_base/components/memo.py b/packages/reflex-base/src/reflex_base/components/memo.py index 979d1501bd5..80334ab6804 100644 --- a/packages/reflex-base/src/reflex_base/components/memo.py +++ b/packages/reflex-base/src/reflex_base/components/memo.py @@ -8,15 +8,17 @@ from collections.abc import Callable, Mapping, Sequence from copy import copy from enum import Enum -from functools import cache, update_wrapper +from functools import cache, partial, update_wrapper from types import UnionType from typing import ( Annotated, Any, ClassVar, Generic, + Protocol, TypeVar, Union, + cast, get_args, get_origin, get_type_hints, @@ -61,6 +63,15 @@ # entry point. ``memo.py`` is imported lazily, after ``environment`` is ready. EMPTY_VAR_COMPONENT: Var[Component] = LiteralVar.create(Component.create()) +# The default JS wrapper applied to a compiled component memo's function +# definition: React's ``memo``, carrying its own import. ``@rx.memo`` accepts a +# ``wrapper=`` override to swap it for another helper, or ``None`` to export +# the bare function component. +DEFAULT_MEMO_WRAPPER: FunctionVar = FunctionStringVar.create( + "memo", + _var_data=VarData(imports={"react": [ImportVar(tag="memo")]}), +) + # Base ``Component`` props a memo accepts without an ``rx.RestProp`` (with a # deprecation warning). Only ``key`` qualifies: React consumes it at the # reconciliation layer, so it takes effect on the rendered element even though @@ -292,6 +303,11 @@ class MemoComponentDefinition(MemoDefinition): # imports collection, so descendants emit their refs/imports/hooks in the # page scope rather than being duplicated inside the memo body. passthrough_hole_child: Component | None = None + # The JS function the compiled function component is wrapped in — React's + # ``memo`` by default. ``None`` exports the bare function component. The + # wrapper's ``VarData`` supplies its imports, so a custom wrapper brings + # its own and ``None`` pulls in nothing. + wrapper: Var | None = DEFAULT_MEMO_WRAPPER @property def component(self) -> Component: @@ -1864,31 +1880,23 @@ def _warn_legacy_base_props(fn_name: str, prop_names: Sequence[str]) -> None: ) -@overload -def memo(fn: Callable[..., Component]) -> _MemoComponentWrapper: ... -@overload -def memo(fn: Callable[..., Var[_MemoVarT]]) -> _MemoFunctionWrapper: ... -def memo(fn: Callable[..., Any]) -> _MemoComponentWrapper | _MemoFunctionWrapper: - """Create a memo from a function. - - The decorated function's body is **not** executed here. Only signature-level - analysis — return annotation, parameter kinds, name-collision registration, - and the deprecation warning for missing annotations — runs at decoration - time. The body is compiled lazily on first read of ``.component`` / - ``.function`` — when the component wrapper is instantiated, or when the - compiler reads the memo (see ``_LazyBody``). Deferring the body keeps - ``@rx.memo`` free of import-time side effects, so a memo whose body - references another module no longer forces that module to load during - import — sidestepping circular-import ordering issues. +def _memo_impl( + fn: Callable[..., Any], + wrapper: Var | None, +) -> _MemoComponentWrapper | _MemoFunctionWrapper: + """Analyze and register a memo definition for a decorated function. Args: fn: The function to memoize. + wrapper: The JS wrapper for a component-returning memo, or ``None`` + for no wrapper. Returns: The wrapped function or component factory. Raises: - TypeError: If the return annotation is not supported. + TypeError: If the return annotation is not supported, or a non-default + ``wrapper`` is given for a var-returning memo. """ hints = get_type_hints(fn, include_extras=True) return_annotation = hints.get("return", inspect.Signature.empty) @@ -1904,6 +1912,13 @@ def memo(fn: Callable[..., Any]) -> _MemoComponentWrapper | _MemoFunctionWrapper f"`rx.Var[...]`, got `{return_annotation}`." ) raise TypeError(msg) + if not is_component and wrapper is not DEFAULT_MEMO_WRAPPER: + msg = ( + "`@rx.memo` only supports `wrapper=` on component-returning memos; " + f"`{fn.__name__}` returns `rx.Var[...]`, which compiles to a plain " + "function." + ) + raise TypeError(msg) defaulted_params: list[str] = [] params = _analyze_params( @@ -1922,8 +1937,9 @@ def memo(fn: Callable[..., Any]) -> _MemoComponentWrapper | _MemoFunctionWrapper # first read of ``.component`` / ``.function`` (see ``_LazyBody``), not here, # so decoration has no import-time side effects. The component placeholder # stands in for re-entrant reads during a recursive memo's own evaluation, - # where the name resolves to ``wrapper`` (already bound by first use). + # where the name resolves to ``memo_callable`` (already bound by first use). definition: MemoComponentDefinition | MemoFunctionDefinition + memo_callable: _MemoComponentWrapper | _MemoFunctionWrapper if is_component: definition = MemoComponentDefinition( fn=fn, @@ -1935,8 +1951,9 @@ def memo(fn: Callable[..., Any]) -> _MemoComponentWrapper | _MemoFunctionWrapper lambda: _evaluate_component_body(fn, params), placeholder=Fragment.create(), ), + wrapper=wrapper, ) - wrapper = _create_component_wrapper(definition) + memo_callable = _create_component_wrapper(definition) else: definition = MemoFunctionDefinition( fn=fn, @@ -1950,13 +1967,77 @@ def memo(fn: Callable[..., Any]) -> _MemoComponentWrapper | _MemoFunctionWrapper source_module=source_module, ), ) - wrapper = _create_function_wrapper(definition) + memo_callable = _create_function_wrapper(definition) _register_memo_definition(definition) - return wrapper + return memo_callable + + +class _MemoDecorator(Protocol): + """The decorator returned by ``rx.memo()`` called with no arguments.""" + + @overload + def __call__(self, fn: Callable[..., Component]) -> _MemoComponentWrapper: ... + @overload + def __call__(self, fn: Callable[..., Var[_MemoVarT]]) -> _MemoFunctionWrapper: ... + + +@overload +def memo(fn: Callable[..., Component]) -> _MemoComponentWrapper: ... +@overload +def memo(fn: Callable[..., Var[_MemoVarT]]) -> _MemoFunctionWrapper: ... +@overload +def memo() -> _MemoDecorator: ... +@overload +def memo( + *, wrapper: Var | None +) -> Callable[[Callable[..., Component]], _MemoComponentWrapper]: ... +def memo( + fn: Callable[..., Any] | None = None, + *, + wrapper: Var | None = DEFAULT_MEMO_WRAPPER, +) -> ( + _MemoComponentWrapper + | _MemoFunctionWrapper + | _MemoDecorator + | Callable[[Callable[..., Component]], _MemoComponentWrapper] +): + """Create a memo from a function. + + The decorated function's body is **not** executed here. Only signature-level + analysis — return annotation, parameter kinds, name-collision registration, + and the deprecation warning for missing annotations — runs at decoration + time. The body is compiled lazily on first read of ``.component`` / + ``.function`` — when the component wrapper is instantiated, or when the + compiler reads the memo (see ``_LazyBody``). Deferring the body keeps + ``@rx.memo`` free of import-time side effects, so a memo whose body + references another module no longer forces that module to load during + import — sidestepping circular-import ordering issues. + + Args: + fn: The function to memoize. When omitted, returns a decorator that + applies the given keyword arguments (``@rx.memo(wrapper=...)``). + wrapper: The JS function the compiled function component is wrapped in. + Defaults to React's ``memo``; pass another ``Var`` (typically an + ``rx.vars.FunctionStringVar`` carrying its own imports) to swap the + wrapper, or ``None`` to export the bare function component. Only + supported on component-returning memos. + + Returns: + The wrapped function or component factory, or — when ``fn`` is omitted + — a decorator applying the keyword arguments. + + Raises: + TypeError: If the return annotation is not supported, or a non-default + ``wrapper`` is given for a var-returning memo. + """ + if fn is None: + return cast("_MemoDecorator", partial(_memo_impl, wrapper=wrapper)) + return _memo_impl(fn, wrapper) __all__ = [ + "DEFAULT_MEMO_WRAPPER", "EMPTY_VAR_COMPONENT", "MEMOS", "MemoComponent", diff --git a/pyi_hashes.json b/pyi_hashes.json index f5b97d72d2b..1ec6e0eb1d6 100644 --- a/pyi_hashes.json +++ b/pyi_hashes.json @@ -120,5 +120,5 @@ "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "58521fcd1b514804f534d97624e82c9a", "reflex/__init__.pyi": "56385a4f0d9431eb0056dbc5553a58f9", "reflex/components/__init__.pyi": "9facd05a776d0641432696bbf8e34388", - "reflex/experimental/memo.pyi": "00c9ab0ff3086150278fb330953eeee8" + "reflex/experimental/memo.pyi": "36e5d5f97eb64e94c0974e909e7e2952" } diff --git a/reflex/compiler/compiler.py b/reflex/compiler/compiler.py index 4e521d47b92..dc18eeaf8b3 100644 --- a/reflex/compiler/compiler.py +++ b/reflex/compiler/compiler.py @@ -424,11 +424,12 @@ def add_function( _extend_imports_in_place(self.imports, memo_imports) -# Imports every memo module needs regardless of its body: ``memo`` from React -# for the wrapper, and ``isTrue`` for prop coercion. Shared by the grouped and -# un-mirrored compile paths so they can't drift apart. +# Imports every memo module needs regardless of its body: ``isTrue`` for prop +# coercion. The component wrapper import (``memo`` from React by default) +# rides on each definition's ``wrapper`` var data instead, so a module whose +# memos swap or drop the default wrapper doesn't import it. Shared by the +# grouped and un-mirrored compile paths so they can't drift apart. _MEMO_BASE_IMPORTS: dict[str, list[ImportVar]] = { - "react": [ImportVar(tag="memo")], f"$/{constants.Dirs.STATE_PATH}": [ImportVar(tag="isTrue")], } diff --git a/reflex/compiler/utils.py b/reflex/compiler/utils.py index fb537e4d176..93c23b0a857 100644 --- a/reflex/compiler/utils.py +++ b/reflex/compiler/utils.py @@ -433,6 +433,14 @@ def compile_experimental_component_memo( imports.setdefault("@emotion/react", []).append(ImportVar("jsx")) + # The wrapper import (``memo`` from React by default) rides on the wrapper + # var itself, so a custom wrapper brings its own imports and ``None`` + # pulls in nothing. + wrapper = definition.wrapper + if wrapper is not None and (wrapper_var_data := wrapper._get_all_var_data()): + for lib, fields in wrapper_var_data.imports: + imports.setdefault(lib, []).extend(fields) + signature_fields = [ field for param in definition.params @@ -456,6 +464,7 @@ def compile_experimental_component_memo( fields=tuple(signature_fields), rest=rest_param.placeholder_name if rest_param is not None else None, ).to_javascript(), + "wrapper": str(wrapper) if wrapper is not None else None, "render": rendered, "hooks": hooks, "custom_code": custom_code, diff --git a/tests/integration/tests_playwright/test_memo.py b/tests/integration/tests_playwright/test_memo.py index 0d43ef526e5..ebc2d3c4431 100644 --- a/tests/integration/tests_playwright/test_memo.py +++ b/tests/integration/tests_playwright/test_memo.py @@ -81,6 +81,12 @@ def keyed_row(label: rx.Var[str]) -> rx.Component: # element id so each row is locatable after reordering. return rx.input(id=label) + @rx.memo(wrapper=None) + def unwrapped_label(value: rx.Var[str]) -> rx.Component: + # Compiled without the React ``memo`` wrapper: a bare function + # component that must still render and follow its prop. + return rx.text(value, id="unwrapped-label") + def index() -> rx.Component: return rx.vstack( rx.text(MemoState.last_value, id="memo-last-value"), @@ -100,6 +106,7 @@ def index() -> rx.Component: ), id="keyed-rows", ), + unwrapped_label(value=MemoState.last_value), ) app = rx.App() @@ -222,3 +229,23 @@ def test_memo_key_preserves_identity_across_reorder( # ... while each row kept the value typed into it, by key, not by slot. for row_id in ("row-a", "row-b", "row-c"): expect(page.locator(f"#{row_id}")).to_have_value(row_id.upper()) + + +def test_memo_wrapper_none_renders_and_updates( + memo_app: AppHarness, page: Page +) -> None: + """A ``wrapper=None`` memo renders as a bare function component. + + The compiled module exports the component without React's ``memo`` + wrapper; it must still mount and re-render when its prop changes. + + Args: + memo_app: Running app harness. + page: Playwright page. + """ + assert memo_app.frontend_url is not None + page.goto(memo_app.frontend_url) + + expect(page.locator("#unwrapped-label")).to_have_text("") + page.locator("#memo-input").fill("unwrapped_update") + expect(page.locator("#unwrapped-label")).to_have_text("unwrapped_update") diff --git a/tests/units/components/test_memo.py b/tests/units/components/test_memo.py index 9fa38e51801..3408e523407 100644 --- a/tests/units/components/test_memo.py +++ b/tests/units/components/test_memo.py @@ -11,6 +11,7 @@ from reflex_base.components.component import Component from reflex_base.components.memo import ( _SPECS, + DEFAULT_MEMO_WRAPPER, EMPTY_VAR_COMPONENT, MEMOS, MemoComponent, @@ -31,7 +32,7 @@ from reflex_base.utils.imports import ImportVar from reflex_base.vars import VarData from reflex_base.vars.base import Var -from reflex_base.vars.function import FunctionVar +from reflex_base.vars.function import FunctionStringVar, FunctionVar import reflex as rx from reflex.compiler import compiler @@ -990,6 +991,129 @@ def test_compile_memo_components_falls_back_when_no_source_module(): assert any(path == exp_path for path, _ in files) +def test_default_memo_wrapper_is_react_memo(): + """The default wrapper is React's ``memo``, carrying its own import.""" + assert str(DEFAULT_MEMO_WRAPPER) == "memo" + var_data = DEFAULT_MEMO_WRAPPER._get_all_var_data() + assert var_data is not None + assert [imp.tag for imp in dict(var_data.imports)["react"]] == ["memo"] + + +def test_component_memo_default_wrapper(): + """Bare ``@rx.memo`` wraps the compiled component in React's ``memo``.""" + + @rx.memo + def default_wrapped(label: rx.Var[str]) -> rx.Component: + return rx.text(label) + + definition = MEMOS["DefaultWrapped", __name__] + assert isinstance(definition, MemoComponentDefinition) + assert definition.wrapper is DEFAULT_MEMO_WRAPPER + + files, imports = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + sym = memo_paths.mirrored_symbol("DefaultWrapped", __name__) + assert f"export const {sym} = memo(({{label:labelRxMemo}}) => {{" in code + assert any(imp.tag == "memo" for imp in imports.get("react", [])) + + +def test_component_memo_wrapper_none_emits_bare_function(): + """``@rx.memo(wrapper=None)`` exports the bare function component.""" + + @rx.memo(wrapper=None) + def unwrapped(label: rx.Var[str]) -> rx.Component: + return rx.text(label) + + definition = MEMOS["Unwrapped", __name__] + assert isinstance(definition, MemoComponentDefinition) + assert definition.wrapper is None + + files, imports = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + sym = memo_paths.mirrored_symbol("Unwrapped", __name__) + assert f"export const {sym} = (({{label:labelRxMemo}}) => {{" in code + assert " = memo(" not in code + assert all(imp.tag != "memo" for imp in imports.get("react", [])) + + # The call site is unaffected by the wrapper choice. + component = unwrapped(label="hi") + assert isinstance(component, MemoComponent) + + +def test_component_memo_custom_wrapper(): + """``@rx.memo(wrapper=...)`` swaps React's ``memo`` for the given helper.""" + track_render = FunctionStringVar.create( + "trackRender", + _var_data=VarData(imports={"my-render-lib": [ImportVar(tag="trackRender")]}), + ) + + @rx.memo(wrapper=track_render) + def tracked(label: rx.Var[str]) -> rx.Component: + return rx.text(label) + + definition = MEMOS["Tracked", __name__] + assert isinstance(definition, MemoComponentDefinition) + assert definition.wrapper is track_render + + files, imports = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + sym = memo_paths.mirrored_symbol("Tracked", __name__) + assert f"export const {sym} = trackRender(({{label:labelRxMemo}}) => {{" in code + assert " = memo(" not in code + assert 'import {trackRender} from "my-render-lib"' in code + assert all(imp.tag != "memo" for imp in imports.get("react", [])) + assert any(imp.tag == "trackRender" for imp in imports.get("my-render-lib", [])) + + +def test_component_memo_wrapper_none_in_unmirrored_module(): + """The per-name fallback module honors ``wrapper=None`` too.""" + definition = MemoComponentDefinition( + fn=lambda: None, + python_name="bare_single", + params=(), + export_name="BareSingle", + _component=_LazyBody.ready(rx.text("hi")), + passthrough_hole_child=None, + wrapper=None, + ) + + files, _ = compiler.compile_memo_components((definition,)) + exp_path = compiler._memo_component_file_path( + compiler_utils.get_memo_components_dir(), "BareSingle" + ) + single_code = next(code for path, code in files if path == exp_path) + assert "export const BareSingle = ((" in single_code + assert " = memo(" not in single_code + + +def test_var_returning_memo_rejects_wrapper(): + """``wrapper=`` is only supported on component-returning memos.""" + with pytest.raises(TypeError, match="only supports `wrapper=`"): + + @rx.memo(wrapper=None) # pyright: ignore[reportArgumentType] + def format_id(value: rx.Var[int]) -> rx.Var[str]: + return value.to(str) + + +def test_memo_decorator_parens_form_matches_bare_decorator(): + """``@rx.memo()`` with no arguments behaves like bare ``@rx.memo``.""" + + @rx.memo() + def parens_component(label: rx.Var[str]) -> rx.Component: + return rx.text(label) + + definition = MEMOS["ParensComponent", __name__] + assert isinstance(definition, MemoComponentDefinition) + assert definition.wrapper is DEFAULT_MEMO_WRAPPER + assert isinstance(parens_component(label="hi"), MemoComponent) + + @rx.memo() + def parens_function(value: rx.Var[int]) -> rx.Var[str]: + return value.to(str) + + assert isinstance(MEMOS["parens_function", __name__], MemoFunctionDefinition) + + def test_compile_memo_components_mirrors_underscore_module_without_error(): """A module named ``_internal`` mirrors normally — no reserved-name restriction. From b3ea6783883a3497eebbc606d715df6cd23d212d Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 00:17:57 +0000 Subject: [PATCH 2/3] Parenthesize non-callee memo wrapper expressions before calling An inline wrapper expression like `(Comp) => observer(Comp)` concatenated directly against the component function would swallow the call into the arrow body instead of invoking the wrapper. Wrap anything that isn't an identifier/member chain in parens; `memo(...)` output is unchanged. Addresses PR #6730 review feedback. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CtJ4TyjTc7nDyDyFn446TH --- .../src/reflex_base/compiler/templates.py | 13 +++++++++ tests/units/components/test_memo.py | 28 +++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/packages/reflex-base/src/reflex_base/compiler/templates.py b/packages/reflex-base/src/reflex_base/compiler/templates.py index 7d358df0f89..ca94aa2c482 100644 --- a/packages/reflex-base/src/reflex_base/compiler/templates.py +++ b/packages/reflex-base/src/reflex_base/compiler/templates.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import re from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Literal @@ -706,6 +707,16 @@ def dynamic_components_module_template( return f"{imports_str}\n{memoized_code}" +# Wrapper expressions that are unambiguous JS callees — identifier or member +# chains like ``memo`` / ``React.memo``. Anything else (an inline arrow +# function, a call expression, bracket access) is parenthesized before the +# component function is appended, so the parens bind as the wrapper's call +# rather than being swallowed by the wrapper expression's own grammar (e.g. +# ``(c) => track(c)`` followed by ``(...)`` would otherwise parse the call as +# part of the arrow body). +_MEMO_WRAPPER_CALLEE_RE = re.compile(r"[A-Za-z_$][\w$]*(?:\.[A-Za-z_$][\w$]*)*") + + def _render_memo_component(component: dict[str, Any]) -> str: """Render the ``export const`` statement for one memoized component. @@ -724,6 +735,8 @@ def _render_memo_component(component: dict[str, Any]) -> str: ) }})""" wrapper = component.get("wrapper") + if wrapper and not _MEMO_WRAPPER_CALLEE_RE.fullmatch(wrapper): + wrapper = f"({wrapper})" export_expr = f"{wrapper}{function_expr}" if wrapper else function_expr return f"\nexport const {component['name']} = {export_expr};\n" diff --git a/tests/units/components/test_memo.py b/tests/units/components/test_memo.py index 3408e523407..fb91a1b387e 100644 --- a/tests/units/components/test_memo.py +++ b/tests/units/components/test_memo.py @@ -1065,6 +1065,34 @@ def tracked(label: rx.Var[str]) -> rx.Component: assert any(imp.tag == "trackRender" for imp in imports.get("my-render-lib", [])) +def test_component_memo_inline_function_wrapper_is_parenthesized(): + """A wrapper that isn't a bare callee is parenthesized before the call. + + An inline arrow expression concatenated directly against the component + function would swallow the call into its own body (``(c) => track(c)((...))``); + the emitted code must invoke the wrapper with the component instead. + """ + inline = FunctionStringVar.create( + "(Comp) => trackRender(Comp)", + _var_data=VarData(imports={"my-render-lib": [ImportVar(tag="trackRender")]}), + ) + + @rx.memo(wrapper=inline) + def inline_wrapped(label: rx.Var[str]) -> rx.Component: + return rx.text(label) + + definition = MEMOS["InlineWrapped", __name__] + assert isinstance(definition, MemoComponentDefinition) + + files, _ = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + sym = memo_paths.mirrored_symbol("InlineWrapped", __name__) + assert ( + f"export const {sym} = ((Comp) => trackRender(Comp))(({{label:labelRxMemo}}) => {{" + in code + ) + + def test_component_memo_wrapper_none_in_unmirrored_module(): """The per-name fallback module honors ``wrapper=None`` too.""" definition = MemoComponentDefinition( From 947ca6cbbdd959cff347e2d4bd37b838c225f0c3 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 00:28:45 +0000 Subject: [PATCH 3/3] Add changelog news fragments for memo wrapper customization Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CtJ4TyjTc7nDyDyFn446TH --- news/6730.feature.md | 1 + packages/reflex-base/news/6730.feature.md | 1 + 2 files changed, 2 insertions(+) create mode 100644 news/6730.feature.md create mode 100644 packages/reflex-base/news/6730.feature.md diff --git a/news/6730.feature.md b/news/6730.feature.md new file mode 100644 index 00000000000..1c037614d4e --- /dev/null +++ b/news/6730.feature.md @@ -0,0 +1 @@ +`@rx.memo` components now compile with a configurable JS wrapper: React's `memo` remains the default, `wrapper=` swaps in a custom function `Var` whose imports ride along into the generated module, and `wrapper=None` emits the bare function component. diff --git a/packages/reflex-base/news/6730.feature.md b/packages/reflex-base/news/6730.feature.md new file mode 100644 index 00000000000..7633597fc10 --- /dev/null +++ b/packages/reflex-base/news/6730.feature.md @@ -0,0 +1 @@ +`@rx.memo` now accepts a `wrapper=` argument controlling the JS function that wraps the compiled component definition: keep the default React `memo`, pass a custom function `Var` (e.g. an `rx.vars.FunctionStringVar` carrying its own imports), or pass `wrapper=None` to export the bare function component.