154 lines
2.5 KiB
Go
154 lines
2.5 KiB
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
|
|
}
|