pkgs/utils/html.go

package utils

import (
	"bytes"
	"golang.org/x/net/html"
	"io"
	"strings"
)

type Node struct {
	*html.Node
}

func Parse(r io.Reader) (*Node, error) {
	n, err := html.Parse(r)
	if err != nil {
		return nil, err
	}
	return &Node{n}, nil
}

func (n *Node) QuerySelector(selector string) *Node {
	if strings.HasPrefix(selector, "#") {
		return n.GetElementById(selector[1:])
	}
	if strings.HasPrefix(selector, ".") {
		return n.GetElementByClass(selector[1:])
	}
	return n.GetElementByTagName(selector)
}

func (n *Node) GetElementById(id string) *Node {
	return n.FindOne(func(n *Node) bool {
		return n.HasAttr("id", id)
	})
}

func (n *Node) GetElementByClass(class string) *Node {
	return n.FindOne(func(n *Node) bool {
		return n.HasClass(class)
	})
}

func (n *Node) GetElementByTagName(name string) *Node {
	return n.FindOne(func (n *Node) bool {
		return n.Type == html.ElementNode && n.Data == name
	})
}

func (n *Node) HasClass(class string) bool {
	return n.HasAttr("class", class)
}

func (n *Node) GetAttr(key string) []string {
	var res []string
	for _, attr := range n.Attr {
		if attr.Key == key {
			res = append(res, attr.Val)
		}
	}
	return res
}

func (n *Node) HasAttr(key, value string) bool {
	for _, attr := range n.Attr {
		if attr.Key == key && attr.Val == value {
			return true
		}
	}
	return false
}

func (n *Node) ForEach(cb func (n *Node)) {
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		cb(&Node{c})
	}
}

func (n *Node) ChildNodes() []*Node {
	var res []*Node
	n.ForEach(func(n *Node) {
		res = append(res, n)
	})
	return res
}

func (n *Node) Children() []*Node {
	var res []*Node
	n.ForEach(func(n *Node) {
		if n.Type == html.ElementNode {
			res = append(res, n)
		}
	})
	return res
}

func (n *Node) Traverse(cb func(n *Node)) {
	var f func(*Node)
	f = func(n *Node) {
		cb(n)
		n.ForEach(f)
	}
	f(n)
}

func (n *Node) FindOne(cb func(n *Node) bool) *Node {
	var res *Node

	var f func(*Node)
	f = func(n *Node) {
		if res != nil {
			return
		}
		if cb(n) {
			res = n
			return
		}
		n.ForEach(f)
	}
	f(n)

	return res
}

func (n *Node) FindMany(cb func(n *Node) bool) []*Node {
	var res []*Node

	var f func(*Node)
	f = func(n *Node) {
		if cb(n) {
			res = append(res, n)
		}
		n.ForEach(f)
	}
	f(n)

	return res
}

func (n *Node) Text() string {
	res := ""
	n.Traverse(func(n *Node) {
		if n.Type == html.TextNode {
			res += n.Data
		}
	})
	return res
}

func (n *Node) Render() (string, error) {
	w := bytes.NewBuffer([]byte{})
	err := html.Render(w, n.Node)
	return w.String(), err
}