@@ -62,6 +62,10 @@ def is_html_well_formed(html: str) -> bool:
6262
6363
6464def assert_has_single_wrapper_element (root_element : Tag , component_name : str ) -> None :
65+ """Assert that there is at least one child in the root element. And that there is only
66+ one root element.
67+ """
68+
6569 # Check that the root element has at least one child
6670 try :
6771 next (root_element .descendants )
@@ -70,8 +74,13 @@ def assert_has_single_wrapper_element(root_element: Tag, component_name: str) ->
7074 f"The '{ component_name } ' component does not appear to have one root element."
7175 ) from None
7276
77+ if "unicorn:view" in root_element .attrs or "u:view" in root_element .attrs :
78+ # If the root element is a direct view, skip the check
79+ return
80+
7381 # Check that there is not more than one root element
7482 parent_element = root_element .parent
83+
7584 tag_count = len ([c for c in parent_element .children if isinstance (c , Tag )])
7685
7786 if tag_count > 1 :
@@ -80,6 +89,37 @@ def assert_has_single_wrapper_element(root_element: Tag, component_name: str) ->
8089 ) from None
8190
8291
92+ def _get_direct_view (tag : Tag ):
93+ return tag .find_next (attrs = {"unicorn:view" : True }) or tag .find_next (attrs = {"u:view" : True })
94+
95+
96+ def get_root_element (soup : BeautifulSoup ) -> Tag :
97+ """Gets the first tag element for the component or the first element with a `unicorn:view` attribute for a direct
98+ view.
99+
100+ Returns:
101+ BeautifulSoup tag element.
102+
103+ Raises `Exception` if an element cannot be found.
104+ """
105+
106+ for element in soup .contents :
107+ if isinstance (element , Tag ) and element .name :
108+ if element .name == "html" :
109+ view_element = _get_direct_view (element )
110+
111+ if not view_element :
112+ raise MissingComponentViewElementError (
113+ "An element with an `unicorn:view` attribute is required for a direct view"
114+ )
115+
116+ return view_element
117+
118+ return element
119+
120+ raise MissingComponentElementError ("No root element for the component was found" )
121+
122+
83123class UnsortedAttributes (HTMLFormatter ):
84124 """
85125 Prevent beautifulsoup from re-ordering attributes.
@@ -230,32 +270,3 @@ def render(self):
230270 def _desoupify (soup ):
231271 soup .smooth ()
232272 return soup .encode (formatter = UnsortedAttributes ()).decode ("utf-8" )
233-
234-
235- def get_root_element (soup : BeautifulSoup ) -> Tag :
236- """
237- Gets the first tag element for the component or the first element with a `unicorn:view` attribute for a direct view.
238-
239- Returns:
240- BeautifulSoup tag element.
241-
242- Raises `Exception` if an element cannot be found.
243- """
244-
245- for element in soup .contents :
246- if isinstance (element , Tag ) and element .name :
247- if element .name == "html" :
248- view_element = element .find_next (attrs = {"unicorn:view" : True }) or element .find_next (
249- attrs = {"u:view" : True }
250- )
251-
252- if not view_element :
253- raise MissingComponentViewElementError (
254- "An element with an `unicorn:view` attribute is required for a direct view"
255- )
256-
257- return view_element
258-
259- return element
260-
261- raise MissingComponentElementError ("No root element for the component was found" )
0 commit comments