feat(compiler): mirror memo output paths to Python source modules

news/6457.feature.md

@@ -0,0 +1 @@
+Auto-memoized (`rx.memo`) components now compile to `.web` output paths that mirror their defining Python source module instead of being bundled into a single shared `components.jsx`. The memo registry is scoped per source module, so same-named components in different modules no longer collide, hot-reloads of a module refresh the correct output, and stale memo files are cleaned up when their source changes.

packages/reflex-base/news/6457.feature.md

@@ -0,0 +1 @@
+Added `reflex_base.utils.memo_paths`, which translates a memo's Python source module into the mirrored `.web` JSX path and `$/...` library specifier used by the compiler. The memo component and compiler plugin now route each memo's compiled output through these helpers so it lands alongside its source module's layout.

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

@@ -37,7 +37,7 @@
 )
 from reflex_base.constants.state import CAMEL_CASE_MEMO_MARKER
 from reflex_base.event import EventChain, EventHandler, no_args_event_spec, run_script
-from reflex_base.utils import console, format
+from reflex_base.utils import console, format, memo_paths
 from reflex_base.utils.imports import ImportVar
 from reflex_base.utils.types import safe_issubclass, typehint_issubclass
 from reflex_base.vars import VarData
@@ -150,6 +150,11 @@ class MemoDefinition:
     fn: Callable[..., Any]
     python_name: str
     params: tuple[MemoParam, ...]
+    # The user-app Python module that defined this memo. When set, the memo's
+    # compiled JSX is emitted to a path mirroring that module and the page-side
+    # import resolves there instead of the legacy per-name ``utils/components``
+    # path. ``kw_only`` so subclasses can keep their own required fields.
+    source_module: str | None = dataclasses.field(default=None, kw_only=True)
 
 
 @dataclasses.dataclass(frozen=True, slots=True)
@@ -225,6 +230,7 @@ def _post_init(self, **kwargs):
 def _get_memo_component_class(
     export_name: str,
     wrapped_component_type: type[Component] = Component,
+    source_module: str | None = None,
 ) -> type[MemoComponent]:
     """Get the component subclass for a memo export.
 
@@ -239,18 +245,24 @@ def _get_memo_component_class(
         wrapped_component_type: The class of the component being memoized.
             Defaults to ``Component`` for memos that don't wrap a user
             component (e.g. function memos, raw passthroughs).
+        source_module: The user-app Python module that defined this memo. When
+            set, the wrapper imports from a path mirroring that module instead
+            of the legacy per-name path under ``utils/components``.
 
     Returns:
         A cached component subclass with the tag set at class definition time.
     """
