Add collapsing of HTML nodes

3 files changed, 177 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
diff --git a/mammoth/html/nodes.py b/mammoth/html/nodes.py
index aae4888..a5a5bde 100644
--- a/mammoth/html/nodes.py
+++ b/mammoth/html/nodes.py
@@ -16,6 +16,10 @@ class Element(Node):
     attributes = cobble.field()
     children = cobble.field()
     collapsible = cobble.field()
+    
+    @property
+    def tag_name(self):
+        return self.tag_names[0]
 
 
 @cobble.data
diff --git a/tests/html/collapse_tests.py b/tests/html/collapse_tests.py
new file mode 100644
index 0000000..52cecd2
--- /dev/null
+++ b/tests/html/collapse_tests.py
@@ -0,0 +1,107 @@
+from nose.tools import istest, assert_equal
+
+from mammoth import html
+
+
+@istest
+def collapsing_does_nothing_to_single_text_node():
+    assert_equal(
+        html.collapse([html.text("Bluebells")]),
+        [html.text("Bluebells")])
+
+
+@istest
+def consecutive_fresh_elements_are_not_collapsed():
+    assert_equal(
+        html.collapse([html.element("p"), html.element("p")]),
+        [html.element("p"), html.element("p")])
+
+
+@istest
+def consecutive_collapsible_elements_are_collapsed_if_they_have_the_same_tag_and_attributes():
+    assert_equal(
+        [html.collapsible_element("p", {}, [html.text("One"), html.text("Two")])],
+        html.collapse([
+            html.collapsible_element("p", {}, [html.text("One")]),
+            html.collapsible_element("p", {}, [html.text("Two")])
+        ]))
+        
+
+@istest
+def elements_with_different_tag_names_are_not_collapsed():
+    assert_equal(
+        [
+            html.collapsible_element("p", {}, [html.text("One")]),
+            html.collapsible_element("div", {}, [html.text("Two")])
+        ],
+        
+        html.collapse([
+            html.collapsible_element("p", {}, [html.text("One")]),
+            html.collapsible_element("div", {}, [html.text("Two")])
+        ]))
+        
+
+@istest
+def elements_with_different_attributes_are_not_collapsed():
+    assert_equal(
+        [
+            html.collapsible_element("p", {"id": "a"}, [html.text("One")]),
+            html.collapsible_element("p", {}, [html.text("Two")])
+        ],
+        
+        html.collapse([
+            html.collapsible_element("p", {"id": "a"}, [html.text("One")]),
+            html.collapsible_element("p", {}, [html.text("Two")])
+        ]))
+
+
+@istest
+def children_of_collapsed_element_can_collapse_with_children_of_previous_element():
+    assert_equal(
+        [
+            html.collapsible_element("blockquote", {}, [
+                html.collapsible_element("p", {}, [
+                    html.text("One"),
+                    html.text("Two")
+                ])
+            ]),
+        ],
+        
+        html.collapse([
+            html.collapsible_element("blockquote", {}, [
+                html.collapsible_element("p", {}, [html.text("One")])
+            ]),
+            html.collapsible_element("blockquote", {}, [
+                html.collapsible_element("p", {}, [html.text("Two")])
+            ]),
+        ]))
+
+
+@istest
+def collapsible_element_can_collapse_into_previous_fresh_element():
+    assert_equal(
+        [html.element("p", {}, [html.text("One"), html.text("Two")])],
+        html.collapse([
+            html.element("p", {}, [html.text("One")]),
+            html.collapsible_element("p", {}, [html.text("Two")])
+        ]))
+
+
+@istest
+def element_with_choice_of_tag_names_can_collapse_into_previous_element_if_it_has_one_of_those_tag_names_as_its_main_tag_name():
+    assert_equal(
+        [html.collapsible_element(["ol"])],
+        html.collapse([
+            html.collapsible_element("ol"),
+            html.collapsible_element(["ul", "ol"])
+        ]))
+
+    assert_equal(
+        [
+            html.collapsible_element(["ul", "ol"]),
+            html.collapsible_element("ol")
+        ],
+        html.collapse([
+            html.collapsible_element(["ul", "ol"]),
+            html.collapsible_element("ol")
+        ]))