node.go

package node

import (
	"io"
	"strings"

	"golang.org/x/net/html"
)

var (
	_ HtmlNode = &htmlNode{}
	_ Node     = &node{}
	_ TextNode = &textNode{}
)

// Node is an interface representing an HTML node.
type Node interface {
	HtmlNode

	// String returns a TextNode if the node has only one child whose type is text, otherwise returns nil.
	String() TextNode

	// Strings return all of the text nodes inside this node.
	Strings() []TextNode

	// StrippedStrings return a list of strings generated by Strings, where strings consisting entirely of
	// whitespace are ignored, and whitespace at the beginning and end of strings is removed.
	StrippedStrings() []string

	// GetText concatenates all of the text node's content.
	GetText() string
}

// TextNode is an interface representing a text node.
type TextNode interface {
	HtmlNode

	// String returns content for text node.
	String() string
}

// HtmlNode is an interface representing an HTML node.
type HtmlNode interface {
	// Raw returns origin *html.Node.
	Raw() *html.Node
	// ToNode converts HtmlNode to Node.
	ToNode() Node
	// ToTextNode converts HtmlNode to TextNode.
	// It will panic if the node type is not text node.
	ToTextNode() TextNode

	// Type returns a NodeType.
	Type() html.NodeType
	// Data returns tag name for element node or content for text node.
	Data() string
	// Attrs returns an Attributes interface for element node.
	Attrs() Attributes
	// HasAttr return whether node has an attribute.
	HasAttr(string) bool
	// HTML renders the node's parse tree as HTML code.
	HTML() string
	// Readable renders unescaped HTML code.
	Readable() string

	// Parent returns the parent of this node.
	Parent() Node
	// FirstChild returns the first child of this node.
	FirstChild() Node
	// LastChild returns the last child of this node.
	LastChild() Node
	// PrevSibling returns the previous node that are on the same level of the parse tree.
	PrevSibling() Node
	// NextSibling returns the next node that are on the same level of the parse tree.
	NextSibling() Node
	// PrevNode returns the node that was parsed immediately before this node.
	PrevNode() Node
	// NextNode returns the node that was parsed immediately after this node.
	NextNode() Node

	// Parents iterate over all of this node's parent recursively.
	Parents() []Node
	// Children return all of this node's direct children.
	Children() []Node
	// Descendants iterate over all of this node's children recursively.
	Descendants() []Node
	// PrevSiblings return all of this node's previous nodes that are on the same level of the parse tree.
	PrevSiblings() []Node
	// NextSiblings return all of this node's next nodes that are on the same level of the parse tree.
	NextSiblings() []Node
	// PrevNodes return all of the nodes that was parsed before this node.
	PrevNodes() []Node
	// NextNodes return all of the nodes that was parsed after this node.
	NextNodes() []Node

	// Finder includes a set of find methods.
	Finder
}

type htmlNode struct {
	*html.Node
}

// Parse returns the parse tree for the HTML from the given Reader.
func Parse(r io.Reader) (Node, error) {
	return ParseWithOptions(r)
}

// ParseWithOptions is like Parse, with options.
func ParseWithOptions(r io.Reader, opts ...html.ParseOption) (Node, error) {
	n, err := html.ParseWithOptions(r, opts...)
	if err != nil {
		return nil, err
	}
	return NewNode(n), nil
}

// ParseHTML returns the parse tree for the HTML from string.
func ParseHTML(s string) (Node, error) {
	return Parse(strings.NewReader(s))
}

func (n *htmlNode) Raw() *html.Node {
	return n.Node
}

func (n *htmlNode) ToNode() Node {
	return &node{n}
}

func (n *htmlNode) ToTextNode() TextNode {
	if n.Type() != html.TextNode {
		panic("node is not TextNode")
	}
	return &textNode{n}
}

func (n *htmlNode) Type() html.NodeType {
	return n.Raw().Type
}

func (n *htmlNode) Data() string {
	return n.Raw().Data
}

func (n *htmlNode) Attrs() Attributes {
	attrs := make(attributes)
	for _, i := range n.Node.Attr {
		if _, ok := attrs[i.Key]; !ok {
			attrs[i.Key] = i.Val
		}
	}
	return attrs
}

func (n *htmlNode) HasAttr(attr string) bool {
	_, ok := n.Attrs().Get(attr)
	return ok
}

