From 96f48bc013cdbcda66d1763e7079c71d2fb258ec Mon Sep 17 00:00:00 2001
From: Juro Oravec <juraj.oravec.josefson@gmail.com>
Date: Sat, 1 Feb 2025 17:39:56 +0100
Subject: [PATCH] refactor: fix on_render_after hook (#941)

---
 src/django_components/component.py          |  28 +++--
 src/django_components/perfutil/component.py | 124 +++++++++++++-----
 tests/test_component.py                     | 133 +++++++++++++++++++-
 3 files changed, 241 insertions(+), 44 deletions(-)

diff --git a/src/django_components/component.py b/src/django_components/component.py
index d83c24ff..4ffdc1fe 100644
--- a/src/django_components/component.py
+++ b/src/django_components/component.py
@@ -232,6 +232,10 @@ class ComponentContext:
     fills: Dict[SlotName, Slot]
     outer_context: Optional[Context]
     registry: ComponentRegistry
+    # When we render a component, the root component, together with all the nested Components,
+    # shares this dictionary for storing callbacks that are called from within `component_post_render`.
+    # This is so that we can pass them all in when the root component is passed to `component_post_render`.
+    post_render_callbacks: Dict[str, Callable[[str], str]]
 
 
 class Component(
@@ -1028,9 +1032,11 @@ class Component(
             parent_id = cast(str, context[_COMPONENT_CONTEXT_KEY])
             parent_comp_ctx = component_context_cache[parent_id]
             component_path = [*parent_comp_ctx.component_path, self.name]
+            post_render_callbacks = parent_comp_ctx.post_render_callbacks
         else:
             parent_id = None
             component_path = [self.name]
+            post_render_callbacks = {}
 
         trace_component_msg(
             "COMP_PREP_START",
@@ -1061,6 +1067,7 @@ class Component(
             default_slot=None,
             outer_context=snapshot_context(self.outer_context) if self.outer_context is not None else None,
             registry=self.registry,
+            post_render_callbacks=post_render_callbacks,
         )
 
         # Instead of passing the ComponentContext directly through the Context, the entry on the Context
@@ -1137,9 +1144,18 @@ class Component(
         )
 
         # Remove component from caches
-        def on_component_rendered(component_id: str) -> None:
-            del component_context_cache[component_id]  # type: ignore[arg-type]
-            unregister_provide_reference(component_id)  # type: ignore[arg-type]
+        def on_component_rendered(html: str) -> str:
+            with self._with_metadata(metadata):
+                # Allow to optionally override/modify the rendered content
+                new_output = self.on_render_after(context_snapshot, template, html)
+                html = new_output if new_output is not None else html
+
+            del component_context_cache[render_id]  # type: ignore[arg-type]
+            unregister_provide_reference(render_id)  # type: ignore[arg-type]
+
+            return html
+
+        post_render_callbacks[render_id] = on_component_rendered
 
         # After the component and all its children are rendered, we resolve
         # all inserted HTML comments into <script> and <link> tags (if render_dependencies=True)
@@ -1161,7 +1177,7 @@ class Component(
             render_id=render_id,
             component_name=self.name,
             parent_id=parent_id,
-            on_component_rendered=on_component_rendered,
+            on_component_rendered_callbacks=post_render_callbacks,
             on_html_rendered=on_html_rendered,
         )
 
@@ -1217,10 +1233,6 @@ class Component(
                 # Get the component's HTML
                 html_content = template.render(context)
 
-                # Allow to optionally override/modify the rendered content
-                new_output = component.on_render_after(context, template, html_content)
-                html_content = new_output if new_output is not None else html_content
-
             # Add necessary HTML attributes to work with JS and CSS variables
             updated_html, child_components = set_component_attrs_for_js_and_css(
                 html_content=html_content,
diff --git a/src/django_components/perfutil/component.py b/src/django_components/perfutil/component.py
index 14c24156..21f1cdcc 100644
--- a/src/django_components/perfutil/component.py
+++ b/src/django_components/perfutil/component.py
@@ -30,10 +30,18 @@ component_context_cache: Dict[str, "ComponentContext"] = {}
 
 class PostRenderQueueItem(NamedTuple):
     content_before_component: str
-    component_id: Optional[str]
+    child_id: Optional[str]
     parent_id: Optional[str]
+    grandparent_id: Optional[str]
     component_name_path: List[str]
 
+    def __repr__(self) -> str:
+        return (
+            f"PostRenderQueueItem(child_id={self.child_id!r}, parent_id={self.parent_id!r}, "
+            f"grandparent_id={self.grandparent_id!r}, component_name_path={self.component_name_path!r}, "
+            f"content_before_component={self.content_before_component[:10]!r})"
+        )
+
 
 # Function that accepts a list of extra HTML attributes to be set on the component's root elements
 # and returns the component's HTML content and a dictionary of child components' IDs
@@ -102,7 +110,7 @@ def component_post_render(
     render_id: str,
     component_name: str,
     parent_id: Optional[str],
-    on_component_rendered: Callable[[str], None],
+    on_component_rendered_callbacks: Dict[str, Callable[[str], str]],
     on_html_rendered: Callable[[str], str],
 ) -> str:
     # Instead of rendering the component's HTML content immediately, we store it,
@@ -146,36 +154,89 @@ def component_post_render(
     # 3. If there were any extra attributes set by the parent component, we apply these to the renderer.
     # 4. We split the content by placeholders, and put the pairs of (content, placeholder_id) into the queue,
     #    repeating this whole process until we've processed all nested components.
-    content_parts: List[str] = []
+    # 5. If the placeholder ID is None, then we've reached the end of the component's HTML content,
+    #    and we can go one level up to continue the process with component's parent.
     process_queue: Deque[PostRenderQueueItem] = deque()
 
     process_queue.append(
         PostRenderQueueItem(
             content_before_component="",
-            component_id=render_id,
+            child_id=render_id,
             parent_id=None,
+            grandparent_id=None,
             component_name_path=[],
         )
     )
 
+    # By looping over the queue below, we obtain bits of rendered HTML, which we then
+    # must all join together into a single final HTML.
+    #
+    # But instead of joining it all up once at the end, we join the bits on component basis.
+    # So if component has a template like this:
+    # ```django
+    # <div>
+    #   Hello
+    #   {% component "table" / %}
+    # </div>
+    # ```
+    #
+    # Then we end up with 3 bits - 1. test before, 2. component, and 3. text after
+    #
+    # We know when we've arrived at component's end, because `child_id` will be set to `None`.
+    # So we can collect the HTML parts by the component ID, and when we hit the end, we join
+    # all the bits that belong to the same component.
+    #
+    # Once the component's HTML is joined, we can call the callback for the component, and
+    # then add the joined HTML to the cache for the parent component to continue the cycle.
+    html_parts_by_component_id: Dict[str, List[str]] = {}
+    content_parts: List[str] = []
+
+    def get_html_parts(component_id: str) -> List[str]:
+        if component_id not in html_parts_by_component_id:
+            html_parts_by_component_id[component_id] = []
+        return html_parts_by_component_id[component_id]
+
     while len(process_queue):
         curr_item = process_queue.popleft()
 
-        # Process content before the component
-        if curr_item.content_before_component:
-            content_parts.append(curr_item.content_before_component)
-
         # In this case we've reached the end of the component's HTML content, and there's
         # no more subcomponents to process.
-        if curr_item.component_id is None:
-            on_component_rendered(curr_item.parent_id)  # type: ignore[arg-type]
+        if curr_item.child_id is None:
+            # Parent ID must NOT be None in this branch
+            if curr_item.parent_id is None:
+                raise RuntimeError("Parent ID is None")
+
+            parent_parts = html_parts_by_component_id.pop(curr_item.parent_id, [])
+
+            # Add the left-over content
+            parent_parts.append(curr_item.content_before_component)
+
+            # Allow to optionally override/modify the rendered content from outside
+            component_html = "".join(parent_parts)
+            on_component_rendered = on_component_rendered_callbacks[curr_item.parent_id]
+            component_html = on_component_rendered(component_html)  # type: ignore[arg-type]
+
+            # Add the component's HTML to parent's parent's HTML parts
+            if curr_item.grandparent_id is not None:
+                target_list = get_html_parts(curr_item.grandparent_id)
+                target_list.append(component_html)
+            else:
+                content_parts.append(component_html)
+
             continue
 
+        # Process content before the component
+        if curr_item.content_before_component:
+            if curr_item.parent_id is None:
+                raise RuntimeError("Parent ID is None")
+            parent_html_parts = get_html_parts(curr_item.parent_id)
+            parent_html_parts.append(curr_item.content_before_component)
+
         # Generate component's content, applying the extra HTML attributes set by the parent component
-        curr_comp_renderer, curr_comp_name = component_renderer_cache.pop(curr_item.component_id)
+        curr_comp_renderer, curr_comp_name = component_renderer_cache.pop(curr_item.child_id)
         # NOTE: This may be undefined, because this is set only for components that
         # are also root elements in their parent's HTML
-        curr_comp_attrs = child_component_attrs.pop(curr_item.component_id, None)
+        curr_comp_attrs = child_component_attrs.pop(curr_item.child_id, None)
 
         full_path = [*curr_item.component_name_path, curr_comp_name]
 
@@ -184,14 +245,14 @@ def component_post_render(
         # NOTE: [1:] because the root component will be yet again added to the error's
         # `components` list in `_render` so we remove the first element from the path.
         with component_error_message(full_path[1:]):
-            curr_comp_content, curr_child_component_attrs = curr_comp_renderer(curr_comp_attrs)
+            curr_comp_content, grandchild_component_attrs = curr_comp_renderer(curr_comp_attrs)
 
         # Exclude the `data-djc-scope-...` attribute from being applied to the child component's HTML
-        for key in list(curr_child_component_attrs.keys()):
+        for key in list(grandchild_component_attrs.keys()):
             if key.startswith("data-djc-scope-"):
-                curr_child_component_attrs.pop(key, None)
+                grandchild_component_attrs.pop(key, None)
 
-        child_component_attrs.update(curr_child_component_attrs)
+        child_component_attrs.update(grandchild_component_attrs)
 
         # Process the component's content
         last_index = 0
@@ -204,31 +265,32 @@ def component_post_render(
             comp_part = match[0]
 
             # Extract the placeholder ID from `<template djc-render-id="a1b3cf"></template>`
-            curr_child_id_match = render_id_pattern.search(comp_part)
-            if curr_child_id_match is None:
+            grandchild_id_match = render_id_pattern.search(comp_part)
+            if grandchild_id_match is None:
                 raise ValueError(f"No placeholder ID found in {comp_part}")
-            curr_child_id = curr_child_id_match.group("render_id")
+            grandchild_id = grandchild_id_match.group("render_id")
             parts_to_process.append(
                 PostRenderQueueItem(
                     content_before_component=part_before_component,
-                    component_id=curr_child_id,
-                    parent_id=curr_item.component_id,
+                    child_id=grandchild_id,
+                    parent_id=curr_item.child_id,
+                    grandparent_id=curr_item.parent_id,
                     component_name_path=full_path,
                 )
             )
 
         # Append any remaining text
-        if last_index < len(curr_comp_content):
-            parts_to_process.append(
-                PostRenderQueueItem(
-                    content_before_component=curr_comp_content[last_index:],
-                    # Setting component_id to None means that this is the last part of the component's HTML
-                    # and we're done with this component
-                    component_id=None,
-                    parent_id=curr_item.component_id,
-                    component_name_path=full_path,
-                )
+        parts_to_process.append(
+            PostRenderQueueItem(
+                content_before_component=curr_comp_content[last_index:],
+                # Setting `child_id` to None means that this is the last part of the component's HTML
+                # and we're done with this component
+                child_id=None,
+                parent_id=curr_item.child_id,
+                grandparent_id=curr_item.parent_id,
+                component_name_path=full_path,
             )
+        )
 
         process_queue.extendleft(reversed(parts_to_process))
 
diff --git a/tests/test_component.py b/tests/test_component.py
index 13ce9c91..bcc68aee 100644
--- a/tests/test_component.py
+++ b/tests/test_component.py
@@ -1248,6 +1248,16 @@ class ComponentRenderTest(BaseTestCase):
 
 class ComponentHookTest(BaseTestCase):
     def test_on_render_before(self):
+        @register("nested")
+        class NestedComponent(Component):
+            template: types.django_html = """
+                {% load component_tags %}
+                Hello from nested
+                <div>
+                    {% slot "content" default / %}
+                </div>
+            """
+
         class SimpleComponent(Component):
             template: types.django_html = """
                 {% load component_tags %}
@@ -1255,6 +1265,10 @@ class ComponentHookTest(BaseTestCase):
                 kwargs: {{ kwargs|safe }}
                 ---
                 from_on_before: {{ from_on_before }}
+                ---
+                {% component "nested" %}
+                    Hello from simple
+                {% endcomponent %}
             """
 
             def get_context_data(self, *args, **kwargs):
@@ -1268,6 +1282,9 @@ class ComponentHookTest(BaseTestCase):
                 context["from_on_before"] = ":)"
 
                 # Insert text into the Template
+                #
+                # NOTE: Users should NOT do this, because this will insert the text every time
+                #       the component is rendered.
                 template.nodelist.append(TextNode("\n---\nFROM_ON_BEFORE"))
 
         rendered = SimpleComponent.render()
@@ -1279,6 +1296,11 @@ class ComponentHookTest(BaseTestCase):
             ---
             from_on_before: :)
             ---
+            Hello from nested
+            <div data-djc-id-a1bc3e data-djc-id-a1bc40>
+                Hello from simple
+            </div>
+            ---
             FROM_ON_BEFORE
             """,
         )
@@ -1287,13 +1309,27 @@ class ComponentHookTest(BaseTestCase):
     def test_on_render_after(self):
         captured_content = None
 
+        @register("nested")
+        class NestedComponent(Component):
+            template: types.django_html = """
+                {% load component_tags %}
+                Hello from nested
+                <div>
+                    {% slot "content" default / %}
+                </div>
+            """
+
         class SimpleComponent(Component):
             template: types.django_html = """
                 {% load component_tags %}
                 args: {{ args|safe }}
                 kwargs: {{ kwargs|safe }}
                 ---
-                from_on_before: {{ from_on_before }}
+                from_on_after: {{ from_on_after }}
+                ---
+                {% component "nested" %}
+                    Hello from simple
+                {% endcomponent %}
             """
 
             def get_context_data(self, *args, **kwargs):
@@ -1305,10 +1341,10 @@ class ComponentHookTest(BaseTestCase):
             # Check that modifying the context or template does nothing
             def on_render_after(self, context: Context, template: Template, content: str) -> None:
                 # Insert value into the Context
-                context["from_on_before"] = ":)"
+                context["from_on_after"] = ":)"
 
                 # Insert text into the Template
-                template.nodelist.append(TextNode("\n---\nFROM_ON_BEFORE"))
+                template.nodelist.append(TextNode("\n---\nFROM_ON_AFTER"))
 
                 nonlocal captured_content
                 captured_content = content
@@ -1321,7 +1357,12 @@ class ComponentHookTest(BaseTestCase):
             args: ()
             kwargs: {}
             ---
-            from_on_before:
+            from_on_after:
+            ---
+            Hello from nested
+            <div data-djc-id-a1bc3e data-djc-id-a1bc40>
+                Hello from simple
+            </div>
             """,
         )
         self.assertHTMLEqual(
@@ -1330,7 +1371,12 @@ class ComponentHookTest(BaseTestCase):
             args: ()
             kwargs: {}
             ---
-            from_on_before:
+            from_on_after:
+            ---
+            Hello from nested
+            <div data-djc-id-a1bc3e data-djc-id-a1bc40>
+                Hello from simple
+            </div>
             """,
         )
 
@@ -1339,6 +1385,16 @@ class ComponentHookTest(BaseTestCase):
     def test_on_render_after_override_output(self):
         captured_content = None
 
+        @register("nested")
+        class NestedComponent(Component):
+            template: types.django_html = """
+                {% load component_tags %}
+                Hello from nested
+                <div>
+                    {% slot "content" default / %}
+                </div>
+            """
+
         class SimpleComponent(Component):
             template: types.django_html = """
                 {% load component_tags %}
@@ -1346,6 +1402,10 @@ class ComponentHookTest(BaseTestCase):
                 kwargs: {{ kwargs|safe }}
                 ---
                 from_on_before: {{ from_on_before }}
+                ---
+                {% component "nested" %}
+                    Hello from simple
+                {% endcomponent %}
             """
 
             def get_context_data(self, *args, **kwargs):
@@ -1369,6 +1429,11 @@ class ComponentHookTest(BaseTestCase):
             kwargs: {}
             ---
             from_on_before:
+            ---
+            Hello from nested
+            <div data-djc-id-a1bc3e data-djc-id-a1bc40>
+                Hello from simple
+            </div>
             """,
         )
         self.assertHTMLEqual(
@@ -1379,5 +1444,63 @@ class ComponentHookTest(BaseTestCase):
             kwargs: {}
             ---
             from_on_before:
+            ---
+            Hello from nested
+            <div data-djc-id-a1bc3e data-djc-id-a1bc40>
+                Hello from simple
+            </div>
             """,
         )
+
+    def test_on_render_before_after_same_context(self):
+        context_in_before = None
+        context_in_after = None
+
+        @register("nested")
+        class NestedComponent(Component):
+            template: types.django_html = """
+                {% load component_tags %}
+                Hello from nested
+                <div>
+                    {% slot "content" default / %}
+                </div>
+            """
+
+        class SimpleComponent(Component):
+            template: types.django_html = """
+                {% load component_tags %}
+                args: {{ args|safe }}
+                kwargs: {{ kwargs|safe }}
+                ---
+                from_on_after: {{ from_on_after }}
+                ---
+                {% component "nested" %}
+                    Hello from simple
+                {% endcomponent %}
+            """
+
+            def get_context_data(self, *args, **kwargs):
+                return {
+                    "args": args,
+                    "kwargs": kwargs,
+                }
+
+            def on_render_before(self, context: Context, template: Template) -> None:
+                context["from_on_before"] = ":)"
+                nonlocal context_in_before
+                context_in_before = context
+
+            # Check that modifying the context or template does nothing
+            def on_render_after(self, context: Context, template: Template, html: str) -> None:
+                context["from_on_after"] = ":)"
+                nonlocal context_in_after
+                context_in_after = context
+
+        SimpleComponent.render()
+
+        self.assertEqual(
+            context_in_before,
+            context_in_after,
+        )
+        self.assertIn("from_on_before", context_in_before)
+        self.assertIn("from_on_after", context_in_after)