diff --git a/packages/reflex-base/src/reflex_base/components/component.py b/packages/reflex-base/src/reflex_base/components/component.py index 79b562ddc78..4e79e612987 100644 --- a/packages/reflex-base/src/reflex_base/components/component.py +++ b/packages/reflex-base/src/reflex_base/components/component.py @@ -41,7 +41,7 @@ unwrap_var_annotation, ) from reflex_base.style import Style, format_as_emotion -from reflex_base.utils import console, format, imports, types +from reflex_base.utils import console, format, imports, memo_paths, types from reflex_base.utils.imports import ImportDict, ImportVar, ParsedImportDict from reflex_base.vars import VarData from reflex_base.vars.base import ( @@ -2092,6 +2092,11 @@ class CustomComponent(Component): doc="The props of the component.", default_factory=dict ) + _source_module: str | None = field( + doc="The user-app Python module that defined this memo, used to mirror its compiled JSX path.", + default=None, + ) + def _post_init(self, **kwargs): """Initialize the custom component. @@ -2156,6 +2161,9 @@ def get_args_spec(key: str) -> types.ArgsSpec | Sequence[types.ArgsSpec]: # Set the tag to the name of the function. self.tag = format.to_title_case(self.component_fn.__name__) + if specifier := memo_paths.library_specifier_for(self._source_module): + self.library = specifier + for key, value in props.items(): # Skip kwargs that are not props. if key not in props_types: @@ -2304,11 +2312,15 @@ def _get_all_app_wrap_components( def _register_custom_component( component_fn: Callable[..., Component], + source_module: str | None = None, ): """Register a custom component to be compiled. Args: component_fn: The function that creates the component. + source_module: The user-app Python module that defined the component, + used to mirror its compiled JSX path. ``None`` falls back to the + legacy ``utils/components`` location. Returns: The custom component. @@ -2331,6 +2343,7 @@ def _register_custom_component( dummy_component = CustomComponent._create( children=[], component_fn=component_fn, + _source_module=source_module, **dummy_props, ) if dummy_component.tag is None: @@ -2351,18 +2364,26 @@ def custom_component( Returns: The decorated function. """ + source_module = memo_paths.capture_source_module(component_fn) @wraps(component_fn) def wrapper(*children, **props) -> CustomComponent: # Remove the children from the props. props.pop("children", None) return CustomComponent._create( - children=list(children), component_fn=component_fn, **props + children=list(children), + component_fn=component_fn, + _source_module=source_module, + **props, ) # Register this component so it can be compiled. - dummy_component = _register_custom_component(component_fn) + dummy_component = _register_custom_component(component_fn, source_module) if tag := dummy_component.tag: + import_specifier = ( + memo_paths.library_specifier_for(source_module) + or f"$/{constants.Dirs.UTILS}/components" + ) object.__setattr__( wrapper, "_as_var", @@ -2371,7 +2392,7 @@ def wrapper(*children, **props) -> CustomComponent: _var_type=type[Component], _var_data=VarData( imports={ - f"$/{constants.Dirs.UTILS}/components": [ImportVar(tag=tag)], + import_specifier: [ImportVar(tag=tag)], "@emotion/react": [ ImportVar(tag="jsx"), ], diff --git a/packages/reflex-base/src/reflex_base/plugins/compiler.py b/packages/reflex-base/src/reflex_base/plugins/compiler.py index ecb55a03d92..f85c1ae1bf0 100644 --- a/packages/reflex-base/src/reflex_base/plugins/compiler.py +++ b/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,9 +767,13 @@ class CompileContext(BaseContext): # ``MemoizeStatefulPlugin``). memoize_wrappers: dict[str, None] = dataclasses.field(default_factory=dict) # Compiler-generated experimental memo definitions for auto-memoized - # stateful wrappers. Stored as ``Any`` to keep ``reflex_base`` decoupled - # from ``reflex.experimental.memo``. - auto_memo_components: dict[str, Any] = dataclasses.field(default_factory=dict) + # stateful wrappers. 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 keep + # ``reflex_base`` decoupled from ``reflex.experimental.memo``. + auto_memo_components: dict[tuple[str, str | None], Any] = dataclasses.field( + default_factory=dict + ) def compile( self, diff --git a/packages/reflex-base/src/reflex_base/utils/memo_paths.py b/packages/reflex-base/src/reflex_base/utils/memo_paths.py new file mode 100644 index 00000000000..9ce4b41b3cd --- /dev/null +++ b/packages/reflex-base/src/reflex_base/utils/memo_paths.py @@ -0,0 +1,223 @@ +"""Mirror user-app Python module paths into the compiler's ``.web`` output. + +The compiler uses these helpers to write each memo's compiled JSX to a path +that mirrors its Python source module, instead of bundling everything into +``.web/utils/components.jsx``. This module owns the small set of helpers that: + +- Read ``fn.__module__`` and reject framework / synthetic modules. +- Walk the live frame stack as a fallback for entry points that don't take a + user-supplied callable (notably ``app.add_page(component)`` with a Component + instance). +- Translate a dotted Python module name into mirrored JSX path segments and + the corresponding ``$/...`` library specifier consumed by the import system. +""" + +from __future__ import annotations + +import importlib.util +import inspect +import sys +from collections.abc import Callable +from pathlib import Path + +# Modules whose names start with one of these prefixes are treated as +# framework code and never mirrored. Mirroring them would emit ``.web/reflex/...`` +# files for memos defined inside the framework's own component packages. +_FRAMEWORK_MODULE_PREFIXES = ( + "reflex.", + "reflex_base.", + "reflex_components_", + "reflex_site_shared.", + "reflex_hosting_cli.", + "reflex_docgen.", +) + +# Bare module names that are treated as framework. Prefix matches above use +# trailing dots, so the bare ``reflex`` package itself is matched here. +_FRAMEWORK_MODULE_NAMES = frozenset({ + "reflex", + "reflex_base", + "reflex_site_shared", + "reflex_hosting_cli", + "reflex_docgen", +}) + + +def _is_framework_module(module_name: str) -> bool: + """Whether ``module_name`` belongs to the framework itself. + + Args: + module_name: The dotted module name. + + Returns: + True if the module is part of the framework and should not be + mirrored under ``.web/``. + """ + if module_name in _FRAMEWORK_MODULE_NAMES: + return True + return module_name.startswith(_FRAMEWORK_MODULE_PREFIXES) + + +def capture_source_module(fn: Callable | None) -> str | None: + """Return the user-app module name for ``fn``, or ``None`` if not user code. + + Reads ``fn.__module__`` directly — Python sets this on every function + definition, and it survives re-exports, decorators that ``functools.wraps`` + correctly, and aliasing. Returns ``None`` for ``__main__``, missing + modules, and framework modules. + + Args: + fn: The user callable whose definition module is wanted. + + Returns: + The dotted module name to mirror under ``.web/``, or ``None`` to fall + back to the legacy un-mirrored output path. + """ + if fn is None: + return None + module_name = getattr(fn, "__module__", None) + if not module_name or module_name == "__main__": + return None + if _is_framework_module(module_name): + return None + return module_name + + +def resolve_user_module_from_frame(skip: int = 0) -> str | None: + """Walk the live frame stack and return the first user-app module name. + + Used only as a fallback for ``app.add_page(component)`` when the caller + passed a pre-built ``Component`` instance instead of a callable, so there + is no ``__module__`` to read directly. + + Args: + skip: Number of frames above the immediate caller to skip before + starting the search. Pass ``1`` to ignore the function that is + calling this helper. + + Returns: + The first frame's module name that is not a framework module, or + ``None`` if no suitable frame exists (e.g. running inside a REPL). + """ + frame = inspect.currentframe() + if frame is None: + return None + frame = frame.f_back + for _ in range(skip): + if frame is None: + return None + frame = frame.f_back + while frame is not None: + module_name = frame.f_globals.get("__name__") + if ( + module_name + and module_name != "__main__" + and not _is_framework_module(module_name) + ): + return module_name + frame = frame.f_back + return None + + +def _segment_is_safe(segment: str) -> bool: + """Whether ``segment`` is a path-safe Python identifier-like fragment. + + Args: + segment: A single dotted-module segment. + + Returns: + True if the segment can be used as a directory or filename without + introducing path traversal or platform-specific quirks. + """ + if not segment or segment in {".", ".."}: + return False + return not any(ch in segment for ch in ("/", "\\", ":", "\0")) + + +def module_to_mirrored_segments(module_name: str | None) -> tuple[str, ...] | None: + """Translate a dotted module name to a tuple of mirrored path segments. + + For a *package* (a module whose import resolves to ``__init__.py``), an + extra ``"index"`` segment is appended so the file lives at + ``/index.jsx`` and submodule files can coexist alongside it as + siblings under ``/``. + + Args: + module_name: The dotted Python module name. ``None`` short-circuits. + + Returns: + A tuple of safe path segments to join under ``.web/``, or ``None`` if + the module name is missing, contains unsafe segments, or cannot be + resolved as a package vs. module. + """ + if not module_name: + return None + segments = module_name.split(".") + if not all(_segment_is_safe(seg) for seg in segments): + return None + # Prefer the live module's __file__ over a fresh find_spec lookup. A user + # can switch a module to a package (or back) between hot-reload compiles, + # and importlib re-binds __file__ when that happens — a cached find_spec + # result wouldn't. + origin: str | None = None + module = sys.modules.get(module_name) + if module is not None: + origin = getattr(module, "__file__", None) + if origin is None: + try: + spec = importlib.util.find_spec(module_name) + except (ImportError, ValueError): + spec = None + if spec is not None: + origin = spec.origin + if origin and origin.endswith("__init__.py"): + return (*segments, "index") + return tuple(segments) + + +def library_specifier_for(source_module: str | None) -> str | None: + """Return the ``$/...`` import specifier mirroring ``source_module``, or None. + + Args: + source_module: The dotted module name a memo was defined in. + + Returns: + The ``$/`` specifier, or ``None`` if no source module was + captured or it can't be safely mirrored. + """ + if source_module is None: + return None + segments = module_to_mirrored_segments(source_module) + if segments is None: + return None + return mirrored_library_specifier(segments) + + +def mirrored_jsx_path(web_dir: Path, segments: tuple[str, ...]) -> Path: + """Build the absolute ``.jsx`` path under ``web_dir`` for ``segments``. + + Args: + web_dir: The project's ``.web`` directory. + segments: Mirrored path segments from + :func:`module_to_mirrored_segments`. + + Returns: + The absolute path the compiler should write the memo module to. + """ + return web_dir.joinpath(*segments).with_suffix(".jsx") + + +def mirrored_library_specifier(segments: tuple[str, ...]) -> str: + """Build the ``$/...`` import specifier for mirrored ``segments``. + + The specifier has no extension; Vite resolves the ``.jsx`` automatically. + + Args: + segments: Mirrored path segments from + :func:`module_to_mirrored_segments`. + + Returns: + A ``$/`` prefixed module specifier suitable for use as a + ``Component.library`` value. + """ + return "$/" + "/".join(segments) diff --git a/pyi_hashes.json b/pyi_hashes.json index 801f1660679..f20a7f66bce 100644 --- a/pyi_hashes.json +++ b/pyi_hashes.json @@ -120,5 +120,5 @@ "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "2c5fadcc014056f041cd4d916137d9e7", "reflex/__init__.pyi": "3a9bb8544cbc338ffaf0a5927d9156df", "reflex/components/__init__.pyi": "f39a2af77f438fa243c58c965f19d42e", - "reflex/experimental/memo.pyi": "9946d9b757f7cef5f53d599194d6e50e" + "reflex/experimental/memo.pyi": "ad3685fc293017ebfe2d7803128aaaa8" } diff --git a/reflex/app.py b/reflex/app.py index 65ae0a8235b..a04b574fccf 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -42,7 +42,7 @@ from reflex_base.event.context import EventContext from reflex_base.event.processor import BaseStateEventProcessor, EventProcessor from reflex_base.registry import RegistrationContext -from reflex_base.utils import console +from reflex_base.utils import console, memo_paths from reflex_base.utils.imports import ImportVar from reflex_base.utils.types import ASGIApp, Message, Receive, Scope, Send from reflex_components_core.base.error_boundary import ErrorBoundary @@ -238,6 +238,7 @@ class UnevaluatedPage: on_load: EventType[()] | None = None meta: Sequence[Mapping[str, Any] | Component] = () context: Mapping[str, Any] = dataclasses.field(default_factory=dict) + _source_module: str | None = None def merged_with(self, other: UnevaluatedPage) -> UnevaluatedPage: """Merge the other page into this one. @@ -256,6 +257,9 @@ def merged_with(self, other: UnevaluatedPage) -> UnevaluatedPage: else other.description, on_load=self.on_load if self.on_load is not None else other.on_load, context=self.context if self.context is not None else other.context, + _source_module=self._source_module + if self._source_module is not None + else other._source_module, ) @@ -864,6 +868,13 @@ def add_page( # Check if the route given is valid verify_route_validity(route) + if isinstance(component, Callable): + source_module = memo_paths.capture_source_module(component) + else: + # The user passed a pre-built Component instance — fall back to + # walking the call stack from add_page's caller. + source_module = memo_paths.resolve_user_module_from_frame(skip=1) + unevaluated_page = UnevaluatedPage( component=component, route=route, @@ -873,6 +884,7 @@ def add_page( on_load=on_load, meta=meta, context=context or {}, + _source_module=source_module, ) if route in self._unevaluated_pages: diff --git a/reflex/compiler/compiler.py b/reflex/compiler/compiler.py index 39ff4931a9e..58bed5f1c70 100644 --- a/reflex/compiler/compiler.py +++ b/reflex/compiler/compiler.py @@ -2,6 +2,8 @@ from __future__ import annotations +import collections +import dataclasses import json import sys from collections.abc import Callable, Iterable, Sequence @@ -24,6 +26,7 @@ from reflex_base.environment import environment from reflex_base.plugins import CompileContext, CompilerHooks, PageContext, Plugin from reflex_base.style import SYSTEM_COLOR_MODE +from reflex_base.utils import memo_paths from reflex_base.utils.exceptions import ReflexError from reflex_base.utils.format import to_title_case from reflex_base.utils.imports import ImportVar @@ -393,78 +396,131 @@ def _compile_component(component: Component) -> str: return templates.component_template(component=component) +@dataclasses.dataclass +class _MemoGroup: + """Accumulator for memos that share a mirrored output path.""" + + components: list[dict] = dataclasses.field(default_factory=list) + functions: list[dict] = dataclasses.field(default_factory=list) + imports: dict[str, list[ImportVar]] = dataclasses.field(default_factory=dict) + dynamic_imports: list[str] = dataclasses.field(default_factory=list) + custom_code: list[str] = dataclasses.field(default_factory=list) + + def add_component( + self, render: dict, memo_imports: dict[str, list[ImportVar]] + ) -> None: + self.components.append(render) + _extend_imports_in_place(self.imports, memo_imports) + self.dynamic_imports.extend(sorted(render.get("dynamic_imports", []) or [])) + self.custom_code.extend(render.get("custom_code", []) or []) + + def add_function( + self, render: dict, memo_imports: dict[str, list[ImportVar]] + ) -> None: + self.functions.append(render) + _extend_imports_in_place(self.imports, memo_imports) + + def _compile_memo_components( components: Iterable[CustomComponent], experimental_memos: Iterable[ExperimentalMemoDefinition] = (), ) -> tuple[list[tuple[str, str]], dict[str, list[ImportVar]]]: - """Compile each memo/custom-component as its own module plus an index. - - Each memo lands in ``.web//.jsx`` with only the imports - it actually uses. Experimental memo wrappers declare their ``library`` as - that per-memo file path so page-side imports resolve directly to the - individual module. + """Compile memos grouped by their source module's mirrored output path. - The ``$/utils/components`` index only re-exports the legacy - ``@rx.memo`` custom components, which are the ones app-level code - (``root.jsx``) imports by name. Keeping experimental memos out of the - index is what lets root's ``import * as utils_components`` avoid - transitively dragging every page-specific memo into the always-loaded - chunk — the tree-shaking win of per-memo files relies on that. + Memos that captured a user-app source module land in a single combined + file at ``.web//.jsx`` so the page-side import surface + matches the source layout. Memos without a source module keep the legacy + behavior — one file per memo at ``.web/utils/components/.jsx`` + re-exported through the ``$/utils/components`` index. Args: components: The components to compile. experimental_memos: The experimental memos to compile. Returns: - A list of ``(path, code)`` pairs to write — one per memo plus one - index — and the aggregated imports across all memo modules. + A list of ``(path, code)`` pairs to write and the aggregated imports + across all memo modules. """ - per_memo_files: list[tuple[str, str]] = [] - # Only legacy custom components go through the index: they are the ones - # root.jsx/custom code imports by name from ``$/utils/components``. - # Experimental memos declare their library per-file (see - # ``_get_experimental_memo_component_class``) so pages import them - # directly and the index stays small. - index_entries: list[tuple[str, str]] = [] + output_files: list[tuple[str, str]] = [] aggregate_imports: dict[str, list[ImportVar]] = {} + legacy_files: list[tuple[str, str]] = [] + legacy_index_entries: list[tuple[str, str]] = [] + legacy_base_dir = utils.get_memo_components_dir() + groups: collections.defaultdict[tuple[str, ...], _MemoGroup] = ( + collections.defaultdict(_MemoGroup) + ) - base_dir = utils.get_memo_components_dir() + def _emit_legacy_component(render: dict, render_imports: dict) -> None: + code, file_imports = _compile_single_memo_component(render, render_imports) + name = render["name"] + legacy_files.append((_memo_component_file_path(legacy_base_dir, name), code)) + _extend_imports_in_place(aggregate_imports, file_imports) + + def _emit_legacy_function(render: dict, render_imports: dict) -> None: + code, file_imports = _compile_single_memo_function(render, render_imports) + legacy_files.append(( + _memo_component_file_path(legacy_base_dir, render["name"]), + code, + )) + _extend_imports_in_place(aggregate_imports, file_imports) for component in components: component_render, component_imports = utils.compile_custom_component(component) - name = component_render["name"] - code, file_imports = _compile_single_memo_component( - component_render, component_imports - ) - path = _memo_component_file_path(base_dir, name) - specifier = _memo_component_index_specifier(name) - per_memo_files.append((path, code)) - index_entries.append((name, specifier)) - _extend_imports_in_place(aggregate_imports, file_imports) + segments = memo_paths.module_to_mirrored_segments(component._source_module) + if segments is None: + name = component_render["name"] + _emit_legacy_component(component_render, component_imports) + legacy_index_entries.append((name, _memo_component_index_specifier(name))) + else: + groups[segments].add_component(component_render, component_imports) for memo in experimental_memos: if isinstance(memo, ExperimentalMemoComponentDefinition): memo_render, memo_imports = utils.compile_experimental_component_memo(memo) - name = memo_render["name"] - code, file_imports = _compile_single_memo_component( - memo_render, memo_imports - ) - path = _memo_component_file_path(base_dir, name) - per_memo_files.append((path, code)) - _extend_imports_in_place(aggregate_imports, file_imports) + segments = memo_paths.module_to_mirrored_segments(memo.source_module) + if segments is None: + _emit_legacy_component(memo_render, memo_imports) + else: + groups[segments].add_component(memo_render, memo_imports) elif isinstance(memo, ExperimentalMemoFunctionDefinition): memo_render, memo_imports = utils.compile_experimental_function_memo(memo) - name = memo_render["name"] - code, file_imports = _compile_single_memo_function( - memo_render, memo_imports - ) - path = _memo_component_file_path(base_dir, name) - per_memo_files.append((path, code)) - _extend_imports_in_place(aggregate_imports, file_imports) + segments = memo_paths.module_to_mirrored_segments(memo.source_module) + if segments is None: + _emit_legacy_function(memo_render, memo_imports) + else: + groups[segments].add_function(memo_render, memo_imports) + + framework_imports: dict[str, list[ImportVar]] = { + "react": [ImportVar(tag="memo")], + f"$/{constants.Dirs.STATE_PATH}": [ImportVar(tag="isTrue")], + } + if groups: + _extend_imports_in_place(aggregate_imports, framework_imports) + _apply_common_imports(aggregate_imports) + + for segments, group in groups.items(): + merged_imports = utils.merge_imports(framework_imports, group.imports) + _apply_common_imports(merged_imports) + # Strip self-imports — when memos in this group reference each other, + # their compiled imports point at this group's own mirrored specifier. + # Importing from the same file would shadow the module's own exports. + merged_imports.pop(memo_paths.mirrored_library_specifier(segments), None) + code = templates.memo_components_template( + imports=utils.compile_imports(merged_imports), + components=group.components, + functions=group.functions, + dynamic_imports=sorted(set(group.dynamic_imports)), + custom_codes=list(dict.fromkeys(group.custom_code)), + ) + output_files.append((utils.get_memo_module_path(segments), code)) + _extend_imports_in_place(aggregate_imports, group.imports) index_path = utils.get_components_path() - index_code = templates.memo_index_template(index_entries) - return [(index_path, index_code), *per_memo_files], aggregate_imports + index_code = templates.memo_index_template(legacy_index_entries) + return ( + [(index_path, index_code), *legacy_files, *output_files], + aggregate_imports, + ) def _compile_single_memo_component( @@ -1132,6 +1188,7 @@ def compile_app( ) compile_results.extend(memo_component_files) all_imports = utils.merge_imports(all_imports, memo_components_imports) + utils.prune_stale_memo_files(path for path, _ in memo_component_files) progress.advance(task) compile_results.append( diff --git a/reflex/compiler/plugins/builtin.py b/reflex/compiler/plugins/builtin.py index a4b326be4ab..c75d0807fd0 100644 --- a/reflex/compiler/plugins/builtin.py +++ b/reflex/compiler/plugins/builtin.py @@ -64,6 +64,7 @@ def eval_page( name=getattr(page_fn, "__name__", page.route), route=page.route, root_component=component, + source_module=getattr(page, "_source_module", None), ) diff --git a/reflex/compiler/plugins/memoize.py b/reflex/compiler/plugins/memoize.py index b596b147a79..e1789113c8a 100644 --- a/reflex/compiler/plugins/memoize.py +++ b/reflex/compiler/plugins/memoize.py @@ -383,8 +383,12 @@ def _build_wrapper( compile_context.memoize_wrappers[tag] = None # Passthrough memo definitions capture app-specific event/state vars, so # they must be rebuilt for each compile instead of shared globally. - wrapper_factory, definition = create_passthrough_component_memo(tag, comp) - compile_context.auto_memo_components[tag] = definition + wrapper_factory, definition = create_passthrough_component_memo( + tag, comp, source_module=page_context.source_module + ) + compile_context.auto_memo_components[ + definition.export_name, definition.source_module + ] = definition wrapper = wrapper_factory() # The wrapper has no structural children at the page level, but parents diff --git a/reflex/compiler/utils.py b/reflex/compiler/utils.py index c2bdf618650..37cc98e6838 100644 --- a/reflex/compiler/utils.py +++ b/reflex/compiler/utils.py @@ -5,9 +5,12 @@ import asyncio import concurrent.futures import copy +import json import operator +import os +import tempfile import traceback -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime from pathlib import Path from typing import Any, TypedDict @@ -17,7 +20,7 @@ from reflex_base.components.component import Component, ComponentStyle, CustomComponent from reflex_base.constants.state import CAMEL_CASE_MEMO_MARKER, FIELD_MARKER from reflex_base.style import Style -from reflex_base.utils import format, imports +from reflex_base.utils import format, imports, memo_paths from reflex_base.utils.imports import ImportVar, ParsedImportDict from reflex_base.vars.base import Field, Var, VarData from reflex_base.vars.function import DestructuredArg @@ -798,6 +801,19 @@ def get_memo_components_dir() -> str: ) +def get_memo_module_path(segments: tuple[str, ...]) -> str: + """Get the on-disk path for a memo module mirrored from a Python module. + + Args: + segments: Mirrored path segments produced by + :func:`reflex_base.utils.memo_paths.module_to_mirrored_segments`. + + Returns: + The absolute path the compiler should write the combined memo file to. + """ + return str(memo_paths.mirrored_jsx_path(get_web_dir(), segments)) + + def add_meta( page: Component, title: str, @@ -862,6 +878,99 @@ def write_file(path: str | Path, code: str): path.write_text(code, encoding="utf-8") +_MEMO_MANIFEST_FILENAME = ".memo-manifest.json" + + +def _read_memo_manifest(web_dir: Path) -> set[str]: + """Read the previous compile's memo file manifest. + + Args: + web_dir: The project's ``.web`` directory. + + Returns: + The set of paths (relative to ``.web``) recorded by the previous + compile, or an empty set if the manifest is absent or invalid. + """ + manifest_path = web_dir / _MEMO_MANIFEST_FILENAME + if not manifest_path.exists(): + return set() + try: + data = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return set() + if not isinstance(data, list): + return set() + return {entry for entry in data if isinstance(entry, str)} + + +def _write_memo_manifest(web_dir: Path, relative_paths: set[str]) -> None: + """Atomically write the new memo file manifest. + + Args: + web_dir: The project's ``.web`` directory. + relative_paths: Paths emitted this run, relative to ``.web``. + """ + manifest_path = web_dir / _MEMO_MANIFEST_FILENAME + web_dir.mkdir(parents=True, exist_ok=True) + fd, tmp_name = tempfile.mkstemp( + prefix=".memo-manifest.", suffix=".json.tmp", dir=str(web_dir) + ) + # Close the raw fd immediately and reopen the file by path. Wrapping the + # fd via os.fdopen() would leak it if the wrap itself raised. + os.close(fd) + tmp_path = Path(tmp_name) + try: + with tmp_path.open("w", encoding="utf-8") as fh: + json.dump(sorted(relative_paths), fh) + tmp_path.replace(manifest_path) + except Exception: + # Best-effort cleanup; manifest write is recoverable on the next run. + tmp_path.unlink(missing_ok=True) + raise + + +def prune_stale_memo_files(emitted_paths: Iterable[str | Path]) -> None: + """Delete memo files written previously that this compile no longer emits. + + Only paths that appear in the previous manifest are considered for + deletion — never a fresh filesystem walk — so files this code did not + emit are never touched. Empty parent directories created by mirrored + output are removed up to (but not including) the ``.web`` root. + + Args: + emitted_paths: Absolute (or ``.web``-relative) paths the current + compile produced for the memo pipeline. + """ + web_dir = get_web_dir() + + emitted_relative: set[str] = set() + for path in emitted_paths: + absolute = Path(path) + if not absolute.is_absolute(): + absolute = web_dir / absolute + try: + relative = absolute.relative_to(web_dir) + except ValueError: + continue + emitted_relative.add(str(relative).replace(os.sep, "/")) + + previous = _read_memo_manifest(web_dir) + for relative in previous - emitted_relative: + target = web_dir / relative + if target.is_file(): + target.unlink() + parent = target.parent + while parent != web_dir and parent.is_relative_to(web_dir): + try: + parent.rmdir() + except OSError: + break + parent = parent.parent + + if emitted_relative != previous: + _write_memo_manifest(web_dir, emitted_relative) + + def empty_dir(path: str | Path, keep_files: list[str] | None = None): """Remove all files and folders in a directory except for the keep_files. diff --git a/reflex/experimental/memo.py b/reflex/experimental/memo.py index 0dbf91bfbf4..c2378b6877c 100644 --- a/reflex/experimental/memo.py +++ b/reflex/experimental/memo.py @@ -22,7 +22,7 @@ SpecialAttributes, ) from reflex_base.constants.state import CAMEL_CASE_MEMO_MARKER -from reflex_base.utils import format +from reflex_base.utils import format, memo_paths from reflex_base.utils.imports import ImportVar from reflex_base.utils.types import safe_issubclass from reflex_base.vars import VarData @@ -62,6 +62,7 @@ class ExperimentalMemoDefinition: fn: Callable[..., Any] python_name: str params: tuple[MemoParam, ...] + source_module: str | None = dataclasses.field(default=None, kw_only=True) @dataclasses.dataclass(frozen=True, slots=True) @@ -150,6 +151,7 @@ def _post_init(self, **kwargs): def _get_experimental_memo_component_class( export_name: str, wrapped_component_type: type[Component] = Component, + source_module: str | None = None, ) -> type[ExperimentalMemoComponent]: """Get the component subclass for an experimental memo export. @@ -164,18 +166,23 @@ def _get_experimental_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. """ + # Per-file fallback gives Vite distinct module boundaries per memo, enabling + # actual code-split by page when no source module is available. + 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/`` 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, } if ( wrapped_component_type._get_app_wrap_components @@ -339,24 +346,29 @@ def _get_rest_param(params: tuple[MemoParam, ...]) -> MemoParam | None: return next((param for param in params if param.is_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 an experimental 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)]}), ) @@ -637,12 +649,14 @@ def _analyze_params( def _create_function_definition( fn: Callable[..., Any], return_annotation: Any, + source_module: str | None = None, ) -> ExperimentalMemoFunctionDefinition: """Create a definition for a var-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 function memo definition. @@ -677,9 +691,12 @@ def _create_function_definition( fn=fn, python_name=fn.__name__, params=params, + source_module=source_module, function=function, imported_var=_imported_function_var( - fn.__name__, _annotation_inner_type(return_annotation) + fn.__name__, + _annotation_inner_type(return_annotation), + source_module=source_module, ), ) @@ -687,12 +704,14 @@ def _create_function_definition( def _create_component_definition( fn: Callable[..., Any], return_annotation: Any, + source_module: str | None = None, ) -> ExperimentalMemoComponentDefinition: """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. @@ -713,6 +732,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), ) @@ -962,7 +982,9 @@ def __call__(self, *children: Any, **props: Any) -> ExperimentalMemoComponent: # Build the component props passed into the memo wrapper. return _get_experimental_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, @@ -1010,6 +1032,7 @@ def _create_component_wrapper( def create_passthrough_component_memo( export_name: str, component: Component, + source_module: str | None = None, ) -> tuple[ Callable[..., ExperimentalMemoComponent], ExperimentalMemoComponentDefinition, @@ -1023,6 +1046,8 @@ def create_passthrough_component_memo( Args: export_name: The exported memo component name. 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. @@ -1056,7 +1081,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 != export_name: replacements["export_name"] = export_name @@ -1089,13 +1114,15 @@ def memo(fn: Callable[..., Any]) -> Callable[..., Any]: ) raise TypeError(msg) + source_module = memo_paths.capture_source_module(fn) + if _is_component_annotation(return_annotation): - definition = _create_component_definition(fn, return_annotation) + definition = _create_component_definition(fn, return_annotation, source_module) _register_memo_definition(definition) return _create_component_wrapper(definition) if _is_var_annotation(return_annotation): - definition = _create_function_definition(fn, return_annotation) + definition = _create_function_definition(fn, return_annotation, source_module) _register_memo_definition(definition) return _create_function_wrapper(definition) diff --git a/reflex/utils/memo_paths.py b/reflex/utils/memo_paths.py new file mode 100644 index 00000000000..31100a76d8e --- /dev/null +++ b/reflex/utils/memo_paths.py @@ -0,0 +1,24 @@ +"""Re-export of ``reflex_base.utils.memo_paths``. + +The helpers live in ``reflex_base`` so the lower-level component package can +use them at decoration time. This module exists so callers in the top-level +``reflex`` package can import them from a familiar location. +""" + +from reflex_base.utils.memo_paths import ( + capture_source_module, + library_specifier_for, + mirrored_jsx_path, + mirrored_library_specifier, + module_to_mirrored_segments, + resolve_user_module_from_frame, +) + +__all__ = [ + "capture_source_module", + "library_specifier_for", + "mirrored_jsx_path", + "mirrored_library_specifier", + "module_to_mirrored_segments", + "resolve_user_module_from_frame", +] diff --git a/tests/units/compiler/test_memoize_plugin.py b/tests/units/compiler/test_memoize_plugin.py index 18fc3b3e5dc..bd552c92f55 100644 --- a/tests/units/compiler/test_memoize_plugin.py +++ b/tests/units/compiler/test_memoize_plugin.py @@ -85,6 +85,7 @@ class FakePage: description: Any = None image: str = "" meta: tuple[dict[str, Any], ...] = () + _source_module: str | None = None def _compile_single_page( @@ -169,7 +170,7 @@ def test_memoize_wrapper_uses_experimental_memo_component_and_call_site() -> Non assert len(ctx.memoize_wrappers) == 1 wrapper_tag = next(iter(ctx.memoize_wrappers)) - assert wrapper_tag in ctx.auto_memo_components + assert (wrapper_tag, None) in ctx.auto_memo_components output = page_ctx.output_code or "" assert f'import {{{wrapper_tag}}} from "$/utils/components/{wrapper_tag}"' in output assert f"jsx({wrapper_tag}," in (page_ctx.output_code or "") @@ -186,7 +187,7 @@ def test_memoize_wrapper_deduped_across_repeated_subtrees() -> None: ) assert len(ctx.memoize_wrappers) == 1 wrapper_tag = next(iter(ctx.memoize_wrappers)) - assert list(ctx.auto_memo_components) == [wrapper_tag] + assert list(ctx.auto_memo_components) == [(wrapper_tag, None)] assert (page_ctx.output_code or "").count( f'import {{{wrapper_tag}}} from "$/utils/components/{wrapper_tag}"' ) == 1 @@ -369,8 +370,13 @@ def test_passthrough_memo_definitions_are_not_shared_globally(monkeypatch) -> No def fake_create_passthrough_component_memo( export_name: str, component: Component, + source_module: str | None = None, ): - definition = SimpleNamespace(export_name=export_name, component=component) + definition = SimpleNamespace( + export_name=export_name, + component=component, + source_module=source_module, + ) return (lambda definition=definition: definition), definition monkeypatch.setattr( @@ -381,7 +387,7 @@ def fake_create_passthrough_component_memo( first_compile = SimpleNamespace(memoize_wrappers={}, auto_memo_components={}) second_compile = SimpleNamespace(memoize_wrappers={}, auto_memo_components={}) - page_context = cast(PageContext, SimpleNamespace()) + page_context = cast(PageContext, SimpleNamespace(source_module=None)) MemoizeStatefulPlugin._build_wrapper( first_component, @@ -394,8 +400,8 @@ def fake_create_passthrough_component_memo( compile_context=second_compile, ) - first_definition = first_compile.auto_memo_components[tag] - second_definition = second_compile.auto_memo_components[tag] + first_definition = first_compile.auto_memo_components[tag, None] + second_definition = second_compile.auto_memo_components[tag, None] assert first_definition.component is first_component assert second_definition.component is second_component assert second_definition is not first_definition @@ -415,13 +421,75 @@ def test_shared_subtree_across_pages_uses_same_tag() -> None: assert len(ctx.memoize_wrappers) == 1 tag = next(iter(ctx.memoize_wrappers)) - assert list(ctx.auto_memo_components) == [tag] + assert list(ctx.auto_memo_components) == [(tag, None)] for route in ("/a", "/b"): output = ctx.compiled_pages[route].output_code or "" assert f'import {{{tag}}} from "$/utils/components/{tag}"' in output assert f"jsx({tag}," in output +def test_shared_subtree_in_distinct_source_modules_emits_per_module() -> None: + """Identical subtrees in different user modules emit one memo per module. + + Regression: the auto-memo registry was keyed by tag only, so when two + pages from different user modules produced the same memoizable subtree + (and therefore the same tag), the second registration overwrote the + first. Only the surviving definition's source module got a mirrored + memo file, leaving the other page importing a tag from a file that + never exported it. + """ + from reflex.compiler.compiler import compile_memo_components + + ctx = CompileContext( + pages=[ + FakePage( + route="/a", + component=lambda: Plain.create(STATE_VAR), + _source_module="memo_collision_test.module_a", + ), + FakePage( + route="/b", + component=lambda: Plain.create(STATE_VAR), + _source_module="memo_collision_test.module_b", + ), + ], + hooks=CompilerHooks(plugins=default_page_plugins()), + ) + with ctx: + ctx.compile() + + assert len(ctx.memoize_wrappers) == 1 + tag = next(iter(ctx.memoize_wrappers)) + # Both source modules survive registration as distinct entries. + assert set(ctx.auto_memo_components) == { + (tag, "memo_collision_test.module_a"), + (tag, "memo_collision_test.module_b"), + } + + page_a_output = ctx.compiled_pages["/a"].output_code or "" + page_b_output = ctx.compiled_pages["/b"].output_code or "" + assert f'import {{{tag}}} from "$/memo_collision_test/module_a"' in page_a_output + assert f'import {{{tag}}} from "$/memo_collision_test/module_b"' in page_b_output + + output_files, _ = compile_memo_components( + components=(), + experimental_memos=tuple(ctx.auto_memo_components.values()), + ) + emitted = {Path(path).as_posix(): code for path, code in output_files} + + def find_emitted(suffix: str) -> str | None: + return next( + (code for path, code in emitted.items() if path.endswith(suffix)), None + ) + + matched_a = find_emitted("memo_collision_test/module_a.jsx") + matched_b = find_emitted("memo_collision_test/module_b.jsx") + assert matched_a is not None, f"missing module_a memo file in {sorted(emitted)}" + assert matched_b is not None, f"missing module_b memo file in {sorted(emitted)}" + assert f"export const {tag} = memo" in matched_a + assert f"export const {tag} = memo" in matched_b + + def test_shared_parent_instance_across_pages_preserves_original() -> None: """A parent instance reused across pages must not have its children rebound. diff --git a/tests/units/compiler/test_plugins.py b/tests/units/compiler/test_plugins.py index 26eb1f39c99..1103e37aa07 100644 --- a/tests/units/compiler/test_plugins.py +++ b/tests/units/compiler/test_plugins.py @@ -1057,7 +1057,7 @@ def test_compile_context_memoize_wrappers_registers_shared_subtree_tag() -> None # Both pages share the same subtree hash, so exactly one wrapper tag is registered. assert len(compile_ctx.memoize_wrappers) == 1 wrapper_tag = next(iter(compile_ctx.memoize_wrappers)) - assert list(compile_ctx.auto_memo_components) == [wrapper_tag] + assert [key for key, _ in compile_ctx.auto_memo_components] == [wrapper_tag] # Each page imports the generated experimental memo component. page_a_code = compile_ctx.compiled_pages["/a"].output_code or "" assert ( @@ -1079,7 +1079,7 @@ def test_compile_context_resets_memoize_wrappers_between_runs() -> None: with ctx: ctx.compile() first_tags = set(ctx.memoize_wrappers) - first_defs = set(ctx.auto_memo_components) + first_defs = {tag for tag, _ in ctx.auto_memo_components} assert first_tags # memoize wrapper was registered assert first_defs == first_tags @@ -1093,7 +1093,7 @@ def test_compile_context_resets_memoize_wrappers_between_runs() -> None: # Same shared component → same tag, not a union across runs. assert set(ctx2.memoize_wrappers) == first_tags - assert set(ctx2.auto_memo_components) == first_tags + assert {tag for tag, _ in ctx2.auto_memo_components} == first_tags page_ctx = ctx2.compiled_pages["/c"] assert "react-moment" in page_ctx.frontend_imports assert "$/utils/stateful_components" not in (page_ctx.output_code or "") diff --git a/tests/units/compiler/test_stale_cleanup.py b/tests/units/compiler/test_stale_cleanup.py new file mode 100644 index 00000000000..51b23487873 --- /dev/null +++ b/tests/units/compiler/test_stale_cleanup.py @@ -0,0 +1,117 @@ +"""Tests for the memo manifest-driven stale file cleanup.""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import patch + +import pytest + +from reflex.compiler import utils as compiler_utils + + +@pytest.fixture +def fake_web_dir(tmp_path: Path): + """Pretend tmp_path is the project's .web directory. + + Args: + tmp_path: The pytest tmp directory. + + Yields: + The path used as ``.web`` for the duration of the test. + """ + web_dir = tmp_path / ".web" + web_dir.mkdir() + with patch.object(compiler_utils, "get_web_dir", return_value=web_dir): + yield web_dir + + +def _seed_manifest(web_dir: Path, paths: list[str]) -> None: + (web_dir / compiler_utils._MEMO_MANIFEST_FILENAME).write_text( + json.dumps(paths), encoding="utf-8" + ) + + +def _seed_files(web_dir: Path, relative_paths: list[str]) -> list[Path]: + written: list[Path] = [] + for rel in relative_paths: + target = web_dir / rel + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text("// memo", encoding="utf-8") + written.append(target) + return written + + +def test_prune_removes_files_dropped_between_runs(fake_web_dir: Path): + survivor, removed = _seed_files( + fake_web_dir, + ["myapp/widgets/buttons.jsx", "myapp/dropped.jsx"], + ) + _seed_manifest(fake_web_dir, ["myapp/widgets/buttons.jsx", "myapp/dropped.jsx"]) + + compiler_utils.prune_stale_memo_files([survivor]) + + assert survivor.exists() + assert not removed.exists() + + +def test_prune_cleans_empty_parent_dirs(fake_web_dir: Path): + survivor, _orphan = _seed_files( + fake_web_dir, + ["myapp/keep.jsx", "myapp/widgets/orphan.jsx"], + ) + _seed_manifest(fake_web_dir, ["myapp/keep.jsx", "myapp/widgets/orphan.jsx"]) + + compiler_utils.prune_stale_memo_files([survivor]) + + assert not (fake_web_dir / "myapp" / "widgets").exists() + assert (fake_web_dir / "myapp").exists() + + +def test_prune_only_touches_manifest_paths(fake_web_dir: Path): + untouched = fake_web_dir / "user_added.jsx" + untouched.write_text("// stays", encoding="utf-8") + [survivor] = _seed_files(fake_web_dir, ["myapp/keep.jsx"]) + # Manifest only mentions the survivor — even if other files exist next to + # it, prune must never delete files outside the manifest's history. + _seed_manifest(fake_web_dir, ["myapp/keep.jsx"]) + + compiler_utils.prune_stale_memo_files([survivor]) + + assert untouched.exists() + assert survivor.exists() + + +def test_prune_writes_new_manifest(fake_web_dir: Path): + [survivor] = _seed_files(fake_web_dir, ["myapp/widgets/buttons.jsx"]) + _seed_manifest(fake_web_dir, []) + + compiler_utils.prune_stale_memo_files([survivor]) + + manifest = json.loads( + (fake_web_dir / compiler_utils._MEMO_MANIFEST_FILENAME).read_text( + encoding="utf-8" + ) + ) + assert manifest == ["myapp/widgets/buttons.jsx"] + + +def test_prune_handles_missing_previous_manifest(fake_web_dir: Path): + [survivor] = _seed_files(fake_web_dir, ["myapp/widgets/buttons.jsx"]) + + # No manifest seeded — should not raise and should still write one. + compiler_utils.prune_stale_memo_files([survivor]) + + assert (fake_web_dir / compiler_utils._MEMO_MANIFEST_FILENAME).exists() + + +def test_prune_ignores_corrupt_manifest(fake_web_dir: Path): + (fake_web_dir / compiler_utils._MEMO_MANIFEST_FILENAME).write_text( + "not json", encoding="utf-8" + ) + [survivor] = _seed_files(fake_web_dir, ["myapp/widgets/buttons.jsx"]) + + compiler_utils.prune_stale_memo_files([survivor]) + + assert survivor.exists() diff --git a/tests/units/components/test_component.py b/tests/units/components/test_component.py index 9ca3495a2a2..7a98a463a8b 100644 --- a/tests/units/components/test_component.py +++ b/tests/units/components/test_component.py @@ -883,6 +883,30 @@ def test_custom_component_hash(my_component): assert {component1, component2} == {component1} +def test_custom_component_captures_source_module(my_component): + """The user module that defined the component is captured on the instance. + + Args: + my_component: A test custom component. + """ + component = rx.memo(my_component)(prop1="test", prop2=1) + assert component._source_module == my_component.__module__ + assert component.library == "$/" + my_component.__module__.replace(".", "/") + + +def test_custom_component_as_var_uses_mirrored_specifier(my_component): + """The wrapper's _as_var imports from the mirrored library specifier. + + Args: + my_component: A test custom component. + """ + wrapper = rx.memo(my_component) + var_data = wrapper._as_var()._var_data # pyright: ignore [reportFunctionMemberAccess] + expected_specifier = "$/" + my_component.__module__.replace(".", "/") + libraries = {lib for lib, _ in var_data.imports} + assert expected_specifier in libraries + + def test_custom_component_wrapper(): """Test that the wrapper of a custom component is correct.""" @@ -1791,8 +1815,10 @@ def outer(): # The imports are only resolved during compilation. _, imports_outer = compile_custom_component(outer_comp) assert "inner" not in imports_outer - assert "$/utils/components" in imports_outer - assert imports_outer["$/utils/components"] == [ImportVar(tag="Wrapper")] + # Memos defined in this module land at the mirrored library specifier. + expected_specifier = "$/" + __name__.replace(".", "/") + assert expected_specifier in imports_outer + assert imports_outer[expected_specifier] == [ImportVar(tag="Wrapper")] def test_custom_component_declare_event_handlers_in_fields(): diff --git a/tests/units/experimental/test_memo.py b/tests/units/experimental/test_memo.py index efb006d545d..321c68bfd09 100644 --- a/tests/units/experimental/test_memo.py +++ b/tests/units/experimental/test_memo.py @@ -2,6 +2,7 @@ from __future__ import annotations +from pathlib import Path from types import SimpleNamespace from typing import Any @@ -389,6 +390,56 @@ def my_card(children: rx.Var[rx.Component], *, title: rx.Var[str]) -> rx.Compone assert "export const MyCard = memo(" in code +def test_compile_memo_components_groups_by_source_module(): + """Memos sharing a source module are concatenated into one mirrored file.""" + mirrored_segments = __name__.split(".") + + @rx.memo + def grouped_first(title: rx.Var[str]) -> rx.Component: + return rx.text(title) + + @rx.memo + def grouped_second(title: rx.Var[str]) -> rx.Component: + return rx.heading(title) + + files, _ = compiler.compile_memo_components( + dict.fromkeys(CUSTOM_COMPONENTS.values()), + tuple(EXPERIMENTAL_MEMOS.values()), + ) + expected_suffix = (Path(*mirrored_segments).with_suffix(".jsx")).as_posix() + + grouped_files = [ + (path, code) + for path, code in files + if Path(path).as_posix().endswith("/" + expected_suffix) + ] + assert len(grouped_files) == 1 + code = grouped_files[0][1] + assert "export const GroupedFirst = memo(" in code + assert "export const GroupedSecond = memo(" in code + # The merged module must carry imports its memos use, not just the + # framework-level ones added by the compiler. + assert "RadixThemesText" in code + assert "RadixThemesHeading" in code + + +def test_compile_memo_components_falls_back_when_no_source_module(): + """Memos with no source module continue to populate the legacy index.""" + legacy_definition = ExperimentalMemoComponentDefinition( + fn=lambda: None, + python_name="legacy_memo", + params=(), + source_module=None, + export_name="LegacyMemo", + component=rx.fragment(), + passthrough_hole_child=None, + ) + + files, _ = compiler.compile_memo_components((), (legacy_definition,)) + paths = [Path(path).as_posix() for path, _ in files] + assert any(path.endswith("utils/components/LegacyMemo.jsx") for path in paths) + + def test_compile_memo_components_extends_imports_without_remerging( monkeypatch: pytest.MonkeyPatch, ): @@ -402,6 +453,7 @@ def noop() -> None: fn=noop, python_name=f"memo_{idx}", params=(), + source_module=None, export_name=f"Memo{idx}", component=rx.fragment(), passthrough_hole_child=None, diff --git a/tests/units/utils/test_memo_paths.py b/tests/units/utils/test_memo_paths.py new file mode 100644 index 00000000000..558edbf176d --- /dev/null +++ b/tests/units/utils/test_memo_paths.py @@ -0,0 +1,88 @@ +"""Tests for source-module capture and mirrored-path translation.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from unittest.mock import patch + +from reflex_base.utils import memo_paths + + +def _user_fn(): + """Stand-in for a user-defined function (defined in this test module).""" + + +def test_capture_source_module_returns_user_module(): + captured = memo_paths.capture_source_module(_user_fn) + # Module name depends on how pytest collected this module; either is fine + # so long as it isn't filtered. + assert captured is not None + assert "test_memo_paths" in captured + + +def test_capture_source_module_filters_framework(): + assert memo_paths.capture_source_module(memo_paths.capture_source_module) is None + + +def test_capture_source_module_filters_main(): + fn = type("F", (), {"__module__": "__main__"}) + assert memo_paths.capture_source_module(fn) is None + + +def test_capture_source_module_filters_missing(): + assert memo_paths.capture_source_module(None) is None + fn = type("F", (), {"__module__": ""}) + assert memo_paths.capture_source_module(fn) is None + + +def test_module_to_mirrored_segments_module(): + spec = importlib.util.find_spec("reflex.experimental.memo") + # Ensure the test runs against a real, non-package module. + assert spec is not None + assert spec.origin + assert not spec.origin.endswith("__init__.py") + segments = memo_paths.module_to_mirrored_segments("reflex.experimental.memo") + assert segments == ("reflex", "experimental", "memo") + + +def test_module_to_mirrored_segments_package_appends_index(): + # `reflex.experimental` is a package — its __init__.py origin should + # cause an "index" segment to be appended. + segments = memo_paths.module_to_mirrored_segments("reflex.experimental") + assert segments == ("reflex", "experimental", "index") + + +def test_module_to_mirrored_segments_unsafe_segment_rejected(): + assert memo_paths.module_to_mirrored_segments("foo..bar") is None + assert memo_paths.module_to_mirrored_segments("foo./bar") is None + assert memo_paths.module_to_mirrored_segments("..secret") is None + + +def test_module_to_mirrored_segments_none(): + assert memo_paths.module_to_mirrored_segments(None) is None + assert memo_paths.module_to_mirrored_segments("") is None + + +def test_mirrored_jsx_path_joins_segments(): + web_dir = Path("/tmp/.web") + path = memo_paths.mirrored_jsx_path(web_dir, ("counter_app", "ui", "buttons")) + assert path == web_dir / "counter_app" / "ui" / "buttons.jsx" + + +def test_mirrored_library_specifier_joins_with_slash(): + spec = memo_paths.mirrored_library_specifier(("counter_app", "ui", "buttons")) + assert spec == "$/counter_app/ui/buttons" + + +def test_resolve_user_module_from_frame_skips_framework(): + captured = memo_paths.resolve_user_module_from_frame() + # Must find a user frame — the test module itself qualifies. + assert captured is not None + assert not captured.startswith("reflex") + + +def test_resolve_user_module_from_frame_returns_none_inside_framework_only(): + # Simulate a stack of framework-only frames. + with patch.object(memo_paths, "_is_framework_module", return_value=True): + assert memo_paths.resolve_user_module_from_frame() is None