Handle component templates with comments or blank strings before a root element.

Author: adamghill

django_unicorn/components/unicorn_template_response.py

@@ -68,7 +68,7 @@ def render(self):
         checksum = generate_checksum(orjson.dumps(frontend_context_variables_dict))
 
         soup = BeautifulSoup(content, features="html.parser")
-        root_element = UnicornTemplateResponse._get_root_element(soup)
+        root_element = get_root_element(soup)
         root_element["unicorn:id"] = self.component.component_id
         root_element["unicorn:name"] = self.component.component_name
         root_element["unicorn:key"] = self.component.component_key
@@ -122,23 +122,23 @@ def render(self):
 
         return response
 
-    @staticmethod
-    def _get_root_element(soup: BeautifulSoup) -> Tag:
-        """
-        Gets the first element.
-
-        Returns:
-            BeautifulSoup element.
-            
-            Raises an Exception if an element cannot be found.
-        """
-        for element in soup.contents:
-            if element.name:
-                return element
-
-        raise Exception("No root element found")
-
     @staticmethod
     def _desoupify(soup):
         soup.smooth()
         return soup.encode(formatter=UnsortedAttributes()).decode("utf-8")
+
+
+def get_root_element(soup: BeautifulSoup) -> Tag:
+    """
+    Gets the first tag element.
+
+    Returns:
+        BeautifulSoup tag element.
+
+        Raises an Exception if an element cannot be found.
+    """
+    for element in soup.contents:
+        if isinstance(element, Tag) and element.name:
+            return element
+
+    raise Exception("No root element found")

django_unicorn/views/__init__.py

@@ -13,6 +13,7 @@
 from bs4 import BeautifulSoup
 
 from django_unicorn.components import UnicornView
+from django_unicorn.components.unicorn_template_response import get_root_element
 from django_unicorn.decorators import timed
 from django_unicorn.errors import RenderNotModified, UnicornCacheError, UnicornViewError
 from django_unicorn.serializer import dumps, loads
@@ -217,6 +218,10 @@ def _process_component_request(
         ):
             raise RenderNotModified()
 
+        # Make sure that partials with comments or blank lines before the root element only return the root element
+        soup = BeautifulSoup(rendered_component, features="html.parser")
+        rendered_component = str(get_root_element(soup))
+
         res.update(
             {"dom": rendered_component, "hash": hash,}
         )

example/unicorn/templates/unicorn/objects.html

@@ -1,3 +1,5 @@
+<!-- test comment -->
+
 <div>
   <div>
     <label>Dictionaries (original)</label>

tests/components/test_unicorn_template_response.py

@@ -0,0 +1,46 @@
+import pytest
+from bs4 import BeautifulSoup
+
+from django_unicorn.components.unicorn_template_response import get_root_element
+
+
+def test_get_root_element():
+    expected = "<div>test</div>"
+
+    component_html = "<div>test</div>"
+    soup = BeautifulSoup(component_html, features="html.parser")
+    actual = get_root_element(soup)
+
+    assert str(actual) == expected
+
+
+def test_get_root_element_with_comment():
+    expected = "<div>test</div>"
+
+    component_html = "<!-- some comment --><div>test</div>"
+    soup = BeautifulSoup(component_html, features="html.parser")
+    actual = get_root_element(soup)
+
+    assert str(actual) == expected
+
+
+def test_get_root_element_with_blank_string():
+    expected = "<div>test</div>"
+
+    component_html = "\n<div>test</div>"
+    soup = BeautifulSoup(component_html, features="html.parser")
+    actual = get_root_element(soup)
+
+    assert str(actual) == expected
+
+
+def test_get_root_element_no_element():
+    expected = "<div>test</div>"
+
+    component_html = "\n"
+    soup = BeautifulSoup(component_html, features="html.parser")
+
+    with pytest.raises(Exception):
+        actual = get_root_element(soup)
+
+        assert str(actual) == expected