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 compiler's auto-memo registry is scoped per source module, so identical-rendering subtrees in different modules each emit their own output instead of one silently overwriting another, 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/compiler/templates.py

@@ -8,6 +8,7 @@
 
 from reflex_base import constants
 from reflex_base.constants import Hooks
+from reflex_base.utils import memo_paths
 from reflex_base.utils.format import format_state_name, json_dumps
 from reflex_base.vars.base import VarData
 
@@ -192,7 +193,7 @@ def app_root_template(
     if hydrate_fallback_export is not None:
         hydrate_fallback_str = (
             f"export {{ {hydrate_fallback_export} as HydrateFallback }} "
-            f'from "$/{constants.Dirs.COMPONENTS_PATH}/{hydrate_fallback_export}";'
+            f'from "{memo_paths.unmirrored_library_specifier(hydrate_fallback_export)}";'
         )
 
     custom_code_str = "\n".join(custom_codes)

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

@@ -39,7 +39,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
@@ -253,6 +253,12 @@ 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 per-name
+    # ``app_components/_internal/<name>`` path used for memos that can't be
+    # mirrored. ``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)
@@ -300,7 +306,7 @@ def component(self) -> Component:
 class MemoComponent(Component):
     """A rendered instance of a memo component."""
 
-    library = f"$/{constants.Dirs.COMPONENTS_PATH}"
+    library = f"$/{constants.Dirs.APP_COMPONENTS_INTERNAL}"
     _memoization_mode = MemoizationMode(disposition=MemoizationDisposition.NEVER)
 
     # The user-authored component class this wrapper stands in for. Populated
@@ -346,6 +352,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.
 
@@ -360,18 +367,21 @@ 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 per-name ``app_components/_internal/<name>`` path.
 
     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_for(source_module, 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 (
@@ -388,6 +398,18 @@ def _get_memo_component_class(
     )
 
 
+def reset_memo_component_classes() -> None:
+    """Clear the cached memo wrapper classes.
+
+    Called at the start of each compile so a memo's ``library`` is recomputed
+    from the current module layout. Without this, a module that switches to a
+    package (or back) between hot-reload compiles would keep serving the
+    library specifier resolved on the first compile, pointing pages at an
+    output path the compiler no longer writes.
+    """
+    _get_memo_component_class.cache_clear()
+
+
 MEMOS: dict[str, MemoDefinition] = {}
 
 
@@ -627,42 +649,48 @@ 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
+            per-name ``app_components/_internal/<name>`` path.
 
     Returns:
         The imported FunctionVar.
     """
+    library = memo_paths.library_for(source_module, 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
+            per-name ``app_components/_internal/<name>`` path.
 
     Returns:
         The component var.
     """
+    library = memo_paths.library_for(source_module, 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")],
             }
         ),
@@ -1311,12 +1339,14 @@ def _evaluate_function_body(
 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.
@@ -1329,6 +1359,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=_LazyBody.ready(_evaluate_component_body(fn, params)),
     )
@@ -1594,7 +1625,9 @@ def __call__(self, *children: Any, **props: Any) -> MemoComponent:
         # Reading ``component`` materializes the deferred body, so ``type(...)``
         # reflects the real wrapped class rather than the placeholder.
         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,
@@ -1608,7 +1641,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(
@@ -1641,6 +1676,7 @@ def _create_component_wrapper(
 
 def create_passthrough_component_memo(
     component: Component,
+    source_module: str | None = None,
 ) -> tuple[
     Callable[..., MemoComponent],
     MemoComponentDefinition,
@@ -1659,6 +1695,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.
@@ -1726,7 +1764,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
@@ -1871,6 +1909,8 @@ def memo(fn: Callable[..., Any]) -> _MemoComponentWrapper | _MemoFunctionWrapper
         defaulted_params=defaulted_params,
     )
 
+    source_module = memo_paths.capture_source_module(fn)
+
     if missing_return or defaulted_params:
         _warn_missing_annotations(fn.__name__, missing_return, defaulted_params)
 
@@ -1885,6 +1925,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=_LazyBody(
                 lambda: _evaluate_component_body(fn, params),
@@ -1897,9 +1938,12 @@ def memo(fn: Callable[..., Any]) -> _MemoComponentWrapper | _MemoFunctionWrapper
             fn=fn,
             python_name=fn.__name__,
             params=params,
+            source_module=source_module,
             _function=_LazyBody(lambda: _evaluate_function_body(fn, params)),
             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/constants/base.py

@@ -32,8 +32,6 @@ class Dirs(SimpleNamespace):
     UTILS = "utils"
     # The name of the state file.
     STATE_PATH = UTILS + "/state"
-    # The name of the components file.
-    COMPONENTS_PATH = UTILS + "/components"
     # The name of the contexts file.
     CONTEXTS_PATH = UTILS + "/context"
     # The name of the output directory.
@@ -48,6 +46,14 @@ class Dirs(SimpleNamespace):
     PAGES = "app"
     # The name of the routes directory.
     ROUTES = "routes"
+    # Subdirectory holding memo modules mirrored from user Python modules,
+    # kept separate from framework output so user module paths can't collide.
+    APP_COMPONENTS = "app_components"
+    # Reserved subdir for memos that can't be mirrored to a user module
+    # (framework components, ``__main__``, unsafe names). The leading underscore
+    # keeps it clear of mirrored user-module trees, since top-level Python
+    # packages conventionally never start with ``_``.
+    APP_COMPONENTS_INTERNAL = APP_COMPONENTS + "/_internal"
     # The name of the env json file.
     ENV_JSON = "env.json"
     # The name of the reflex json file.

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,

packages/reflex-base/src/reflex_base/utils/imports.py

@@ -6,6 +6,16 @@
 from collections import defaultdict
 from collections.abc import Mapping, Sequence
 
+# Absolute import paths beginning with one of these reserved ``.web``
+# subdirectories are rewritten to ``$``-prefixed module specifiers.
+ABSOLUTE_IMPORT_PREFIXES = (
+    "/utils/",
+    "/components/",
+    "/styles/",
+    "/public/",
+    "/app_components/",
+)
+
 
 def merge_parsed_imports(
     *imports: ImmutableParsedImportDict,
@@ -42,11 +52,7 @@ def merge_imports(
             import_dict if isinstance(import_dict, tuple) else import_dict.items()
         ):
             # If the lib is an absolute path, we need to prefix it with a $
-            lib = (
-                "$" + lib
-                if lib.startswith(("/utils/", "/components/", "/styles/", "/public/"))
-                else lib
-            )
+            lib = "$" + lib if lib.startswith(ABSOLUTE_IMPORT_PREFIXES) else lib
             if isinstance(fields, (list, tuple, set)):
                 all_imports[lib].extend(
                     ImportVar(field) if isinstance(field, str) else field