Support xml directive

Author: subchen

document.go

@@ -1,17 +1,33 @@
 package xmldom
 
+import "bytes"
+
 type Document struct {
-	ProcInst string
-	Root     *Node
+	ProcInst   string
+	Directives []string
+	Root       *Node
 }
 
 func (d *Document) XML() string {
-	return d.ProcInst + d.Root.XML()
+	buf := new(bytes.Buffer)
+	buf.WriteString(d.ProcInst)
+	for _, directive := range d.Directives {
+		buf.WriteString(directive)
+	}
+	buf.WriteString(d.Root.XML())
+	return buf.String()
 }
 
 func (d *Document) XMLPretty() string {
+	buf := new(bytes.Buffer)
 	if len(d.ProcInst) > 0 {
-		return d.ProcInst + "\n" + d.Root.XMLPretty()
+		buf.WriteString(d.ProcInst)
+		buf.WriteByte('\n')
+	}
+	for _, directive := range d.Directives {
+		buf.WriteString(directive)
+		buf.WriteByte('\n')
 	}
-	return d.Root.XMLPretty()
+	buf.WriteString(d.Root.XMLPretty())
+	return buf.String()
 }

dom.go

@@ -70,6 +70,8 @@ func Parse(r io.Reader) (*Document, error) {
 			}
 		case xml.ProcInst:
 			doc.ProcInst = stringifyProcInst(&token)
+		case xml.Directive:
+			doc.Directives = append(doc.Directives, stringifyDirective(&token))
 		}
 
 		// get the next token

example_test.go

@@ -8,6 +8,7 @@ import (
 
 const (
 	ExampleXml = `<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE junit SYSTEM "junit-result.dtd">
 <testsuites>
 	<testsuite tests="2" failures="0" time="0.009" name="github.com/subchen/go-xmldom">
 		<properties>
@@ -85,14 +86,15 @@ func ExampleDocument_XML() {
 	doc := xmldom.Must(xmldom.ParseXML(ExampleXml))
 	fmt.Println(doc.XML())
 	// Output:
-	// <?xml version="1.0" encoding="UTF-8"?><testsuites><testsuite tests="2" failures="0" time="0.009" name="github.com/subchen/go-xmldom"><properties><property name="go.version">go1.8.1</property></properties><testcase classname="go-xmldom" name="ExampleParseXML" time="0.004" /><testcase classname="go-xmldom" name="ExampleParse" time="0.005" /></testsuite></testsuites>
+	// <?xml version="1.0" encoding="UTF-8"?><!DOCTYPE junit SYSTEM "junit-result.dtd"><testsuites><testsuite tests="2" failures="0" time="0.009" name="github.com/subchen/go-xmldom"><properties><property name="go.version">go1.8.1</property></properties><testcase classname="go-xmldom" name="ExampleParseXML" time="0.004" /><testcase classname="go-xmldom" name="ExampleParse" time="0.005" /></testsuite></testsuites>
 }
 
 func ExampleDocument_XMLPretty() {
 	doc := xmldom.Must(xmldom.ParseXML(ExampleXml))
 	fmt.Println(doc.XMLPretty())
 	// Output:
 	// <?xml version="1.0" encoding="UTF-8"?>
+	// <!DOCTYPE junit SYSTEM "junit-result.dtd">
 	// <testsuites>
 	//   <testsuite tests="2" failures="0" time="0.009" name="github.com/subchen/go-xmldom">
 	//     <properties>

print.go

@@ -14,6 +14,13 @@ func stringifyProcInst(pi *xml.ProcInst) string {
 	return fmt.Sprintf("<?%s %s?>", pi.Target, string(pi.Inst))
 }
 
+func stringifyDirective(directive *xml.Directive) string {
+	if directive == nil {
+		return ""
+	}
+	return fmt.Sprintf("<!%s>", string(*directive))
+}
+
 func printXML(buf *bytes.Buffer, n *Node, level int, indent string) {
 	pretty := len(indent) > 0