preserve prop names to some degree with memo (#5777)

* preserve prop names to some degree with memo

* fix tests

Author: adhami3310

reflex/app.py

@@ -1281,12 +1281,12 @@ def memoized_toast_provider():
 
         # Compile custom components.
         (
-            custom_components_output,
-            custom_components_result,
-            custom_components_imports,
-        ) = compiler.compile_components(dict.fromkeys(CUSTOM_COMPONENTS.values()))
-        compile_results.append((custom_components_output, custom_components_result))
-        all_imports.update(custom_components_imports)
+            memo_components_output,
+            memo_components_result,
+            memo_components_imports,
+        ) = compiler.compile_memo_components(dict.fromkeys(CUSTOM_COMPONENTS.values()))
+        compile_results.append((memo_components_output, memo_components_result))
+        all_imports.update(memo_components_imports)
         progress.advance(task)
 
         with console.timing("Collect all imports and app wraps"):

reflex/compiler/compiler.py

@@ -336,7 +336,7 @@ def _compile_component(component: Component | StatefulComponent) -> str:
     return templates.component_template(component=component)
 
 
-def _compile_components(
+def _compile_memo_components(
     components: Iterable[CustomComponent],
 ) -> tuple[str, dict[str, list[ImportVar]]]:
     """Compile the components.
@@ -376,7 +376,7 @@ def _compile_components(
 
     # Compile the components page.
     return (
-        templates.custom_component_template(
+        templates.memo_components_template(
             imports=utils.compile_imports(imports),
             components=component_renders,
             dynamic_imports=sorted(dynamic_imports),
@@ -565,7 +565,7 @@ def compile_page(path: str, component: BaseComponent) -> tuple[str, str]:
     return output_path, code
 
 
-def compile_components(
+def compile_memo_components(
     components: Iterable[CustomComponent],
 ) -> tuple[str, str, dict[str, list[ImportVar]]]:
     """Compile the custom components.
@@ -580,7 +580,7 @@ def compile_components(
     output_path = utils.get_components_path()
 
     # Compile the components.
-    code, imports = _compile_components(components)
+    code, imports = _compile_memo_components(components)
     return output_path, code, imports
 
 

reflex/compiler/templates.py

@@ -8,6 +8,7 @@
 
 from reflex import constants
 from reflex.constants import Hooks
+from reflex.constants.state import CAMEL_CASE_MEMO_MARKER
 from reflex.utils.format import format_state_name, json_dumps
 from reflex.vars.base import VarData
 
@@ -637,7 +638,7 @@ def stateful_components_template(imports: list[_ImportDict], memoized_code: str)
     return f"{imports_str}\n{memoized_code}"
 
 
-def custom_component_template(
+def memo_components_template(
     imports: list[_ImportDict],
     components: list[dict[str, Any]],
     dynamic_imports: Iterable[str],
@@ -661,7 +662,7 @@ def custom_component_template(
     components_code = ""
     for component in components:
         components_code += f"""
-export const {component["name"]} = memo(({{ {", ".join(component.get("props", []))} }}) => {{
+export const {component["name"]} = memo(({{ {",".join([f"{prop}:{prop}{CAMEL_CASE_MEMO_MARKER}" for prop in component.get("props", [])])} }}) => {{
     {_render_hooks(component.get("hooks", {}))}
     return(
         {_RenderUtils.render(component["render"])}

reflex/components/component.py

@@ -37,7 +37,7 @@
     PageNames,
 )
 from reflex.constants.compiler import SpecialAttributes
-from reflex.constants.state import FRONTEND_EVENT_STATE, MEMO_MARKER
+from reflex.constants.state import CAMEL_CASE_MEMO_MARKER, FRONTEND_EVENT_STATE
 from reflex.event import (
     EventCallback,
     EventChain,
@@ -1983,7 +1983,7 @@ def get_args_spec(key: str) -> types.ArgsSpec | Sequence[types.ArgsSpec]:
 
         super()._post_init(
             event_triggers={
-                key + MEMO_MARKER: EventChain.create(
+                key: EventChain.create(
                     value=props[key],
                     args_spec=get_args_spec(key),
                     key=key,
@@ -1994,9 +1994,7 @@ def get_args_spec(key: str) -> types.ArgsSpec | Sequence[types.ArgsSpec]:
         )
 
         to_camel_cased_props = {
-            format.to_camel_case(key + MEMO_MARKER): None
-            for key in props
-            if key not in event_types
+            format.to_camel_case(key): None for key in props if key not in event_types
         }
         self.get_props = lambda: to_camel_cased_props  # pyright: ignore [reportIncompatibleVariableOverride]
 
@@ -2011,7 +2009,7 @@ def get_args_spec(key: str) -> types.ArgsSpec | Sequence[types.ArgsSpec]:
             if key not in props_types:
                 continue
 
-            camel_cased_key = format.to_camel_case(key + MEMO_MARKER)
+            camel_cased_key = format.to_camel_case(key)
 
             # Get the type based on the annotation.
             type_ = props_types[key]
@@ -2101,11 +2099,13 @@ def get_prop_vars(self) -> list[Var | Callable]:
         """
         return [
             Var(
-                _js_expr=name,
+                _js_expr=name + CAMEL_CASE_MEMO_MARKER,
                 _var_type=(prop._var_type if isinstance(prop, Var) else type(prop)),
             ).guess_type()
             if isinstance(prop, Var) or not isinstance(prop, EventChain)
-            else CustomComponent._get_event_spec_from_args_spec(name, prop)
+            else CustomComponent._get_event_spec_from_args_spec(
+                name + CAMEL_CASE_MEMO_MARKER, prop
+            )
             for name, prop in self.props.items()
         ]
 

reflex/constants/state.py

@@ -16,3 +16,4 @@ class StateManagerMode(str, Enum):
 
 FIELD_MARKER = "_rx_state_"
 MEMO_MARKER = "_rx_memo_"
+CAMEL_CASE_MEMO_MARKER = "RxMemo"

tests/units/components/markdown/test_markdown.py

@@ -169,7 +169,7 @@ def test_create_map_fn_var_subclass(cls, fn_body, fn_args, explicit_return, expe
         (
             "code",
             {"codeblock": syntax_highlighter_memoized_component(CodeBlock)},
-            r"""(({node, inline, className, children, ...props}) => { const match = (className || '').match(/language-(?<lang>.*)/); let _language = match ? match[1] : '';  ;             return inline ? (                 jsx(RadixThemesCode,{...props},children)             ) : (                 jsx(CodeBlock,{codeRxMemo:((Array.isArray(children)) ? children.join("\n") : children),languageRxMemo:_language,...props},)             );         })""",
+            r"""(({node, inline, className, children, ...props}) => { const match = (className || '').match(/language-(?<lang>.*)/); let _language = match ? match[1] : '';  ;             return inline ? (                 jsx(RadixThemesCode,{...props},children)             ) : (                 jsx(CodeBlock,{code:((Array.isArray(children)) ? children.join("\n") : children),language:_language,...props},)             );         })""",
         ),
         (
             "code",
@@ -178,7 +178,7 @@ def test_create_map_fn_var_subclass(cls, fn_body, fn_args, explicit_return, expe
                     ShikiHighLevelCodeBlock
                 )
             },
-            r"""(({node, inline, className, children, ...props}) => { const match = (className || '').match(/language-(?<lang>.*)/); let _language = match ? match[1] : '';  ;             return inline ? (                 jsx(RadixThemesCode,{...props},children)             ) : (                 jsx(CodeBlock,{codeRxMemo:((Array.isArray(children)) ? children.join("\n") : children),languageRxMemo:_language,...props},)             );         })""",
+            r"""(({node, inline, className, children, ...props}) => { const match = (className || '').match(/language-(?<lang>.*)/); let _language = match ? match[1] : '';  ;             return inline ? (                 jsx(RadixThemesCode,{...props},children)             ) : (                 jsx(CodeBlock,{code:((Array.isArray(children)) ? children.join("\n") : children),language:_language,...props},)             );         })""",
         ),
     ],
 )

tests/units/components/test_component.py

@@ -825,7 +825,7 @@ def test_create_custom_component(my_component):
     """
     component = rx.memo(my_component)(prop1="test", prop2=1)
     assert component.tag == "MyComponent"
-    assert set(component.get_props()) == {"prop1RxMemo", "prop2RxMemo"}
+    assert set(component.get_props()) == {"prop1", "prop2"}
     assert component.tag in CUSTOM_COMPONENTS