func (n *htmlNode) HTML() string {
	var b strings.Builder
	html.Render(&b, n.Raw())
	return b.String()
}

func (n *htmlNode) Readable() string {
	return html.UnescapeString(n.HTML())
}

func (n *htmlNode) Parent() Node {
	return NewNode(n.Node.Parent)
}

func (n *htmlNode) FirstChild() Node {
	return NewNode(n.Node.FirstChild)
}

func (n *htmlNode) LastChild() Node {
	return NewNode(n.Node.LastChild)
}

func (n *htmlNode) PrevSibling() Node {
	return NewNode(n.Node.PrevSibling)
}

func (n *htmlNode) NextSibling() Node {
	return NewNode(n.Node.NextSibling)
}

func (n *htmlNode) PrevNode() Node {
	if prev := n.PrevSibling(); prev != nil {
		for node := prev.LastChild(); node != nil; node = prev.LastChild() {
			prev = node
		}
		return prev
	}
	return n.Parent()
}

func (n *htmlNode) NextNode() Node {
	if child := n.FirstChild(); child != nil {
		return child
	}
	if next := n.NextSibling(); next != nil {
		return next
	}
	for node := n.Parent(); node != nil; node = node.Parent() {
		if node := node.NextSibling(); node != nil {
			return node
		}
	}
	return nil
}

func (n *htmlNode) Parents() (parents []Node) {
	parent := n.Parent()
	for parent != nil {
		parents = append(parents, parent)
		parent = parent.Parent()
	}
	return
}

func (n *htmlNode) Children() (children []Node) {
	child := n.FirstChild()
	for child != nil {
		children = append(children, child)
		child = child.NextSibling()
	}
	return
}

func (n *htmlNode) Descendants() (nodes []Node) {
	var f func(Node)
	f = func(node Node) {
		if node.Raw() != n.Raw() {
			nodes = append(nodes, node)
		}
		for node := node.FirstChild(); node != nil; node = node.NextSibling() {
			f(node)
		}
	}
	f(n.ToNode())
	return
}

func (n *htmlNode) PrevSiblings() (prevSiblings []Node) {
	prevSibling := n.PrevSibling()
	for prevSibling != nil {
		prevSiblings = append(prevSiblings, prevSibling)
		prevSibling = prevSibling.PrevSibling()
	}
	return
}

func (n *htmlNode) NextSiblings() (nextSiblings []Node) {
	nextSibling := n.NextSibling()
	for nextSibling != nil {
		nextSiblings = append(nextSiblings, nextSibling)
		nextSibling = nextSibling.NextSibling()
	}
	return
}

func (n *htmlNode) PrevNodes() (prevNodes []Node) {
	prevNode := n.PrevNode()
	for prevNode != nil {
		prevNodes = append(prevNodes, prevNode)
		prevNode = prevNode.PrevNode()
	}
	return
}

func (n *htmlNode) NextNodes() (nextNodes []Node) {
	nextNode := n.NextNode()
	for nextNode != nil {
		nextNodes = append(nextNodes, nextNode)
		nextNode = nextNode.NextNode()
	}
	return
}

type node struct {
	*htmlNode
}

// NewNode returns a Node with the specified *html.Node.
func NewNode(n *html.Node) Node {
	if n == nil {
		return nil
	}
	return &node{&htmlNode{n}}
}

func (n *node) String() TextNode {
	if n.Type() == html.TextNode {
		return &textNode{n.htmlNode}
	}
	if children := n.Children(); len(children) == 1 {
		if node := children[0].String(); node != nil {
			return node
		}
	}
	return nil
}

func (n *node) Strings() (strings []TextNode) {
	if n.Type() == html.TextNode {
		strings = append(strings, &textNode{n.htmlNode})
	}
	for _, node := range n.Descendants() {
		if node.Type() == html.TextNode {
			strings = append(strings, node.ToTextNode())
		}
	}
	return
}

func (n *node) StrippedStrings() (stripped []string) {
	for _, i := range n.Strings() {
		if s := strings.TrimSpace(i.String()); s != "" {
			stripped = append(stripped, s)
		}
	}
	return
}

func (n *node) GetText() string {
	var b strings.Builder
	for _, i := range n.Strings() {
		b.WriteString(i.String())
	}
	return b.String()
}

type textNode struct {
	*htmlNode
}

func (n *textNode) String() string {
	return n.Data()
}