+    # With a source module the memo is grouped into a file mirroring its
+    # Python module; otherwise each memo gets its own per-file module so Vite
+    # has distinct module boundaries per memo, enabling code-split by page.
+    library = (
+        memo_paths.library_specifier_for(source_module)
+        or f"$/{constants.Dirs.COMPONENTS_PATH}/{export_name}"
+    )
     attrs: dict[str, Any] = {
         "__module__": __name__,
         "tag": export_name,
-        # Point each memo at its own per-file module so pages import directly
-        # from ``$/utils/components/<name>`` rather than through the index.
-        # Per-file import paths give Vite distinct module boundaries per
-        # memo, enabling actual code-split by page.
-        "library": f"$/{constants.Dirs.COMPONENTS_PATH}/{export_name}",
+        "library": library,
         "_wrapped_component_type": wrapped_component_type,
     }
     if (
@@ -458,42 +470,54 @@ def _get_rest_param(params: tuple[MemoParam, ...]) -> MemoParam | None:
     return next((p for p in params if p.kind is MemoParamKind.REST), None)
 
 
-def _imported_function_var(name: str, return_type: Any) -> FunctionVar:
+def _imported_function_var(
+    name: str, return_type: Any, source_module: str | None = None
+) -> FunctionVar:
     """Create the imported FunctionVar for a memo.
 
     Args:
         name: The exported function name.
         return_type: The return type of the function.
+        source_module: The user-app Python module that defined the memo. When
+            set, the import resolves to the mirrored module file instead of the
+            legacy per-name path.
 
     Returns:
         The imported FunctionVar.
     """
+    library = (
+        memo_paths.library_specifier_for(source_module)
+        or f"$/{constants.Dirs.COMPONENTS_PATH}/{name}"
+    )
     return FunctionStringVar.create(
         name,
         _var_type=ReflexCallable[Any, return_type],
-        _var_data=VarData(
-            imports={
-                f"$/{constants.Dirs.COMPONENTS_PATH}/{name}": [ImportVar(tag=name)]
-            }
-        ),
+        _var_data=VarData(imports={library: [ImportVar(tag=name)]}),
     )
 
 
-def _component_import_var(name: str) -> Var:
+def _component_import_var(name: str, source_module: str | None = None) -> Var:
     """Create the imported component var for a memo component.
 
     Args:
         name: The exported component name.
+        source_module: The user-app Python module that defined the memo. When
+            set, the import resolves to the mirrored module file instead of the
+            legacy per-name path.
 
     Returns:
         The component var.
     """
+    library = (
+        memo_paths.library_specifier_for(source_module)
+        or f"$/{constants.Dirs.COMPONENTS_PATH}/{name}"
+    )
     return Var(
         name,
         _var_type=type[Component],
         _var_data=VarData(
             imports={
-                f"$/{constants.Dirs.COMPONENTS_PATH}/{name}": [ImportVar(tag=name)],
+                library: [ImportVar(tag=name)],
                 "@emotion/react": [ImportVar(tag="jsx")],
             }
         ),
@@ -1082,12 +1106,14 @@ def _build_args_function(
 def _create_component_definition(
     fn: Callable[..., Any],
     return_annotation: Any,
+    source_module: str | None = None,
 ) -> MemoComponentDefinition:
     """Create a definition for a component-returning memo.
 
     Args:
         fn: The function to analyze.
         return_annotation: The return annotation.
+        source_module: The user-app Python module that defined the memo.
 
     Returns:
         The component memo definition.
@@ -1108,6 +1134,7 @@ def _create_component_definition(
         fn=fn,
         python_name=fn.__name__,
         params=params,
+        source_module=source_module,
         export_name=format.to_title_case(fn.__name__),
         component=_lift_rest_props(component),
     )
@@ -1358,7 +1385,9 @@ def __call__(self, *children: Any, **props: Any) -> MemoComponent:
 
         # Build the component props passed into the memo wrapper.
         return _get_memo_component_class(
-            definition.export_name, type(definition.component)
+            definition.export_name,
+            type(definition.component),
+            definition.source_module,
         )._create(
             children=list(children),
             memo_definition=definition,
@@ -1372,7 +1401,9 @@ def _as_var(self) -> Var:
         Returns:
             The imported component var.
         """
-        return _component_import_var(self._definition.export_name)
+        return _component_import_var(
+            self._definition.export_name, self._definition.source_module
+        )
 
 
 def _create_function_wrapper(
@@ -1405,6 +1436,7 @@ def _create_component_wrapper(
 
 def create_passthrough_component_memo(
     component: Component,
+    source_module: str | None = None,
 ) -> tuple[
     Callable[..., MemoComponent],
     MemoComponentDefinition,
@@ -1423,6 +1455,8 @@ def create_passthrough_component_memo(
 
     Args:
         component: The component to wrap.
+        source_module: The user-app Python module that triggered creation of
+            this memo (typically the page that contained the wrapped subtree).
 
     Returns:
         The callable memo wrapper and its component definition.
@@ -1490,7 +1524,7 @@ def passthrough(children: Var[Component]) -> Component:
     passthrough.__qualname__ = passthrough.__name__
     passthrough.__module__ = __name__
 
-    definition = _create_component_definition(passthrough, Component)
+    definition = _create_component_definition(passthrough, Component, source_module)
     replacements: dict[str, Any] = {}
     if definition.export_name != tag:
         replacements["export_name"] = tag
@@ -1698,6 +1732,8 @@ def memo(fn: Callable[..., Any]) -> _MemoComponentWrapper | _MemoFunctionWrapper
         defaulted_params=defaulted_params,
     )
 
+    source_module = memo_paths.capture_source_module(fn)
+
     # Construct the wrapper against a placeholder body so the user's body can
     # self-reference the memo during eager evaluation; the real body is patched
     # in after eval completes (see `_bind_self_reference`).
@@ -1707,6 +1743,7 @@ def memo(fn: Callable[..., Any]) -> _MemoComponentWrapper | _MemoFunctionWrapper
             fn=fn,
             python_name=fn.__name__,
             params=params,
+            source_module=source_module,
             export_name=format.to_title_case(fn.__name__),
             component=Fragment.create(),
         )
@@ -1716,11 +1753,14 @@ def memo(fn: Callable[..., Any]) -> _MemoComponentWrapper | _MemoFunctionWrapper
             fn=fn,
             python_name=fn.__name__,
             params=params,
+            source_module=source_module,
             function=ArgsFunctionOperation.create(
                 args_names=(), return_expr=LiteralVar.create(None)
             ),
             imported_var=_imported_function_var(
-                fn.__name__, _annotation_inner_type(return_annotation)
+                fn.__name__,
+                _annotation_inner_type(return_annotation),
+                source_module=source_module,
             ),
         )
         wrapper = _create_function_wrapper(definition)

packages/reflex-base/src/reflex_base/plugins/compiler.py

@@ -689,6 +689,7 @@ class PageContext(BaseContext):
     frontend_imports: ParsedImportDict = dataclasses.field(default_factory=dict)
     output_path: str | None = None
     output_code: str | None = None
+    source_module: str | None = None
     # Stack of ``id(component)`` for components whose subtree is
     # memoize-suppressed. Populated by ``MemoizeStatefulPlugin`` when it
     # encounters a ``MemoizationLeaf``-style snapshot boundary and popped on
@@ -766,8 +767,13 @@ class CompileContext(BaseContext):
     # ``MemoizeStatefulPlugin``).
     memoize_wrappers: dict[str, None] = dataclasses.field(default_factory=dict)
     # Compiler-generated memo definitions for auto-memoized stateful wrappers.
-    # Stored as ``Any`` to avoid an import cycle with ``reflex_base.components.memo``.
-    auto_memo_components: dict[str, Any] = dataclasses.field(default_factory=dict)
+    # Keyed by ``(tag, source_module)`` so identical-rendering subtrees from
+    # different user modules each get their own entry and emit into the right
+    # mirrored memo file. Stored as ``Any`` to avoid an import cycle with
+    # ``reflex_base.components.memo``.
+    auto_memo_components: dict[tuple[str, str | None], Any] = dataclasses.field(
+        default_factory=dict
+    )
 
     def compile(
         self,