forked from nkanaev/yarr
-
Notifications
You must be signed in to change notification settings - Fork 0
/
query.go
86 lines (72 loc) · 1.63 KB
/
query.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
package htmlutil
import (
"regexp"
"strings"
"golang.org/x/net/html"
)
var nodeNameRegex = regexp.MustCompile(`\w+|\*`)
func FindNodes(node *html.Node, match func(*html.Node) bool) []*html.Node {
nodes := make([]*html.Node, 0)
queue := make([]*html.Node, 0)
queue = append(queue, node)
for len(queue) > 0 {
var n *html.Node
n, queue = queue[0], queue[1:]
if match(n) {
nodes = append(nodes, n)
}
for c := n.FirstChild; c != nil; c = c.NextSibling {
queue = append(queue, c)
}
}
return nodes
}
func Query(node *html.Node, sel string) []*html.Node {
matcher := NewMatcher(sel)
return FindNodes(node, matcher.Match)
}
func Closest(node *html.Node, sel string) *html.Node {
matcher := NewMatcher(sel)
for cur := node; cur != nil; cur = cur.Parent {
if matcher.Match(cur) {
return cur
}
}
return nil
}
func NewMatcher(sel string) Matcher {
multi := MultiMatch{}
parts := strings.Split(sel, ",")
for _, part := range parts {
part := strings.TrimSpace(part)
if nodeNameRegex.MatchString(part) {
multi.Add(ElementMatch{Name: part})
} else {
panic("unsupported selector: " + part)
}
}
return multi
}
type Matcher interface {
Match(*html.Node) bool
}
type ElementMatch struct {
Name string
}
func (m ElementMatch) Match(n *html.Node) bool {
return n.Type == html.ElementNode && (n.Data == m.Name || m.Name == "*")
}
type MultiMatch struct {
matchers []Matcher
}
func (m *MultiMatch) Add(matcher Matcher) {
m.matchers = append(m.matchers, matcher)
}
func (m MultiMatch) Match(n *html.Node) bool {
for _, matcher := range m.matchers {
if matcher.Match(n) {
return true
}
}
return false
}