Add collapsing of HTML nodes

1 files changed, 66 insertions, 2 deletions

diff --git a/mammoth/html/__init__.py b/mammoth/html/__init__.py
index 905b423..54bb87d 100644
--- a/mammoth/html/__init__.py
+++ b/mammoth/html/__init__.py
@@ -7,13 +7,22 @@ def text(value):
 
 
 def element(tag_name, attributes=None, children=None):
+    return _element(tag_name, attributes, children, collapsible=False)
+
+
+def collapsible_element(tag_name, attributes=None, children=None):
+    return _element(tag_name, attributes, children, collapsible=True)
+
+
+def _element(tag_names, attributes, children, collapsible):
+    if not isinstance(tag_names, list):
+        tag_names = [tag_names]
     if attributes is None:
         attributes = {}
     if children is None:
         children = []
+    return Element(tag_names, attributes, children, collapsible=collapsible)
     
-    return Element([tag_name], attributes, children, collapsible=False)
-
 
 def self_closing_element(tag_name, attributes=None):
     if attributes is None:
@@ -56,3 +65,58 @@ class StripEmpty(NodeVisitor):
     
     def visit_force_write(self, node):
         return [node]
+
+
+def collapse(nodes):
+    collapsed = []
+    
+    for node in nodes:
+        _collapsing_add(collapsed, node)
+    
+    return collapsed
+
+class _CollapseNode(NodeVisitor):
+    def visit_text_node(self, node):
+        return node
+    
+    def visit_element(self, element):
+        return Element(
+            element.tag_names,
+            element.attributes,
+            collapse(element.children),
+            collapsible=element.collapsible)
+    
+    def visit_self_closing_element(self, element):
+        return element
+    
+    def visit_force_write(self, node):
+        return node
+    
+_collapse_node = _CollapseNode().visit
+
+
+def _collapsing_add(collapsed, node):
+    collapsed_node = _collapse_node(node)
+    if not _try_collapse(collapsed, collapsed_node):
+        collapsed.append(collapsed_node)
+    
+def _try_collapse(collapsed, node):
+    if not collapsed:
+        return False
+
+    last = collapsed[-1]
+    if not isinstance(last, Element) or not isinstance(node, Element):
+        return False
+    
+    if not node.collapsible:
+        return False
+        
+    if not _is_match(last, node):
+        return False
+    
+    for child in node.children:
+        _collapsing_add(last.children, child)
+    return True
+
+def _is_match(first, second):
+    return first.tag_name in second.tag_names and first.attributes == second.attributes