Auto-memoize user @rx.memo components with stateful props

Merged
MA
1mo ago
+292-4

Type of change

  • New feature (non-breaking change which adds functionality)
  • This change requires a documentation update

Description

Extends the compiler's auto-memoization pass to wrap user-authored @rx.memo components when they're called with state-bound props or event handlers. This ensures state hooks compile into a generated wrapper instead of the page module, preventing unnecessary full-page re-renders on state changes.

Key changes:

  1. Auto-memo wrapper generation (reflex/compiler/plugins/memoize.py):

    • Modified _should_memoize() to treat @rx.memo components as eligible for wrapping when bound to state
    • Generated wrappers opt out of being wrapped again via MemoizationDisposition.NEVER

  2. Memo component definition (packages/reflex-base/src/reflex_base/components/memo.py):

    • Added auto_memo_wrapper flag to MemoComponentDefinition to distinguish compiler-generated wrappers from user code
    • Updated _get_memo_component_class() to apply NEVER disposition only to auto-generated wrappers, allowing user @rx.memo components to be wrapped when needed
    • Clarified docstring explaining the interaction between @rx.memo and auto-memoization

  3. Passthrough wrapper creation (packages/reflex-base/src/reflex_base/components/memo.py):

    • Set auto_memo_wrapper=True when creating passthrough wrappers via create_passthrough_component_memo()

  4. Documentation (docs/library/other/memo.md):

    • Added explanation of how state binding at call sites works with auto-memoization

Behavior

Before: A @rx.memo component called with a state Var would pull that state hook into the page module, causing the entire page to re-render on state changes.

After: The compiler detects this pattern and generates a wrapper that holds the state hooks. The page renders the wrapper, which re-renders on state changes and passes the prop to the memoized component. React's memo then skips the component unless the prop value actually changed.

Testing

  • Added 6 comprehensive unit tests covering:
    • Auto-memo wrapper opt-out behavior
    • User memo with stateful props gets wrapped
    • User memo with static props stays inline
    • Event handler props use useCallback in wrapper
    • Children render in page scope, not captured
    • Memo inside foreach renders in foreach body
  • Added 1 integration test (test_memo_stateful_prop_and_children_update) verifying the feature works end-to-end with both stateful props and children

All tests pass with adequate coverage.

Checklist

  • Tests pass with adequate coverage
  • uv run ruff check . and uv run ruff format . clean
  • uv run pyright reflex tests passes
  • pyi_hashes.json updated
  • Documentation updated

https://claude.ai/code/session_01FNfPvVUGGMUQHh63V4Phbe