Implement Single and SingleMatcher

PuerkitoBio · Jun 13, 2021 · 0126a1f · 0126a1f
1 parent e1f2d60
commit 0126a1f
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 1 deletion.
diff --git a/example_test.go b/example_test.go
@@ -80,3 +80,31 @@ func ExampleNewDocumentFromReader_string() {
 
 	// Output: Header
 }
+
+func ExampleSingle() {
+	html := `
+<html>
+  <body>
+    <div>1</div>
+    <div>2</div>
+    <div>3</div>
+  </body>
+</html>
+`
+	doc, err := goquery.NewDocumentFromReader(strings.NewReader(html))
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// By default, the selector string selects all matching nodes
+	multiSel := doc.Find("div")
+	fmt.Println(multiSel.Text())
+
+	// Using goquery.Single, only the first match is selected
+	singleSel := doc.FindMatcher(goquery.Single("div"))
+	fmt.Println(singleSel.Text())
+
+	// Output:
+	// 123
+	// 1
+}
diff --git a/type.go b/type.go
@@ -7,7 +7,6 @@ import (
 	"net/url"
 
 	"github.com/andybalholm/cascadia"
-
 	"golang.org/x/net/html"
 )
 
@@ -122,6 +121,30 @@ type Matcher interface {
 	Filter([]*html.Node) []*html.Node
 }
 
+// Single compiles a selector string to a Matcher that stops after the first
+// match is found.
+//
+// By default, Selection.Find and other functions that accept a selector string
+// will use all matches corresponding to that selector. By using the Matcher
+// returned by Single, at most the first match will be used.
+func Single(selector string) Matcher {
+	return singleMatcher{compileMatcher(selector)}
+}
+
+// SingleMatcher returns a Matcher matches the same nodes as m, but that stops
+// after the first match is found.
+//
+// By default, Selection.FindMatcher and other functions that accept a Matcher
+// will use all corresponding matches. By using the Matcher returned by
+// SingleMatcher, at most the first match will be used.
+func SingleMatcher(m Matcher) Matcher {
+	if _, ok := m.(singleMatcher); ok {
+		// m is already a singleMatcher
+		return m
+	}
+	return singleMatcher{m}
+}
+
 // compileMatcher compiles the selector string s and returns
 // the corresponding Matcher. If s is an invalid selector string,
 // it returns a Matcher that fails all matches.
@@ -133,6 +156,30 @@ func compileMatcher(s string) Matcher {
 	return cs
 }
 
+type singleMatcher struct {
+	Matcher
+}
+
+func (m singleMatcher) MatchAll(n *html.Node) []*html.Node {
+	// Optimized version - stops finding at the first match (cascadia-compiled
+	// matchers all use this code path).
+	if mm, ok := m.Matcher.(interface{ MatchFirst(*html.Node) *html.Node }); ok {
+		node := mm.MatchFirst(n)
+		if node == nil {
+			return nil
+		}
+		return []*html.Node{node}
+	}
+
+	// Fallback version, for e.g. test mocks that don't provide the MatchFirst
+	// method.
+	nodes := m.Matcher.MatchAll(n)
+	if len(nodes) > 0 {
+		return nodes[:1:1]
+	}
+	return nil
+}
+
 // invalidMatcher is a Matcher that always fails to match.
 type invalidMatcher struct{}