diff --git a/main.go b/main.go index 67471d2..e9dfa55 100644 --- a/main.go +++ b/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 { - displayer.Display(currNodes) - } else if printNumber { + if printNumber { fmt.Println(len(currNodes)) } else { - for _, s := range currNodes { - // defined in `printing.go` - PrintNode(s, 0) - } + displayer.Display(currNodes) } } diff --git a/printing.go b/printing.go index 8bb2ec0..51880a9 100644 --- a/printing.go +++ b/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 or
? // 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) - printIndent(level) + t.printChildren(n, level+1) + t.printIndent(level) if printColor { tokenColor.Print("