x/net/html: add Node.All() and Node.ChildNodes()

Fixes #62113
golang · Jun 22, 2024 · 74b75cf · 74b75cf
1 parent 9617c63
commit 74b75cf
Show file tree

Hide file tree

Showing 3 changed files with 130 additions and 6 deletions.
diff --git a/html/example_test.go b/html/example_test.go
@@ -2,6 +2,8 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+//go:build goexperiment.rangefunc
+
 // This example demonstrates parsing HTML data and walking the resulting tree.
 package html_test
 
@@ -19,8 +21,7 @@ func ExampleParse() {
 	if err != nil {
 		log.Fatal(err)
 	}
-	var f func(*html.Node)
-	f = func(n *html.Node) {
+	for n := range doc.All() {
 		if n.Type == html.ElementNode && n.Data == "a" {
 			for _, a := range n.Attr {
 				if a.Key == "href" {
@@ -29,11 +30,8 @@ func ExampleParse() {
 				}
 			}
 		}
-		for c := n.FirstChild; c != nil; c = c.NextSibling {
-			f(c)
-		}
 	}
-	f(doc)
+
 	// Output:
 	// foo
 	// /bar/baz

diff --git a/html/iter.go b/html/iter.go
@@ -0,0 +1,51 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build goexperiment.rangefunc
+
+package html
+
+import "iter"
+
+// ChildNodes returns a sequence yielding the immediate children of n.
+//
+// Mutating a Node while iterating through its ChildNodes may have unexpected results.
+func (n *Node) ChildNodes() iter.Seq[*Node] {
+	return func(yield func(*Node) bool) {
+		if n == nil {
+			return
+		}
+		for c := n.FirstChild; c != nil; c = c.NextSibling {
+			if !yield(c) {
+				return
+			}
+		}
+	}
+
+}
+
+func (n *Node) all(yield func(*Node) bool) bool {
+	if n == nil {
+		return true
+	}
+	if !yield(n) {
+		return false
+	}
+
+	for c := range n.ChildNodes() {
+		if !c.all(yield) {
+			return false
+		}
+	}
+	return true
+}
+
+// All returns a sequence yielding all descendents of n in depth-first pre-order.
+//
+// Mutating a Node while iterating through it or its descendents may have unexpected results.
+func (n *Node) All() iter.Seq[*Node] {
+	return func(yield func(*Node) bool) {
+		n.all(yield)
+	}
+}
diff --git a/html/iter_test.go b/html/iter_test.go
@@ -0,0 +1,75 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build goexperiment.rangefunc
+
+package html
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestNode_ChildNodes(t *testing.T) {
+	tests := []struct {
+		in   string
+		want string
+	}{
+		{"", ""},
+		{"<a></a>", ""},
+		{"<a><b></b></a>", "b"},
+		{"<a>b</a>", "b"},
+		{"<a><!--b--></a>", "b"},
+		{"<a>b<c></c>d</a>", "b c d"},
+		{"<a>b<c><!--d--></c>e</a>", "b c e"},
+		{"<a><b><c>d<!--e-->f</c></b>g<!--h--><i>j</i></a>", "b g h i"},
+	}
+	for _, test := range tests {
+		doc, err := Parse(strings.NewReader(test.in))
+		if err != nil {
+			t.Fatal(err)
+		}
+		// Drill to <html><head></head><body> test.in
+		n := doc.FirstChild.FirstChild.NextSibling.FirstChild
+		var results []string
+		for c := range n.ChildNodes() {
+			results = append(results, c.Data)
+		}
+		if got := strings.Join(results, " "); got != test.want {
+			t.Errorf("unexpected children yielded by ChildNodes; want: %q got: %q", test.want, got)
+		}
+	}
+}
+
+func TestNode_All(t *testing.T) {
+	tests := []struct {
+		in   string
+		want string
+	}{
+		{"", ""},
+		{"<a></a>", "a"},
+		{"<a><b></b></a>", "a b"},
+		{"<a>b</a>", "a b"},
+		{"<a><!--b--></a>", "a b"},
+		{"<a>b<c></c>d</a>", "a b c d"},
+		{"<a>b<c><!--d--></c>e</a>", "a b c d e"},
+		{"<a><b><c>d<!--e-->f</c></b>g<!--h--><i>j</i></a>", "a b c d e f g h i j"},
+	}
+	for _, test := range tests {
+		doc, err := Parse(strings.NewReader(test.in))
+		if err != nil {
+			t.Fatal(err)
+		}
+		// Drill to <html><head></head><body> test.in
+		n := doc.FirstChild.FirstChild.NextSibling.FirstChild
+		var results []string
+		for c := range n.All() {
+			results = append(results, c.Data)
+		}
+		if got := strings.Join(results, " "); got != test.want {
+			t.Errorf("unexpected children yielded by All; want: %q got: %q",
+				test.want, got)
+		}
+	}
+}