Do all HTML writing in HtmlWriter

6 files changed, 64 insertions, 19 deletions

diff --git a/src/main/java/org/zwobble/mammoth/html/Html.java b/src/main/java/org/zwobble/mammoth/html/Html.java
index bd139c8..c22911d 100644
--- a/src/main/java/org/zwobble/mammoth/html/Html.java
+++ b/src/main/java/org/zwobble/mammoth/html/Html.java
@@ -9,7 +9,7 @@ import static org.zwobble.mammoth.util.MammothMaps.map;
 public class Html {
     public static String write(List<HtmlNode> nodes) {
         StringBuilder builder = new StringBuilder();
-        nodes.forEach(node -> node.write(builder));
+        nodes.forEach(node -> HtmlWriter.write(node, builder));
         return builder.toString();
     }
 
diff --git a/src/main/java/org/zwobble/mammoth/html/HtmlElement.java b/src/main/java/org/zwobble/mammoth/html/HtmlElement.java
index 8111507..637b009 100644
--- a/src/main/java/org/zwobble/mammoth/html/HtmlElement.java
+++ b/src/main/java/org/zwobble/mammoth/html/HtmlElement.java
@@ -14,19 +14,20 @@ public class HtmlElement implements HtmlNode {
         this.children = children;
     }
 
-    @Override
-    public void write(StringBuilder builder) {
-        builder.append("<").append(tagName);
-
-        HtmlWriter.generateAttributes(attributes, builder);
+    public String getTagName() {
+        return tagName;
+    }
 
-        builder.append(">");
+    public Map<String, String> getAttributes() {
+        return attributes;
+    }
 
-        children.forEach(node -> node.write(builder));
+    public List<HtmlNode> getChildren() {
+        return children;
+    }
 
-        builder
-            .append("</")
-            .append(tagName)
-            .append(">");
+    @Override
+    public void visit(Visitor visitor) {
+        visitor.visit(this);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/html/HtmlNode.java b/src/main/java/org/zwobble/mammoth/html/HtmlNode.java
index 9977c66..a5fc622 100644
--- a/src/main/java/org/zwobble/mammoth/html/HtmlNode.java
+++ b/src/main/java/org/zwobble/mammoth/html/HtmlNode.java
@@ -1,5 +1,11 @@
 package org.zwobble.mammoth.html;
 
 public interface HtmlNode {
-    void write(StringBuilder builder);
+    void visit(Visitor visitor);
+
+    interface Visitor {
+        void visit(HtmlElement element);
+        void visit(HtmlSelfClosingElement element);
+        void visit(HtmlTextNode node);
+    }
 }
diff --git a/src/main/java/org/zwobble/mammoth/html/HtmlSelfClosingElement.java b/src/main/java/org/zwobble/mammoth/html/HtmlSelfClosingElement.java
index f8bea1d..a1ddf2b 100644
--- a/src/main/java/org/zwobble/mammoth/html/HtmlSelfClosingElement.java
+++ b/src/main/java/org/zwobble/mammoth/html/HtmlSelfClosingElement.java
@@ -11,14 +11,16 @@ public class HtmlSelfClosingElement implements HtmlNode {
         this.attributes = attributes;
     }
 
+    public String getTagName() {
+        return tagName;
+    }
+
     public Map<String, String> getAttributes() {
         return attributes;
     }
 
     @Override
-    public void write(StringBuilder builder) {
-        builder.append("<").append(tagName);
-        HtmlWriter.generateAttributes(attributes, builder);
-        builder.append(" />");
+    public void visit(Visitor visitor) {
+        visitor.visit(this);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/html/HtmlTextNode.java b/src/main/java/org/zwobble/mammoth/html/HtmlTextNode.java
index bc8d1f8..37dceec 100644
--- a/src/main/java/org/zwobble/mammoth/html/HtmlTextNode.java
+++ b/src/main/java/org/zwobble/mammoth/html/HtmlTextNode.java
@@ -7,8 +7,12 @@ public class HtmlTextNode implements HtmlNode {
         this.value = value;
     }
 
+    public String getValue() {
+        return value;
+    }
+
     @Override
-    public void write(StringBuilder builder) {
-        builder.append(HtmlWriter.escapeText(value));
+    public void visit(Visitor visitor) {
+        visitor.visit(this);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/html/HtmlWriter.java b/src/main/java/org/zwobble/mammoth/html/HtmlWriter.java
index 50bf85b..bb98925 100644
--- a/src/main/java/org/zwobble/mammoth/html/HtmlWriter.java
+++ b/src/main/java/org/zwobble/mammoth/html/HtmlWriter.java
@@ -5,6 +5,38 @@ import java.util.Map;
 import static org.zwobble.mammoth.util.MammothLists.orderedBy;
 
 public class HtmlWriter {
+    public static void write(HtmlNode node, StringBuilder builder) {
+        node.visit(new HtmlNode.Visitor() {
+            @Override
+            public void visit(HtmlElement element) {
+                builder.append("<").append(element.getTagName());
+
+                HtmlWriter.generateAttributes(element.getAttributes(), builder);
+
+                builder.append(">");
+
+                element.getChildren().forEach(node -> write(node, builder));
+
+                builder
+                    .append("</")
+                    .append(element.getTagName())
+                    .append(">");
+            }
+
+            @Override
+            public void visit(HtmlSelfClosingElement element) {
+                builder.append("<").append(element.getTagName());
+                HtmlWriter.generateAttributes(element.getAttributes(), builder);
+                builder.append(" />");
+            }
+
+            @Override
+            public void visit(HtmlTextNode node) {
+                builder.append(HtmlWriter.escapeText(node.getValue()));
+            }
+        });
+    }
+
     public static String escapeText(String text) {
         return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;");
     }