Dedupe custom codes emitted by rx.memo components (#4793)

* Dedupe custom codes emitted by rx.memo components

* Update variable names

* Add test case for multiple memo components with custom code

Author: masenf

reflex/.templates/jinja/web/pages/custom_component.js.jinja2

@@ -1,11 +1,14 @@
 {% extends "web/pages/base_page.js.jinja2" %}
 {% from "web/pages/macros.js.jinja2" import renderHooks %}
-{% block export %}
-{% for component in components %}
 
-{% for custom_code in component.custom_code %}
+{% block declaration %}
+{% for custom_code in custom_codes %}
 {{custom_code}}
 {% endfor %}
+{% endblock %}
+
+{% block export %}
+{% for component in components %}
 
 export const {{component.name}} = memo(({ {{-component.props|join(", ")-}} }) => {
     {{ renderHooks(component.hooks) }}

reflex/compiler/compiler.py

@@ -247,12 +247,19 @@ def _compile_components(
         for comp_import in comp_render["dynamic_imports"]
     }
 
+    custom_codes = {
+        comp_custom_code: None
+        for comp_render in component_renders
+        for comp_custom_code in comp_render.get("custom_code", [])
+    }
+
     # Compile the components page.
     return (
         templates.COMPONENTS.render(
             imports=utils.compile_imports(imports),
             components=component_renders,
             dynamic_imports=dynamic_imports,
+            custom_codes=custom_codes,
         ),
         imports,
     )

tests/integration/test_memo.py

@@ -0,0 +1,67 @@
+"""Integration tests for rx.memo components."""
+
+from typing import Generator
+
+import pytest
+from selenium.webdriver.common.by import By
+
+from reflex.testing import AppHarness
+
+
+def MemoApp():
+    """Reflex app with memo components."""
+    import reflex as rx
+
+    class FooComponent(rx.Fragment):
+        def add_custom_code(self) -> list[str]:
+            return [
+                "const foo = 'bar'",
+            ]
+
+    @rx.memo
+    def foo_component(t: str):
+        return FooComponent.create(t, rx.Var("foo"))
+
+    @rx.memo
+    def foo_component2(t: str):
+        return FooComponent.create(t, rx.Var("foo"))
+
+    def index() -> rx.Component:
+        return rx.vstack(
+            foo_component(t="foo"), foo_component2(t="bar"), id="memo-custom-code"
+        )
+
+    app = rx.App()
+    app.add_page(index)
+
+
+@pytest.fixture()
+def memo_app(tmp_path) -> Generator[AppHarness, None, None]:
+    """Start MemoApp app at tmp_path via AppHarness.
+
+    Args:
+        tmp_path: pytest tmp_path fixture
+
+    Yields:
+        running AppHarness instance
+    """
+    with AppHarness.create(
+        root=tmp_path,
+        app_source=MemoApp,
+    ) as harness:
+        yield harness
+
+
+@pytest.mark.asyncio
+async def test_memo_app(memo_app: AppHarness):
+    """Render various memo'd components and assert on the output.
+
+    Args:
+        memo_app: harness for MemoApp app
+    """
+    assert memo_app.app_instance is not None, "app is not running"
+    driver = memo_app.frontend()
+
+    # check that the output matches
+    memo_custom_code_stack = driver.find_element(By.ID, "memo-custom-code")
+    assert memo_custom_code_stack.text == "foobarbarbar"