add iterators for QuerySelector implementations

crhntr · crhntr · commit 9806151e774b · 2024-08-28T16:36:41.000-07:00
diff --git a/document.go b/document.go
@@ -3,6 +3,7 @@ package dom
 import (
 	"strings"
 
+	"github.com/andybalholm/cascadia"
 	"golang.org/x/net/html"
 	"golang.org/x/net/html/atom"
 
@@ -32,6 +33,14 @@ func (d *Document) QuerySelector(query string) spec.Element {
 func (d *Document) QuerySelectorAll(query string) spec.NodeList[spec.Element] {
 	return querySelectorAll(d.node, query, false)
 }
+
+func (d *Document) QuerySelectorEach(query string) spec.NodeIterator[spec.Element] {
+	m := cascadia.MustCompile(query)
+	return func(yield func(spec.Element) bool) {
+		querySelectorEach(d.node, m, yield)
+	}
+}
+
 func (d *Document) Contains(other spec.Node) bool { return contains(d.node, other) }
 
 // TextContent returns an empty string.
diff --git a/document_test.go b/document_test.go
@@ -303,7 +303,7 @@ func TestDocument_QuerySelectorAll(t *testing.T) {
 	document := &Document{node: parsedDocument}
 
 	result := document.QuerySelectorAll(`#1 [data-find-me]`)
-	assert.Equal(t, result.Length(), 2)
+	require.Equal(t, result.Length(), 2)
 	assert.Equal(t, "x1", result.Item(0).ID())
 	assert.Equal(t, "x2", result.Item(1).ID())
 }
diff --git a/element.go b/element.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"strings"
 
+	"github.com/andybalholm/cascadia"
 	"golang.org/x/net/html"
 
 	"github.com/crhntr/dom/spec"
@@ -13,6 +14,13 @@ type Element struct {
 	node *html.Node
 }
 
+func (e *Element) QuerySelectorEach(query string) spec.NodeIterator[spec.Element] {
+	m := cascadia.MustCompile(query)
+	return func(yield func(spec.Element) bool) {
+		querySelectorEach(e.node, m, yield)
+	}
+}
+
 // NewNode
 
 func (e *Element) NodeType() spec.NodeType         { return nodeType(e.node.Type) }
diff --git a/element_test.go b/element_test.go
@@ -13,6 +13,8 @@ import (
 	"github.com/crhntr/dom/spec"
 )
 
+var _ spec.Element = (*Element)(nil)
+
 func TestElement_NodeType(t *testing.T) {
 	// language=html
 	parsedDocument, err := html.Parse(strings.NewReader(`<!DOCTYPE html><html lang="us-en"><head><title></title></head><body><span></span></body</html>`))
diff --git a/fragment.go b/fragment.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"slices"
 
+	"github.com/andybalholm/cascadia"
 	"golang.org/x/net/html"
 
 	"github.com/crhntr/dom/spec"
@@ -127,3 +128,17 @@ func (d *DocumentFragment) QuerySelectorAll(query string) spec.NodeList[spec.Ele
 	}
 	return slices.Clip(list)
 }
+
+func (d *DocumentFragment) QuerySelectorEach(query string) spec.NodeIterator[spec.Element] {
+	m := cascadia.MustCompile(query)
+	return func(yield func(spec.Element) bool) {
+		for _, n := range d.nodes {
+			if m.Match(n) {
+				if !yield(&Element{node: n}) {
+					return
+				}
+			}
+			querySelectorEach(n, m, yield)
+		}
+	}
+}
diff --git a/node.go b/node.go
@@ -435,11 +435,16 @@ func querySelector(node *html.Node, query string, includeParent bool) spec.Eleme
 }
 
 func querySelectorAll(node *html.Node, query string, includeParent bool) nodeListHTMLElements {
-	q := cascadia.MustCompile(query)
-	results := cascadia.QueryAll(node, cascadia.MustCompile(query))
-	if includeParent && node.Type == html.ElementNode && q.Match(node) {
+	m := cascadia.MustCompile(query)
+	var results []*html.Node
+	if includeParent && m.Match(node) {
 		results = slices.Insert(results, 0, node)
 	}
+	querySelectorEach(node, m, func(element spec.Element) bool {
+		el := element.(*Element)
+		results = append(results, el.node)
+		return true
+	})
 	return results
 }
 
@@ -488,3 +493,14 @@ func getAttribute(node *html.Node, name string) string {
 	}
 	return ""
 }
+
+func querySelectorEach(n *html.Node, m cascadia.Matcher, yield func(spec.Element) bool) {
+	for c := n.FirstChild; c != nil; c = c.NextSibling {
+		if m.Match(c) {
+			if !yield(&Element{node: c}) {
+				return
+			}
+		}
+		querySelectorEach(c, m, yield)
+	}
+}
diff --git a/spec/interfaces.go b/spec/interfaces.go
@@ -188,6 +188,7 @@ type ElementQueries interface {
 
 	QuerySelector(query string) Element
 	QuerySelectorAll(query string) NodeList[Element]
+	QuerySelectorEach(query string) NodeIterator[Element]
 }
 
 // Element is based on
@@ -248,6 +249,7 @@ type DocumentFragment interface {
 
 	QuerySelector(query string) Element
 	QuerySelectorAll(query string) NodeList[Element]
+	QuerySelectorEach(query string) NodeIterator[Element]
 }
 
 type Comment interface {
@@ -256,3 +258,5 @@ type Comment interface {
 	Data() string
 	SetData() string
 }
+
+type NodeIterator[N Node] func(func(N) bool)

Original file line number	Diff line number	Diff line change
`@@ -303,7 +303,7 @@ func TestDocument_QuerySelectorAll(t *testing.T) {`
`303`	`303`	`document := &Document{node: parsedDocument}`
`304`	`304`
`305`	`305`	result := document.QuerySelectorAll(`#1 [data-find-me]`)
`306`		`- assert.Equal(t, result.Length(), 2)`
	`306`	`+ require.Equal(t, result.Length(), 2)`
`307`	`307`	`assert.Equal(t, "x1", result.Item(0).ID())`
`308`	`308`	`assert.Equal(t, "x2", result.Item(1).ID())`
`309`	`309`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,8 @@ import (`
`13`	`13`	`"github.com/crhntr/dom/spec"`
`14`	`14`	`)`
`15`	`15`
	`16`	`+var _ spec.Element = (*Element)(nil)`
	`17`	`+`
`16`	`18`	`func TestElement_NodeType(t *testing.T) {`
`17`	`19`	`// language=html`
`18`	`20`	parsedDocument, err := html.Parse(strings.NewReader(`<!DOCTYPE html><html lang="us-en"><head><title></title></head><body><span></span></body</html>`))
Original file line number	Diff line number	Diff line change
`@@ -188,6 +188,7 @@ type ElementQueries interface {`
`188`	`188`
`189`	`189`	`QuerySelector(query string) Element`
`190`	`190`	`QuerySelectorAll(query string) NodeList[Element]`
	`191`	`+ QuerySelectorEach(query string) NodeIterator[Element]`
`191`	`192`	`}`
`192`	`193`
`193`	`194`	`// Element is based on`
`@@ -248,6 +249,7 @@ type DocumentFragment interface {`
`248`	`249`
`249`	`250`	`QuerySelector(query string) Element`
`250`	`251`	`QuerySelectorAll(query string) NodeList[Element]`
	`252`	`+ QuerySelectorEach(query string) NodeIterator[Element]`
`251`	`253`	`}`
`252`	`254`
`253`	`255`	`type Comment interface {`
`@@ -256,3 +258,5 @@ type Comment interface {`
`256`	`258`	`Data() string`
`257`	`259`	`SetData() string`
`258`	`260`	`}`
	`261`	`+`
	`262`	`+type NodeIterator[N Node] func(func(N) bool)`