Implement collapsing using HTML AST

1 files changed, 30 insertions, 3 deletions

diff --git a/mammoth/html_paths.py b/mammoth/html_paths.py
index 1428015..e043e23 100644
--- a/mammoth/html_paths.py
+++ b/mammoth/html_paths.py
@@ -1,4 +1,6 @@
-import collections
+import cobble
+
+from . import html
 
 
 def path(elements):
@@ -13,7 +15,32 @@ def element(names, class_names=None, fresh=None):
     return HtmlPathElement(names, class_names, fresh)
 
 
-HtmlPath = collections.namedtuple("HtmlPath", ["elements"])
-HtmlPathElement = collections.namedtuple("HtmlPathElement", ["names", "class_names", "fresh"])
+@cobble.data
+class HtmlPath(object):
+    elements = cobble.field()
+    
+    def wrap(self, nodes):
+        for element in reversed(self.elements):
+            nodes = element.wrap(nodes)
+        
+        return nodes
+
+@cobble.data
+class HtmlPathElement(object):
+    names = cobble.field()
+    class_names = cobble.field()
+    fresh = cobble.field()
+
+    def wrap(self, nodes):
+        if self.class_names:
+            attributes = {"class": " ".join(self.class_names)}
+        else:
+            attributes = {}
+        element = html.element(
+            self.names,
+            attributes,
+            nodes,
+            collapsible=not self.fresh)
+        return [element]
 
 empty = path([])