diff --git a/docs/library/other/memo.md b/docs/library/other/memo.md index 2c51ca4638f..00dbfc34b1b 100644 --- a/docs/library/other/memo.md +++ b/docs/library/other/memo.md @@ -4,7 +4,10 @@ import reflex as rx # Memo -The `@rx.memo` decorator turns a function into a memoized React component. The compiler emits the function as its own module, and React's `memo` only re-renders it when its declared props change. Reach for it when a subtree is expensive to render and depends on a narrow slice of state. +The `@rx.memo` decorator emits a component or `rx.Var`-returning function in +its own module and memoizes it at runtime. React component memos re-render only +when their declared props change, while function memos reuse cached results for +unchanged arguments. ## Requirements @@ -47,6 +50,79 @@ def index(): `expensive_component` re-renders only when `label` changes — bumping `DemoState.count` does not invalidate it. +Use `name=` to override the generated memo name, which is useful for lambdas: + +```python +named_memo = rx.memo(name="named_memo")(lambda label: rx.text(label)) +``` + +The name must have valid JavaScript identifier characters. Reflex appends its +memo marker to explicit names so JavaScript keywords remain safe. + +For props whose values are recreated but equal on each render, use +`by_value=True` to compare their serialized values: + +```python +@rx.memo(by_value=True) +def settings_panel(settings: rx.Var[dict[str, str]]) -> rx.Component: + return rx.text(settings["title"]) +``` + +If `wrapper=` is also supplied, the value-based memo wrapper is applied around +the custom wrapper rather than replacing it. + +By default, an explicit memo is the auto-memoization boundary for its children. +Pass `recursive=True` when hook-bearing child components should also be +auto-memoized independently. This is useful when the explicit memo provides a +reusable module boundary but state reads should still re-render as close to +their use as possible: + +```python +class MetricsState(rx.State): + request_count: int = 0 + + +def live_request_count() -> rx.Component: + return rx.text("Requests: ", MetricsState.request_count) + + +@rx.memo(recursive=True) +def dashboard_shell() -> rx.Component: + return rx.card( + rx.heading("Dashboard"), + live_request_count(), + ) + + +def analytics_page() -> rx.Component: + return dashboard_shell() + + +def admin_page() -> rx.Component: + return dashboard_shell() +``` + +`dashboard_shell` is emitted once as a reusable component module. Because it is +recursive, the state-bearing `live_request_count` subtree gets its own nearby +auto-memo boundary instead of making the whole dashboard shell depend directly +on `MetricsState.request_count`. + +## Memoized Functions + +A function returning `rx.Var[...]` is emitted as a JavaScript function and +memoizes its return value by argument identity. Repeated calls with the same +arguments reuse the cached result, including calls made while rendering: + +```python +@rx.memo +def format_total(total: rx.Var[int]) -> rx.Var[str]: + return "$" + total.to(str) +``` + +Use `by_value=True` to key that cache by serialized argument values, or +`wrapper=None` to emit a plain uncached function. Custom wrappers work for +function memos as well. `recursive=True` only applies to component memos. + ## With State Variables Props can be ordinary Vars. The memoized component re-renders when those Vars change: diff --git a/news/+memo-value-recursive.feature.md b/news/+memo-value-recursive.feature.md new file mode 100644 index 00000000000..ee9a8f255bf --- /dev/null +++ b/news/+memo-value-recursive.feature.md @@ -0,0 +1 @@ +Extend `@rx.memo` with `by_value=True` for component props and function arguments, `recursive=True` to auto-memoize reactive descendants, and `name=` to override generated names. Function memos now cache return values by argument identity by default; pass `wrapper=None` to emit a plain function. diff --git a/packages/reflex-base/news/+memo-value-recursive.feature.md b/packages/reflex-base/news/+memo-value-recursive.feature.md new file mode 100644 index 00000000000..ee9a8f255bf --- /dev/null +++ b/packages/reflex-base/news/+memo-value-recursive.feature.md @@ -0,0 +1 @@ +Extend `@rx.memo` with `by_value=True` for component props and function arguments, `recursive=True` to auto-memoize reactive descendants, and `name=` to override generated names. Function memos now cache return values by argument identity by default; pass `wrapper=None` to emit a plain function. diff --git a/packages/reflex-base/src/reflex_base/compiler/templates.py b/packages/reflex-base/src/reflex_base/compiler/templates.py index a4e2ff725c4..1fd4c854d13 100644 --- a/packages/reflex-base/src/reflex_base/compiler/templates.py +++ b/packages/reflex-base/src/reflex_base/compiler/templates.py @@ -909,6 +909,23 @@ def dynamic_components_module_template( _MEMO_WRAPPER_CALLEE_RE = re.compile(r"[A-Za-z_$][\w$]*(?:\.[A-Za-z_$][\w$]*)*") +def _apply_memo_wrapper(expression: str, wrapper: str | None) -> str: + """Apply an optional memo wrapper to a JavaScript function expression. + + Args: + expression: The function expression to wrap. + wrapper: The wrapper expression, or ``None`` for a bare function. + + Returns: + The wrapped JavaScript expression. + """ + if not wrapper: + return expression + if not _MEMO_WRAPPER_CALLEE_RE.fullmatch(wrapper): + wrapper = f"({wrapper})" + return f"{wrapper}{expression}" + + def _render_memo_component(component: dict[str, Any]) -> str: """Render the ``export const`` statement for one memoized component. @@ -932,10 +949,7 @@ def _render_memo_component(component: dict[str, Any]) -> str: {_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 + export_expr = _apply_memo_wrapper(function_expr, component.get("wrapper")) name = component["name"] # ``display_name`` is resolved by the caller (``compile_experimental_component_memo``), # which is the layer that knows the memo's clean export name — the JS symbol @@ -974,9 +988,10 @@ def memo_components_template( functions_code = "" for function in functions: - functions_code += ( - f"\nexport const {function['name']} = {function['function']};\n" + function_expr = _apply_memo_wrapper( + function["function"], function.get("wrapper") ) + functions_code += f"\nexport const {function['name']} = {function_expr};\n" return f""" {imports_str} @@ -1037,10 +1052,11 @@ def memo_single_function_template( The rendered standalone function memo module code. """ imports_str = "\n".join([_RenderUtils.get_import(imp) for imp in imports]) + function_expr = _apply_memo_wrapper(function["function"], function.get("wrapper")) return f""" {imports_str} -export const {function["name"]} = {function["function"]}; +export const {function["name"]} = {function_expr}; """ diff --git a/packages/reflex-base/src/reflex_base/components/memo.py b/packages/reflex-base/src/reflex_base/components/memo.py index a1a9d2aaf8b..75d90230d38 100644 --- a/packages/reflex-base/src/reflex_base/components/memo.py +++ b/packages/reflex-base/src/reflex_base/components/memo.py @@ -73,6 +73,88 @@ _var_data=VarData(imports={"react": [ImportVar(tag="memo")]}), ) +# React's default memo comparator compares each prop by identity. This wrapper +# opts into comparing the serialized props instead, which is useful for memo +# props represented by freshly-created but equal objects. +BY_VALUE_MEMO_WRAPPER: FunctionVar = FunctionStringVar.create( + "(Component) => memo(Component, (prevProps, nextProps) => " + "JSON.stringify(prevProps) === JSON.stringify(nextProps))", + _var_data=VarData(imports={"react": [ImportVar(tag="memo")]}), +) + +# Function memos share one module-level cache. Identity mode uses a Map trie so +# multiple call sites with different argument tuples remain cached concurrently. +_DEFAULT_FUNCTION_MEMO_WRAPPER: FunctionVar = FunctionStringVar.create( + "(fn) => { const resultKey = Symbol(); const cache = new Map(); " + "return (...args) => { let node = cache; for (const arg of args) { " + "if (!node.has(arg)) node.set(arg, new Map()); node = node.get(arg); } " + "if (!node.has(resultKey)) node.set(resultKey, fn(...args)); " + "return node.get(resultKey); }; }" +) + +# Value mode uses serialized argument tuples as cache keys. +_BY_VALUE_FUNCTION_MEMO_WRAPPER: FunctionVar = FunctionStringVar.create( + "(fn) => { const cache = new Map(); return (...args) => { " + "const key = JSON.stringify(args); " + "if (!cache.has(key)) cache.set(key, fn(...args)); return cache.get(key); }; }" +) + + +def _is_valid_js_identifier(name: str) -> bool: + """Return whether a name contains valid JavaScript identifier characters. + + Args: + name: The name to validate. + + Returns: + Whether ``name`` is structurally valid as a JavaScript identifier. + """ + if name.isidentifier(): + return True + if not name or name[0] not in "_$": + return False + return all( + char in "_$" or "a" <= char <= "z" or "A" <= char <= "Z" or "0" <= char <= "9" + for char in name[1:] + ) + + +def _validate_memo_name(name: str) -> None: + """Validate a memo name before adding its JavaScript-safe suffix. + + Args: + name: The unsuffixed JavaScript identifier to validate. + + Raises: + ValueError: If ``name`` contains invalid JavaScript identifier characters. + """ + if not name or not _is_valid_js_identifier(name): + msg = ( + f"`@rx.memo` name {name!r} must contain valid JavaScript identifier " + "characters." + ) + raise ValueError(msg) + + +def _compose_memo_wrappers(outer: Var, inner: Var) -> FunctionVar: + """Compose two memo wrappers while preserving their imports. + + Args: + outer: The wrapper applied last. + inner: The wrapper applied directly to the memo body. + + Returns: + A wrapper equivalent to ``outer(inner(value))``. + """ + return FunctionStringVar.create( + f"(value) => ({outer})(({inner})(value))", + _var_data=VarData.merge( + outer._get_all_var_data(), + inner._get_all_var_data(), + ), + ) + + # 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 @@ -290,8 +372,10 @@ class MemoDefinition: class MemoFunctionDefinition(MemoDefinition): """A memo that compiles to a JavaScript function.""" + export_name: str _function: _LazyBody[ArgsFunctionOperation] imported_var: FunctionVar + wrapper: Var | None = _DEFAULT_FUNCTION_MEMO_WRAPPER @property def function(self) -> ArgsFunctionOperation: @@ -325,6 +409,10 @@ class MemoComponentDefinition(MemoDefinition): # wrapper's ``VarData`` supplies its imports, so a custom wrapper brings # its own and ``None`` pulls in nothing. wrapper: Var | None = DEFAULT_MEMO_WRAPPER + # Whether components in the explicit memo's children are eligible for + # compiler auto-memoization. The default keeps the explicit memo as the + # memoization boundary. + recursive: bool = False # Set for definitions the compiler's auto-memoize pass creates (see # ``create_passthrough_component_memo``). Instances of such a definition # are the auto-memo boundary itself, so the pass must not wrap them again. @@ -386,6 +474,9 @@ class MemoComponent(Component): # introspection (e.g. compile telemetry) can recover the underlying type # without parsing the wrapper's auto-generated class name. _wrapped_component_type: ClassVar[type[Component] | None] = None + # Whether children of an explicit memo are eligible for compiler + # auto-memoization. Auto-generated wrappers always set this to ``True``. + _memo_recursive: ClassVar[bool] = True def _validate_component_children(self, children: list[Component]) -> None: """Skip direct parent/child validation for memo wrapper instances. @@ -430,6 +521,7 @@ def _get_memo_component_class( wrapped_component_type: type[Component] = Component, source_module: str | None = None, auto_memo_wrapper: bool = False, + recursive: bool = False, ) -> type[MemoComponent]: """Get the component subclass for a memo export. @@ -452,6 +544,8 @@ def _get_memo_component_class( boundary, so they opt out of being auto-memoized themselves; user-authored ``@rx.memo`` components do not, so their stateful props land in a generated wrapper instead of the page module. + recursive: Whether components inside an explicit memo are eligible for + compiler auto-memoization. Returns: A cached component subclass with the tag set at class definition time. @@ -465,8 +559,9 @@ def _get_memo_component_class( "tag": symbol, "library": library, "_wrapped_component_type": wrapped_component_type, + "_memo_recursive": auto_memo_wrapper or recursive, } - if auto_memo_wrapper: + if auto_memo_wrapper or not recursive: attrs["_memoization_mode"] = MemoizationMode( disposition=MemoizationDisposition.NEVER ) @@ -512,9 +607,9 @@ def _memo_registry_key(definition: MemoDefinition) -> tuple[str, str | None]: Returns: The ``(name, source_module)`` registry key for the memo. """ - if isinstance(definition, MemoComponentDefinition): + if isinstance(definition, (MemoComponentDefinition, MemoFunctionDefinition)): return definition.export_name, definition.source_module - return definition.python_name, definition.source_module + raise TypeError(type(definition)) def _is_memo_reregistration( @@ -1789,6 +1884,7 @@ def __call__(self, *children: Any, **props: Any) -> MemoComponent: type(component), definition.source_module, definition.auto_memo_wrapper, + definition.recursive, )._create( children=list(children), memo_definition=definition, @@ -2113,21 +2209,36 @@ def _warn_legacy_base_props(fn_name: str, prop_names: Sequence[str]) -> None: def _memo_impl( fn: Callable[..., Any], wrapper: Var | None, + by_value: bool, + recursive: bool, + name: str | 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. + wrapper: The JS wrapper for the emitted function, or ``None`` for no + wrapper. + by_value: Whether to compare props by serialized value instead of + React's default prop identity comparison. + recursive: Whether to auto-memoize hook-bearing components inside the + explicitly memoized component. + name: Optional override for the compiled memo name. Returns: The wrapped function or component factory. Raises: - TypeError: If the return annotation is not supported, or a non-default - ``wrapper`` is given for a var-returning memo. + TypeError: If the return annotation is not supported, ``recursive`` is + enabled for a var-returning memo, or ``by_value`` is combined with + ``wrapper=None``. + ValueError: If the memo name is not a valid JavaScript identifier. """ + memo_name = fn.__name__ if name is None else name + if not isinstance(memo_name, str): + msg = "`@rx.memo` name must be a string." + raise ValueError(msg) + hints = get_type_hints(fn, include_extras=True) return_annotation = hints.get("return", inspect.Signature.empty) missing_return = return_annotation is inspect.Signature.empty @@ -2142,13 +2253,35 @@ def _memo_impl( 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." - ) + export_name = format.to_title_case(memo_name) if is_component else memo_name + _validate_memo_name(export_name) + if name is not None: + export_name += CAMEL_CASE_MEMO_MARKER + if not is_component: + if recursive: + msg = ( + "`@rx.memo` only supports `recursive=True` on component-returning " + f"memos; `{fn.__name__}` returns `rx.Var[...]`." + ) + raise TypeError(msg) + if wrapper is DEFAULT_MEMO_WRAPPER: + wrapper = _DEFAULT_FUNCTION_MEMO_WRAPPER + if by_value and wrapper is None: + msg = "`by_value=True` requires a memo wrapper; it cannot use `wrapper=None`." raise TypeError(msg) + if by_value: + assert wrapper is not None + value_wrapper = ( + BY_VALUE_MEMO_WRAPPER if is_component else _BY_VALUE_FUNCTION_MEMO_WRAPPER + ) + default_wrapper = ( + DEFAULT_MEMO_WRAPPER if is_component else _DEFAULT_FUNCTION_MEMO_WRAPPER + ) + wrapper = ( + value_wrapper + if wrapper is default_wrapper + else _compose_memo_wrappers(value_wrapper, wrapper) + ) defaulted_params: list[str] = [] missing_params: list[str] = [] @@ -2176,10 +2309,10 @@ def _memo_impl( rest_target_fields: set[str] = set() definition = MemoComponentDefinition( fn=fn, - python_name=fn.__name__, + python_name=memo_name, params=params, source_module=source_module, - export_name=format.to_title_case(fn.__name__), + export_name=export_name, _component=_LazyBody( lambda: _evaluate_component_body(fn, params, rest_target_fields), placeholder=Fragment.create(), @@ -2187,20 +2320,23 @@ def _memo_impl( _rest_target_fields=rest_target_fields, _runtime_inferred_params=frozenset(missing_params), wrapper=wrapper, + recursive=recursive, ) memo_callable = _create_component_wrapper(definition) else: definition = MemoFunctionDefinition( fn=fn, - python_name=fn.__name__, + python_name=memo_name, params=params, source_module=source_module, + export_name=export_name, _function=_LazyBody(lambda: _evaluate_function_body(fn, params)), imported_var=_imported_function_var( - fn.__name__, + export_name, _annotation_inner_type(return_annotation), source_module=source_module, ), + wrapper=wrapper, ) memo_callable = _create_function_wrapper(definition) @@ -2225,18 +2361,20 @@ def memo(fn: Callable[..., Var[_MemoVarT]]) -> _MemoFunctionWrapper: ... def memo() -> _MemoDecorator: ... @overload def memo( - *, wrapper: Var | None -) -> Callable[[Callable[..., Component]], _MemoComponentWrapper]: ... + *, + wrapper: Var | None = DEFAULT_MEMO_WRAPPER, + by_value: bool = False, + recursive: bool = False, + name: str | None = None, +) -> _MemoDecorator: ... def memo( fn: Callable[..., Any] | None = None, *, wrapper: Var | None = DEFAULT_MEMO_WRAPPER, -) -> ( - _MemoComponentWrapper - | _MemoFunctionWrapper - | _MemoDecorator - | Callable[[Callable[..., Component]], _MemoComponentWrapper] -): + by_value: bool = False, + recursive: bool = False, + name: str | None = None, +) -> _MemoComponentWrapper | _MemoFunctionWrapper | _MemoDecorator: """Create a memo from a function. The decorated function's body is **not** executed here. Only signature-level @@ -2252,26 +2390,47 @@ def memo( 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. + wrapper: The JS function the compiled component or function is wrapped + in. Component memos default to React's ``memo``; function memos + default to a return-value cache keyed by argument identity. Pass + another ``Var`` (typically an ``rx.vars.FunctionStringVar`` + carrying its own imports) to swap the wrapper, or ``None`` to + export the bare component or function. + by_value: When ``True``, compare serialized props by value instead of + identity for component memos, and cache function results by + serialized argument values. + recursive: When ``True``, allow compiler auto-memoization of + hook-bearing components inside this explicit component memo. Not + supported on function memos. + name: Optional compiled memo name. This is useful when decorating a + lambda, whose default name is ````. 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. + TypeError: If the return annotation is not supported, ``recursive`` is + enabled for a var-returning memo, or ``by_value`` is combined with + ``wrapper=None``. + ValueError: If the memo name is not a valid JavaScript identifier. """ if fn is None: - return cast("_MemoDecorator", partial(_memo_impl, wrapper=wrapper)) - return _memo_impl(fn, wrapper) + return cast( + "_MemoDecorator", + partial( + _memo_impl, + wrapper=wrapper, + by_value=by_value, + recursive=recursive, + name=name, + ), + ) + return _memo_impl(fn, wrapper, by_value, recursive, name) __all__ = [ + "BY_VALUE_MEMO_WRAPPER", "DEFAULT_MEMO_WRAPPER", "EMPTY_VAR_COMPONENT", "MEMOS", diff --git a/pyi_hashes.json b/pyi_hashes.json index 65f8c8ae669..599af41d63b 100644 --- a/pyi_hashes.json +++ b/pyi_hashes.json @@ -120,5 +120,5 @@ "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "f170ac685b6ba5892370166c80684db3", "reflex/__init__.pyi": "a3e1782fab4a9aed55f66cc98af8c217", "reflex/components/__init__.pyi": "9facd05a776d0641432696bbf8e34388", - "reflex/experimental/memo.pyi": "27a73a66e238746e5da5accf99a8fdfd" + "reflex/experimental/memo.pyi": "3a98bc2becb828377484e9192aede170" } diff --git a/reflex/compiler/compiler.py b/reflex/compiler/compiler.py index 2304a74e02f..d871f8788ab 100644 --- a/reflex/compiler/compiler.py +++ b/reflex/compiler/compiler.py @@ -465,15 +465,79 @@ def add_function( # 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. +# coercion. Component and function wrapper imports ride on each definition's +# ``wrapper`` var data instead, so a module whose memos swap or drop their +# wrappers doesn't import them. Shared by the grouped and un-mirrored compile +# paths so they can't drift apart. _MEMO_BASE_IMPORTS: dict[str, list[ImportVar]] = { f"$/{constants.Dirs.STATE_PATH}": [ImportVar(tag="isTrue")], } +def _prepare_recursive_memos( + memos: Iterable[MemoDefinition], +) -> list[tuple[MemoDefinition, Component | None]]: + """Auto-memoize hook-bearing descendants inside recursive memo bodies. + + Args: + memos: The memo definitions requested for compilation. + + Returns: + Definitions paired with an optional compiler-transformed component + body, including any nested auto-memo definitions discovered while + walking recursive bodies. + """ + hooks = CompilerHooks(plugins=(MemoizeStatefulPlugin(),)) + compile_context = CompileContext(pages=(), hooks=hooks) + pending = collections.deque(memos) + queued = { + (type(memo), memo.export_name, memo.source_module) + for memo in pending + if isinstance(memo, (MemoComponentDefinition, MemoFunctionDefinition)) + } + prepared: list[tuple[MemoDefinition, Component | None]] = [] + + with compile_context: + while pending: + memo = pending.popleft() + compiled_body: Component | None = None + if ( + isinstance(memo, MemoComponentDefinition) + and memo.recursive + and not memo.auto_memo_wrapper + ): + page_context = PageContext( + name=memo.python_name, + route=memo.export_name, + root_component=memo.component, + source_module=memo.source_module, + ) + with page_context: + transformed = hooks.compile_component( + memo.component, + page_context=page_context, + compile_context=compile_context, + ) + if not isinstance(transformed, Component): + msg = "A recursive memo body must compile to a Component." + raise TypeError(msg) + compiled_body = transformed + + for generated in compile_context.auto_memo_components.values(): + key = ( + type(generated), + generated.export_name, + generated.source_module, + ) + if key not in queued: + queued.add(key) + pending.append(generated) + + prepared.append((memo, compiled_body)) + + return prepared + + def _compile_memo_components( memos: Iterable[MemoDefinition] = (), ) -> tuple[list[tuple[str, str]], dict[str, list[ImportVar]]]: @@ -512,9 +576,13 @@ def _emit_unmirrored( )) _extend_imports_in_place(aggregate_imports, file_imports) - for memo in memos: + for memo, compiled_body in _prepare_recursive_memos(memos): if isinstance(memo, MemoComponentDefinition): - memo_render, memo_imports = utils.compile_experimental_component_memo(memo) + memo_render, memo_imports = ( + utils.compile_experimental_component_memo(memo) + if compiled_body is None + else utils.compile_experimental_component_memo(memo, compiled_body) + ) segments = memo_paths.module_to_mirrored_segments(memo.source_module) if segments is None: _emit_unmirrored( diff --git a/reflex/compiler/plugins/memoize.py b/reflex/compiler/plugins/memoize.py index a50e3e30569..748037d3923 100644 --- a/reflex/compiler/plugins/memoize.py +++ b/reflex/compiler/plugins/memoize.py @@ -24,7 +24,7 @@ from typing import Any from reflex_base.components.component import BaseComponent, Component -from reflex_base.components.memo import create_passthrough_component_memo +from reflex_base.components.memo import MemoComponent, create_passthrough_component_memo from reflex_base.components.memoize_helpers import ( MemoizationStrategy, _is_structural_memoization_child, @@ -152,7 +152,9 @@ def _should_memoize(component: Component) -> bool: strategy = get_memoization_strategy(component) - if component._memoization_mode.disposition == MemoizationDisposition.NEVER: + if component._memoization_mode.disposition == MemoizationDisposition.NEVER and not ( + isinstance(component, MemoComponent) and not type(component)._memo_recursive + ): return False if isinstance(component, Bare): # A stateful value will be wrapped in a separate component. Match the @@ -198,7 +200,14 @@ def _should_memoize(component: Component) -> bool: if strategy is MemoizationStrategy.SNAPSHOT and not is_snapshot_boundary(component): return True - if is_snapshot_boundary(component) and _subtree_has_reactive_data(component): + # A non-recursive explicit memo is already the boundary for its children. + # Its own reactive props were handled above, but descendants must not make + # the memo itself eligible for another wrapper. + if ( + is_snapshot_boundary(component) + and _subtree_has_reactive_data(component) + and not isinstance(component, MemoComponent) + ): return True # Components with event triggers are always memoized (to wrap callbacks). @@ -268,6 +277,11 @@ def enter_component( return None if page_context.memoize_suppressor_stack: return None + if isinstance(comp, MemoComponent) and not type(comp)._memo_recursive: + # A non-recursive explicit memo remains a passthrough component, + # but its children must not be auto-memoized independently. + page_context.memoize_suppressor_stack.append(id(comp)) + return None strategy = get_memoization_strategy(comp) if strategy is not MemoizationStrategy.SNAPSHOT: return None diff --git a/reflex/compiler/utils.py b/reflex/compiler/utils.py index c9908b5e843..fe84052c9f3 100644 --- a/reflex/compiler/utils.py +++ b/reflex/compiler/utils.py @@ -376,15 +376,18 @@ def _app_style() -> ComponentStyle | Style: def compile_experimental_component_memo( definition: MemoComponentDefinition, + component: Component | None = None, ) -> tuple[dict, ParsedImportDict]: """Compile a memo component. Args: definition: The component memo definition. + component: An optional compiler-transformed memo body. Returns: A tuple of the compiled component definition and its imports. """ + memo_body = definition.component if component is None else component hole_child = definition.passthrough_hole_child if hole_child is not None: # Passthrough memo: shallow-copy the root only — ``render.children`` @@ -393,7 +396,7 @@ def compile_experimental_component_memo( # we skip the O(n) deepcopy + recursive style pass. Descendants are # rendered AND styled in the page scope, not here, so only the root # needs app-level style merged. - render = copy.copy(definition.component) + render = copy.copy(memo_body) _apply_root_style(render) hooks = _root_only_hooks(render) @@ -413,7 +416,7 @@ def compile_experimental_component_memo( render.children = [hole_child] rendered = render.render() else: - render = _apply_component_style_for_compile(copy.deepcopy(definition.component)) + render = _apply_component_style_for_compile(copy.deepcopy(memo_body)) hooks = render._get_all_hooks() rendered = render.render() custom_code = render._get_all_custom_code() @@ -543,20 +546,26 @@ def compile_experimental_function_memo( if var_data := function._get_all_var_data(): # Un-mirrored per-file memo modules live at ``$/utils/components/``; # strip only a self-import to this function memo's own module. - self_module = memo_paths.unmirrored_library_specifier(definition.python_name) + self_module = memo_paths.unmirrored_library_specifier(definition.export_name) imports = { lib: list(fields) for lib, fields in dict(var_data.imports).items() if lib != self_module } + 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) + return ( { "kind": "function", "name": memo_paths.library_and_symbol( - definition.source_module, definition.python_name + definition.source_module, definition.export_name )[1], "function": str(function), + "wrapper": str(wrapper) if wrapper is not None else None, }, imports, ) diff --git a/tests/units/compiler/test_memoize_plugin.py b/tests/units/compiler/test_memoize_plugin.py index c7b8615423d..be41b3d8821 100644 --- a/tests/units/compiler/test_memoize_plugin.py +++ b/tests/units/compiler/test_memoize_plugin.py @@ -11,6 +11,7 @@ from reflex_base.components.component import Component from reflex_base.components.component import field as component_field from reflex_base.components.memo import ( + MEMOS, MemoComponent, MemoComponentDefinition, create_passthrough_component_memo, @@ -645,6 +646,68 @@ def static_card(label: rx.Var[str]) -> Component: assert f"jsx({static_card(label='static').tag}," in page_output +def test_user_memo_recursive_controls_descendant_auto_memoization() -> None: + """Only recursive explicit memos auto-memoize reactive children.""" + + @rx.memo(recursive=False) + def non_recursive_card(children: rx.Var[Component]) -> Component: + return rx.box(children) + + non_recursive_instance = non_recursive_card(WithProp.create(label=STATE_VAR)) + assert ( + non_recursive_instance._memoization_mode.disposition + is MemoizationDisposition.NEVER + ) + + ctx, _page_ctx = _compile_single_page( + lambda: non_recursive_card(WithProp.create(label=STATE_VAR)) + ) + assert not ctx.auto_memo_components + + @rx.memo(recursive=True) + def recursive_card_with_child(children: rx.Var[Component]) -> Component: + return rx.box(children) + + recursive_instance = recursive_card_with_child(WithProp.create(label=STATE_VAR)) + assert ( + recursive_instance._memoization_mode.disposition + is not MemoizationDisposition.NEVER + ) + + ctx, _page_ctx = _compile_single_page( + lambda: recursive_card_with_child(WithProp.create(label=STATE_VAR)) + ) + assert ctx.auto_memo_components + + +def test_recursive_user_memo_auto_memoizes_stateful_body_descendant() -> None: + """A state read authored inside a recursive memo gets a nested boundary.""" + from reflex.compiler.compiler import compile_memo_components + + @rx.memo + def non_recursive_dashboard() -> Component: + return Plain.create(WithProp.create(label=STATE_VAR)) + + non_recursive_definition = MEMOS["NonRecursiveDashboard", __name__] + assert isinstance(non_recursive_definition, MemoComponentDefinition) + files, _imports = compile_memo_components((non_recursive_definition,)) + assert sum(content.count("export const ") for _, content in files) == 1 + + @rx.memo(recursive=True) + def recursive_dashboard() -> Component: + return Plain.create(WithProp.create(label=STATE_VAR)) + + definition = MEMOS["RecursiveDashboard", __name__] + assert isinstance(definition, MemoComponentDefinition) + + files, _imports = compile_memo_components((definition,)) + code = "\n".join(content for _, content in files) + outer_symbol = memo_paths.mirrored_symbol("RecursiveDashboard", __name__) + assert f"export const {outer_symbol} = memo(" in code + assert code.count("export const ") == 2 + assert '.displayName = "WithProp";' in code + + def test_user_memo_event_trigger_usecallback_leaves_page_scope() -> None: """A memo's event-handler prop is memoized inside the generated wrapper. diff --git a/tests/units/components/test_memo.py b/tests/units/components/test_memo.py index 457f12e3744..7264f4ddb58 100644 --- a/tests/units/components/test_memo.py +++ b/tests/units/components/test_memo.py @@ -12,6 +12,8 @@ import pytest from reflex_base.components.component import Component from reflex_base.components.memo import ( + _BY_VALUE_FUNCTION_MEMO_WRAPPER, + _DEFAULT_FUNCTION_MEMO_WRAPPER, _SPECS, DEFAULT_MEMO_WRAPPER, EMPTY_VAR_COMPONENT, @@ -184,10 +186,8 @@ def merge_styles( sym = memo_paths.mirrored_symbol("merge_styles", __name__) files, _ = compiler.compile_memo_components(tuple(MEMOS.values())) code = "\n".join(c for _, c in files) - assert ( - f"export const {sym} = (({{base, ...overrides}}) => ({{...base, ...overrides}}));" - in code - ) + assert f"export const {sym} = " in code + assert "(({base, ...overrides}) => ({...base, ...overrides}))" in code with pytest.raises(TypeError, match="Do not pass `overrides=` directly"): merge_styles(base=base, overrides={"color": "red"}) @@ -275,7 +275,8 @@ def label_slot( sym = memo_paths.mirrored_symbol("label_slot", __name__) files, _ = compiler.compile_memo_components(tuple(MEMOS.values())) code = "\n".join(c for _, c in files) - assert f"export const {sym} = (({{children, label, ...rest}}) => label);" in code + assert f"export const {sym} = " in code + assert "(({children, label, ...rest}) => label)" in code def test_memo_munges_legacy_bare_type_param(): @@ -1193,6 +1194,73 @@ def default_wrapped(label: rx.Var[str]) -> rx.Component: assert any(imp.tag == "memo" for imp in imports.get("react", [])) +def test_component_memo_by_value_uses_prop_equality(): + """``by_value=True`` emits React memo with a value comparator.""" + + @rx.memo(by_value=True) + def value_wrapped(label: rx.Var[str]) -> rx.Component: + return rx.text(label) + + definition = MEMOS["ValueWrapped", __name__] + assert isinstance(definition, MemoComponentDefinition) + + files, imports = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + sym = memo_paths.mirrored_symbol("ValueWrapped", __name__) + assert ( + f"export const {sym} = ((Component) => memo(Component, " + "(prevProps, nextProps) => JSON.stringify(prevProps) === " + "JSON.stringify(nextProps)))" in code + ) + assert any(imp.tag == "memo" for imp in imports.get("react", [])) + + +def test_component_memo_name_overrides_lambda_name(): + """``name=`` provides a stable exported name for lambda memos.""" + named_lambda = rx.memo(name="named_lambda")( + lambda label: rx.text(label), # pyright: ignore[reportUnknownLambdaType] + ) + + definition = MEMOS["NamedLambdaRxMemo", __name__] + assert isinstance(definition, MemoComponentDefinition) + assert definition.python_name == "named_lambda" + assert definition.fn.__name__ == "" + assert isinstance(named_lambda(label="label"), MemoComponent) + + files, _imports = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + sym = memo_paths.mirrored_symbol("NamedLambdaRxMemo", __name__) + assert f"export const {sym} = memo(" in code + + +@pytest.mark.parametrize("memo_name", ["bad name", "bad;name", "1bad"]) +def test_memo_name_must_be_a_valid_js_identifier(memo_name: str): + """Memo names that would produce invalid JavaScript are rejected.""" + with pytest.raises(ValueError, match="valid JavaScript identifier"): + + @rx.memo(name=memo_name) + def invalid_name(label: rx.Var[str]) -> rx.Component: + return rx.text(label) + + +def test_function_memo_name_appends_marker_to_js_keyword(): + """An explicit name gets a suffix that makes JS keywords safe.""" + + @rx.memo(name="await") + def keyword_named(value: rx.Var[int]) -> rx.Var[str]: + return value.to(str) + + definition = MEMOS["awaitRxMemo", __name__] + assert isinstance(definition, MemoFunctionDefinition) + assert definition.python_name == "await" + assert definition.export_name == "awaitRxMemo" + + files, _imports = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + sym = memo_paths.mirrored_symbol("awaitRxMemo", __name__) + assert f"export const {sym} = " in code + + def test_component_memo_wrapper_none_emits_bare_function(): """``@rx.memo(wrapper=None)`` exports the bare function component.""" @@ -1345,15 +1413,120 @@ def test_component_memo_wrapper_none_in_unmirrored_module(): 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=`"): +def test_var_returning_memo_default_wrapper_memoizes_by_identity(): + """Function memos cache results by argument identity by default.""" - @rx.memo(wrapper=None) # pyright: ignore[reportArgumentType] - def format_id(value: rx.Var[int]) -> rx.Var[str]: + @rx.memo + def format_id(value: rx.Var[int]) -> rx.Var[str]: + return value.to(str) + + definition = MEMOS["format_id", __name__] + assert isinstance(definition, MemoFunctionDefinition) + assert definition.wrapper is _DEFAULT_FUNCTION_MEMO_WRAPPER + + files, _imports = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + assert "const cache = new Map()" in code + + +def test_var_returning_memo_wrapper_none_emits_bare_function(): + """``wrapper=None`` disables memoization for function memos.""" + + @rx.memo(wrapper=None) + def format_id(value: rx.Var[int]) -> rx.Var[str]: + return value.to(str) + + definition = MEMOS["format_id", __name__] + assert isinstance(definition, MemoFunctionDefinition) + assert definition.wrapper is None + + files, _imports = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + assert "Object.is" not in code + assert "JSON.stringify" not in code + + +def test_var_returning_memo_by_value_uses_argument_equality(): + """``by_value=True`` caches function results by serialized arguments.""" + + @rx.memo(by_value=True) + def format_value(value: rx.Var[int]) -> rx.Var[str]: + return value.to(str) + + definition = MEMOS["format_value", __name__] + assert isinstance(definition, MemoFunctionDefinition) + assert definition.wrapper is _BY_VALUE_FUNCTION_MEMO_WRAPPER + + files, _imports = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + assert "JSON.stringify(args)" in code + + +def test_by_value_composes_with_an_explicit_wrapper(): + """The value comparator wraps the result of a custom wrapper.""" + track_render = FunctionStringVar.create( + "trackRender", + _var_data=VarData(imports={"my-render-lib": [ImportVar(tag="trackRender")]}), + ) + + @rx.memo(by_value=True, 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 not track_render + assert "trackRender" in str(definition.wrapper) + assert "JSON.stringify" in str(definition.wrapper) + + files, imports = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + assert "trackRender" in code + assert "JSON.stringify(prevProps)" in code + assert any(imp.tag == "trackRender" for imp in imports["my-render-lib"]) + + +def test_var_returning_memo_custom_wrapper_and_by_value_compose(): + """Function memos compose a custom wrapper with value memoization.""" + trace = FunctionStringVar.create( + "trace", + _var_data=VarData(imports={"trace-lib": [ImportVar(tag="trace")]}), + ) + + @rx.memo(by_value=True, wrapper=trace) + def format_value(value: rx.Var[int]) -> rx.Var[str]: + return value.to(str) + + definition = MEMOS["format_value", __name__] + assert isinstance(definition, MemoFunctionDefinition) + assert "trace" in str(definition.wrapper) + assert "JSON.stringify(args)" in str(definition.wrapper) + + files, imports = compiler.compile_memo_components((definition,)) + code = "\n".join(c for _, c in files) + assert "trace" in code + assert "JSON.stringify(args)" in code + assert any(imp.tag == "trace" for imp in imports["trace-lib"]) + + +def test_var_returning_memo_rejects_recursive(): + """``recursive`` applies only to component memos.""" + with pytest.raises(TypeError, match="only supports `recursive=True`"): + + @rx.memo(recursive=True) + def format_value(value: rx.Var[int]) -> rx.Var[str]: return value.to(str) +def test_by_value_rejects_wrapper_none(): + """Value memoization requires a wrapper for either memo kind.""" + with pytest.raises(TypeError, match="requires a memo wrapper"): + + @rx.memo(by_value=True, wrapper=None) + def unwrapped(label: rx.Var[str]) -> rx.Component: + return rx.text(label) + + def test_memo_decorator_parens_form_matches_bare_decorator(): """``@rx.memo()`` with no arguments behaves like bare ``@rx.memo``."""