ENG-10089: Support custom JS wrappers and bare functions in @rx.memo

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. |

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.

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.

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,40 @@ 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.
+
+    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")
+    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"
+
+
 def memo_components_template(
     imports: list[_ImportDict],
     components: list[dict[str, Any]],
@@ -729,16 +764,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 +805,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}

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",

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"
 }

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")],
 }
 

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,