default print is now a displayer

main.go

@@ -128,7 +128,8 @@ func main() {
 	}
 	inputStream.Close()
 	if len(cmds) == 0 {
-		PrintNode(root, 0)
+		t := TreeDisplayer{indentString}
+		t.Display([]*html.Node{root})
 		os.Exit(0)
 	}
 	selectors := make([]selector.Selector, len(cmds))
@@ -141,6 +142,8 @@ func main() {
 				displayer = d
 				selectors = selectors[0 : len(cmds)-1]
 				break
+			} else {
+				displayer = TreeDisplayer{indentString}
 			}
 		}
 		selectors[i], err = selector.NewSelector(cmd)
@@ -152,14 +155,9 @@ func main() {
 	for _, selector := range selectors {
 		currNodes = selector.Select(currNodes)
 	}
-	if displayer != nil {
+	if printNumber {
-		displayer.Display(currNodes)
-	} else if printNumber {
 		fmt.Println(len(currNodes))
 	} else {
-		for _, s := range currNodes {
+		displayer.Display(currNodes)
-			// defined in `printing.go`
-			PrintNode(s, 0)
-		}
 	}
 }

printing.go

@@ -26,12 +26,6 @@ func init() {
 	color.Output = colorable.NewColorableStdout()
 }
 
-func printIndent(level int) {
-	for ; level > 0; level-- {
-		fmt.Print(indentString)
-	}
-}
-
 // Is this node a tag with no end tag such as <meta> or <br>?
 // http://www.w3.org/TR/html-markup/syntax.html#syntax-elements
 func isVoidElement(n *html.Node) bool {
@@ -45,34 +39,50 @@ func isVoidElement(n *html.Node) bool {
 
 }
 
-func printChildren(n *html.Node, level int) {
+type TreeDisplayer struct {
+	IndentString string
+}
+
+func (t TreeDisplayer) Display(nodes []*html.Node) {
+	for _, node := range nodes {
+		t.printNode(node, 0)
+	}
+}
+
+func (t TreeDisplayer) printChildren(n *html.Node, level int) {
 	if maxPrintLevel > -1 {
 		if level >= maxPrintLevel {
-			printIndent(level)
+			t.printIndent(level)
 			fmt.Println("...")
 			return
 		}
 	}
 	child := n.FirstChild
 	for child != nil {
-		PrintNode(child, level)
+		t.printNode(child, level)
 		child = child.NextSibling
 	}
 }
 
+func (t TreeDisplayer) printIndent(level int) {
+	for ; level > 0; level-- {
+		fmt.Print(indentString)
+	}
+}
+
 // Print a node and all of it's children to `maxlevel`.
-func PrintNode(n *html.Node, level int) {
+func (t TreeDisplayer) printNode(n *html.Node, level int) {
 	switch n.Type {
 	case html.TextNode:
 		s := n.Data
 		if !whitespaceRegexp.MatchString(s) {
 			s = preWhitespace.ReplaceAllString(s, "")
 			s = postWhitespace.ReplaceAllString(s, "")
-			printIndent(level)
+			t.printIndent(level)
 			fmt.Println(s)
 		}
 	case html.ElementNode:
-		printIndent(level)
+		t.printIndent(level)
 		if printColor {
 			tokenColor.Print("<")
 			tagColor.Printf("%s", n.Data)
@@ -95,8 +105,8 @@ func PrintNode(n *html.Node, level int) {
 			fmt.Print(">\n")
 		}
 		if !isVoidElement(n) {
-			printChildren(n, level+1)
+			t.printChildren(n, level+1)
-			printIndent(level)
+			t.printIndent(level)
 			if printColor {
 				tokenColor.Print("</")
 				tagColor.Printf("%s", n.Data)
@@ -106,6 +116,6 @@ func PrintNode(n *html.Node, level int) {
 			}
 		}
 	case html.CommentNode, html.DoctypeNode, html.DocumentNode:
-		printChildren(n, level)
+		t.printChildren(n, level)
 	}
 }