internal/testing/htmlcheck/query.go - pkgsite - Git at Google

 // Copyright 2023 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.

 package htmlcheck

 import (
 	"errors"
 	"fmt"
 	"strconv"
 	"strings"

 	"golang.org/x/net/html"
 )

 // A selector represents a parsed css selector that can be used in a query.
 // The atoms all match against a given element and next matches against
 // children of that element. So, for example "div#id a" parses into a selector that
 // has atoms for matching the div and the id and a next that points to another
 // selector that has an atom for "a".
 type selector struct {
 	atoms []selectorAtom
 	next  *selector
 }

 // String returns a string used for debugging test failures.
 func (s *selector) String() string {
 	if s == nil {
 		return "nil"
 	}
 	str := "["
 	for i, atom := range s.atoms {
 		str += fmt.Sprintf("%#v", atom)
 		if i != len(s.atoms)-1 {
 			str += ","
 		}
 	}
 	str += "]->" + s.next.String()
 	return str
 }

 // selectorAtom represents a part of a selector that individually
 // matches a single element name, id, class, or attribute value.
 type selectorAtom interface {
 	match(n *html.Node) bool
 }

 // query returns the first node in n that matches the given selector,
 // or nil if there are no nodes matching the selector.
 func query(n *html.Node, selector *selector) *html.Node {
 	allMatch := true
 	for _, atom := range selector.atoms {
 		if !atom.match(n) {
 			allMatch = false
 			break
 		}
 	}
 	if allMatch {
 		if selector.next != nil {
 			if result := queryChildren(n, selector.next); result != nil {
 				return result
 			}
 		} else {
 			return n
 		}
 	}
 	return queryChildren(n, selector)
 }

 func queryChildren(n *html.Node, selector *selector) *html.Node {
 	child := n.FirstChild
 	for child != nil {
 		if result := query(child, selector); result != nil {
 			return result
 		}
 		child = child.NextSibling
 	}
 	return nil
 }

 // parse parses the string into a selector. It matches the following
 // atoms: element, #id, .class, [attribute="value"]. It allows the atoms
 // to be combined where they all need to match (for example, a#id) and
 // for nested selectors to be combined with a space.
 // For simplicity, the selector must not have any non-ASCII bytes.
 func parse(s string) (*selector, error) {
 	sel := &selector{}
 	if !isAscii(s) {
 		return nil, errors.New("non ascii byte found in selector string")
 	}
 	for len(s) > 0 {
 		switch {
 		case isLetter(s[0]):
 			ident, rest := consumeIdentifier(s)
 			sel.atoms = append(sel.atoms, &elementSelector{ident})
 			s = rest
 		case s[0] == '.':
 			ident, rest := consumeIdentifier(s[1:])
 			if len(ident) == 0 {
 				return nil, errors.New("no class name after '.'")
 			}
 			sel.atoms = append(sel.atoms, &classSelector{ident})
 			s = rest
 		case s[0] == '#':
 			ident, rest := consumeIdentifier(s[1:])
 			if len(ident) == 0 {
 				return nil, errors.New("no id name after '#'")
 			}
 			sel.atoms = append(sel.atoms, &idSelector{ident})
 			s = rest
 		case s[0] == '[':
 			attributeSelector, rest, err := parseAttributeSelector(s)
 			if err != nil {
 				return nil, err
 			}
 			sel.atoms = append(sel.atoms, attributeSelector)
 			s = rest
 		case s[0] == ' ':
 			s = strings.TrimLeft(s, " ")
 			next, err := parse(s)
 			if err != nil {
 				return nil, err
 			}
 			sel.next = next
 			return sel, nil
 		case s[0] == ':':
 			atom, rest, err := parsePseudoClass(s)
 			if err != nil {
 				return nil, err
 			}
 			sel.atoms = append(sel.atoms, atom)
 			s = rest
 		default:
 			return nil, fmt.Errorf("unexpected character '%v' in input", string(s[0]))
 		}
 	}
 	return sel, nil
 }

 // parsePseudoClass parses a :nth-child(n) or :nth-of-type(n). It only supports those functions and
 // number arguments to those functions, not even, odd, or an+b.
 func parsePseudoClass(s string) (selectorAtom, string, error) {
 	if s[0] != ':' {
 		return nil, "", errors.New("expected ':' at beginning of pseudo class")
 	}
 	ident, rest := consumeIdentifier(s[1:])
 	if len(ident) == 0 {
 		return nil, "", errors.New("expected identifier after : in pseudo class")
 	}
 	if ident != "nth-of-type" && ident != "nth-child" {
 		return nil, "", errors.New("only :nth-of-type() and :nth-child() pseudoclasses are supported")
 	}
 	s = rest
 	if len(s) == 0 || s[0] != '(' {
 		return nil, "", errors.New("expected '(' after :nth-of-type or nth-child")
 	}
 	numstr, rest := consumeNumber(s[1:])
 	if len(numstr) == 0 {
 		return nil, "", errors.New("only number arguments are supported for :nth-of-type() or :nth-child()")
 	}
 	num, err := strconv.Atoi(numstr)
 	if err != nil {
 		// This shouldn't ever happen because consumeNumber should only return valid numbers and we
 		// check that the length is greater than 0.
 		panic(fmt.Errorf("unexpected parse error for number: %v", err))
 	}
 	s = rest
 	if len(s) == 0 || s[0] != ')' {
 		return nil, "", errors.New("expected ')' after number argument to nth-of-type or nth-child")
 	}
 	rest = s[1:]
 	switch ident {
 	case "nth-of-type":
 		return &nthOfType{num}, rest, nil
 	case "nth-child":
 		return &nthChild{num}, rest, nil
 	default:
 		panic("we should only have allowed nth-of-type or nth-child up to this point")
 	}
 }

 // parseAttributeSelector parses an attribute selector of the form [attribute-name="attribute=value"]
 func parseAttributeSelector(s string) (*attributeSelector, string, error) {
 	if s[0] != '[' {
 		return nil, "", errors.New("expected '[' at beginning of attribute selector")
 	}
 	ident, rest := consumeIdentifier(s[1:])
 	if len(ident) == 0 {
 		return nil, "", errors.New("expected attribute name after '[' in attribute selector")
 	}
 	attributeName := ident
 	s = rest
 	if len(s) == 0 || s[0] != '=' {
 		return nil, "", errors.New("expected '=' after attribute name in attribute selector")
 	}
 	s = s[1:]
 	if len(s) == 0 || s[0] != '"' {
 		return nil, "", errors.New("expected '\"' after = in attribute selector")
 	}
 	s = s[1:]
 	i := 0
 	for ; i < len(s) && s[i] != '"'; i++ {
 	}
 	attributeValue, s := s[:i], s[i:]
 	if len(s) == 0 || s[0] != '"' {
 		return nil, "", errors.New("expected '\"' after attribute value")
 	}
 	s = s[1:]
 	if len(s) == 0 || s[0] != ']' {
 		return nil, "", errors.New("expected ']' at end of attribute selector")
 	}
 	s = s[1:]
 	return &attributeSelector{attribute: attributeName, value: attributeValue}, s, nil
 }

 func isLetter(b byte) bool {
 	return ('a' <= b && b <= 'z') || ('A' <= b && b <= 'Z')
 }

 func isNumber(b byte) bool {
 	return ('0' <= b && b <= '9')
 }

 // consumeIdentifier consumes and returns a identifier at the beginning
 // of the given string, and the rest of the string.
 func consumeIdentifier(s string) (letters, rest string) {
 	i := 0
 	for ; i < len(s); i++ {
 		// must start with letter or hyphen or underscore
 		if i == 0 {
 			if !(isLetter(s[i]) || s[i] == '-' || s[i] == '_') {
 				break
 			}
 		} else {
 			if !(isLetter(s[i]) || isNumber(s[i]) || s[i] == '-' || s[i] == '_') {
 				break
 			}
 		}
 		// CSS doesn't allow identifiers to start with two hyphens or a hyphen
 		// followed by a digit, but we'll allow it.
 	}
 	return s[:i], s[i:]
 }

 // consumeNumber consumes and returns a (0-9)+ number at the beginning
 // of the given string, and the rest of the string.
 func consumeNumber(s string) (letters, rest string) {
 	i := 0
 	for ; i < len(s) && isNumber(s[i]); i++ {
 	}
 	return s[:i], s[i:]
 }

 func isAscii(s string) bool {
 	for i := 0; i < len(s); i++ {
 		if s[i] > 127 {
 			return false
 		}
 	}
 	return true
 }

 // elementSelector matches a node that has the given element name.
 type elementSelector struct {
 	name string
 }

 func (s *elementSelector) match(n *html.Node) bool {
 	if n.Type != html.ElementNode {
 		return false
 	}
 	return n.Data == s.name
 }

 type attributeSelector struct {
 	attribute, value string
 }

 // attributeSelector matches a node with an attribute that has a given value.
 func (s *attributeSelector) match(n *html.Node) bool {
 	if n.Type != html.ElementNode {
 		return false
 	}
 	for _, attr := range n.Attr {
 		if attr.Key == s.attribute {
 			return attr.Val == s.value
 		}
 	}
 	return false
 }

 // idSelector matches an element that has the given id.
 type idSelector struct {
 	id string
 }

 func (s *idSelector) match(n *html.Node) bool {
 	if n.Type != html.ElementNode {
 		return false
 	}
 	for _, attr := range n.Attr {
 		if attr.Key == "id" {
 			return attr.Val == s.id
 		}
 	}
 	return false
 }

 // classSelector matches an element that has the given class set on it.
 type classSelector struct {
 	class string
 }

 func (s *classSelector) match(n *html.Node) bool {
 	if n.Type != html.ElementNode {
 		return false
 	}
 	for _, attr := range n.Attr {
 		if attr.Key == "class" {
 			for _, f := range strings.Fields(attr.Val) {
 				if f == s.class {
 					return true
 				}
 			}
 			break
 		}
 	}
 	return false
 }

 // nthOfType implements the :nth-of-type() pseudoclass.
 type nthOfType struct {
 	n int
 }

 func (s *nthOfType) match(n *html.Node) bool {
 	if n.Type != html.ElementNode || n.Parent == nil {
 		return false
 	}
 	curChild := n.Parent.FirstChild
 	i := 0
 	for {
 		if curChild.Type == html.ElementNode && curChild.Data == n.Data {
 			i++
 			if i == s.n {
 				break
 			}
 		}
 		if curChild.NextSibling == nil {
 			break
 		}
 		curChild = curChild.NextSibling
 	}
 	if i != s.n {
 		// there were fewer than n children of this element type.
 		return false
 	}
 	return curChild == n
 }

 // nthChild implements the :nth-child() pseudoclass
 type nthChild struct {
 	n int
 }

 func (s *nthChild) match(n *html.Node) bool {
 	if n.Type != html.ElementNode || n.Parent == nil {
 		return false
 	}
 	curChild := n.Parent.FirstChild
 	// Advance to next element node.
 	for curChild.Type != html.ElementNode && curChild.NextSibling != nil {
 		curChild = curChild.NextSibling
 	}
 	i := 1
 	for ; i < s.n; i++ {
 		if curChild.NextSibling == nil {
 			break
 		}
 		curChild = curChild.NextSibling
 		// Advance to next element node.
 		for curChild.Type != html.ElementNode && curChild.NextSibling != nil {
 			curChild = curChild.NextSibling
 		}
 	}
 	if i != s.n {
 		// there were fewer than n children.
 		return false
 	}
 	return curChild == n
 }
	// Copyright 2023 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.

	package htmlcheck

	import (
	"errors"
	"fmt"
	"strconv"
	"strings"

	"golang.org/x/net/html"
	)

	// A selector represents a parsed css selector that can be used in a query.
	// The atoms all match against a given element and next matches against
	// children of that element. So, for example "div#id a" parses into a selector that
	// has atoms for matching the div and the id and a next that points to another
	// selector that has an atom for "a".
	type selector struct {
	atoms []selectorAtom
	next *selector
	}

	// String returns a string used for debugging test failures.
	func (s *selector) String() string {
	if s == nil {
	return "nil"
	}
	str := "["
	for i, atom := range s.atoms {
	str += fmt.Sprintf("%#v", atom)
	if i != len(s.atoms)-1 {
	str += ","
	}
	}
	str += "]->" + s.next.String()
	return str
	}

	// selectorAtom represents a part of a selector that individually
	// matches a single element name, id, class, or attribute value.
	type selectorAtom interface {
	match(n *html.Node) bool
	}

	// query returns the first node in n that matches the given selector,
	// or nil if there are no nodes matching the selector.
	func query(n html.Node, selector selector) *html.Node {
	allMatch := true
	for _, atom := range selector.atoms {
	if !atom.match(n) {
	allMatch = false
	break
	}
	}
	if allMatch {
	if selector.next != nil {
	if result := queryChildren(n, selector.next); result != nil {
	return result
	}
	} else {
	return n
	}
	}
	return queryChildren(n, selector)
	}

	func queryChildren(n html.Node, selector selector) *html.Node {
	child := n.FirstChild
	for child != nil {
	if result := query(child, selector); result != nil {
	return result
	}
	child = child.NextSibling
	}
	return nil
	}

	// parse parses the string into a selector. It matches the following
	// atoms: element, #id, .class, [attribute="value"]. It allows the atoms
	// to be combined where they all need to match (for example, a#id) and
	// for nested selectors to be combined with a space.
	// For simplicity, the selector must not have any non-ASCII bytes.
	func parse(s string) (*selector, error) {
	sel := &selector{}
	if !isAscii(s) {
	return nil, errors.New("non ascii byte found in selector string")
	}
	for len(s) > 0 {
	switch {
	case isLetter(s[0]):
	ident, rest := consumeIdentifier(s)
	sel.atoms = append(sel.atoms, &elementSelector{ident})
	s = rest
	case s[0] == '.':
	ident, rest := consumeIdentifier(s[1:])
	if len(ident) == 0 {
	return nil, errors.New("no class name after '.'")
	}
	sel.atoms = append(sel.atoms, &classSelector{ident})
	s = rest
	case s[0] == '#':
	ident, rest := consumeIdentifier(s[1:])
	if len(ident) == 0 {
	return nil, errors.New("no id name after '#'")
	}
	sel.atoms = append(sel.atoms, &idSelector{ident})
	s = rest
	case s[0] == '[':
	attributeSelector, rest, err := parseAttributeSelector(s)
	if err != nil {
	return nil, err
	}
	sel.atoms = append(sel.atoms, attributeSelector)
	s = rest
	case s[0] == ' ':
	s = strings.TrimLeft(s, " ")
	next, err := parse(s)
	if err != nil {
	return nil, err
	}
	sel.next = next
	return sel, nil
	case s[0] == ':':
	atom, rest, err := parsePseudoClass(s)
	if err != nil {
	return nil, err
	}
	sel.atoms = append(sel.atoms, atom)
	s = rest
	default:
	return nil, fmt.Errorf("unexpected character '%v' in input", string(s[0]))
	}
	}
	return sel, nil
	}

	// parsePseudoClass parses a :nth-child(n) or :nth-of-type(n). It only supports those functions and
	// number arguments to those functions, not even, odd, or an+b.
	func parsePseudoClass(s string) (selectorAtom, string, error) {
	if s[0] != ':' {
	return nil, "", errors.New("expected ':' at beginning of pseudo class")
	}
	ident, rest := consumeIdentifier(s[1:])
	if len(ident) == 0 {
	return nil, "", errors.New("expected identifier after : in pseudo class")
	}
	if ident != "nth-of-type" && ident != "nth-child" {
	return nil, "", errors.New("only :nth-of-type() and :nth-child() pseudoclasses are supported")
	}
	s = rest
	if len(s) == 0 \|\| s[0] != '(' {
	return nil, "", errors.New("expected '(' after :nth-of-type or nth-child")
	}
	numstr, rest := consumeNumber(s[1:])
	if len(numstr) == 0 {
	return nil, "", errors.New("only number arguments are supported for :nth-of-type() or :nth-child()")
	}
	num, err := strconv.Atoi(numstr)
	if err != nil {
	// This shouldn't ever happen because consumeNumber should only return valid numbers and we
	// check that the length is greater than 0.
	panic(fmt.Errorf("unexpected parse error for number: %v", err))
	}
	s = rest
	if len(s) == 0 \|\| s[0] != ')' {
	return nil, "", errors.New("expected ')' after number argument to nth-of-type or nth-child")
	}
	rest = s[1:]
	switch ident {
	case "nth-of-type":
	return &nthOfType{num}, rest, nil
	case "nth-child":
	return &nthChild{num}, rest, nil
	default:
	panic("we should only have allowed nth-of-type or nth-child up to this point")
	}
	}

	// parseAttributeSelector parses an attribute selector of the form [attribute-name="attribute=value"]
	func parseAttributeSelector(s string) (*attributeSelector, string, error) {
	if s[0] != '[' {
	return nil, "", errors.New("expected '[' at beginning of attribute selector")
	}
	ident, rest := consumeIdentifier(s[1:])
	if len(ident) == 0 {
	return nil, "", errors.New("expected attribute name after '[' in attribute selector")
	}
	attributeName := ident
	s = rest
	if len(s) == 0 \|\| s[0] != '=' {
	return nil, "", errors.New("expected '=' after attribute name in attribute selector")
	}
	s = s[1:]
	if len(s) == 0 \|\| s[0] != '"' {
	return nil, "", errors.New("expected '\"' after = in attribute selector")
	}
	s = s[1:]
	i := 0
	for ; i < len(s) && s[i] != '"'; i++ {
	}
	attributeValue, s := s[:i], s[i:]
	if len(s) == 0 \|\| s[0] != '"' {
	return nil, "", errors.New("expected '\"' after attribute value")
	}
	s = s[1:]
	if len(s) == 0 \|\| s[0] != ']' {
	return nil, "", errors.New("expected ']' at end of attribute selector")
	}
	s = s[1:]
	return &attributeSelector{attribute: attributeName, value: attributeValue}, s, nil
	}

	func isLetter(b byte) bool {
	return ('a' <= b && b <= 'z') \|\| ('A' <= b && b <= 'Z')
	}

	func isNumber(b byte) bool {
	return ('0' <= b && b <= '9')
	}

	// consumeIdentifier consumes and returns a identifier at the beginning
	// of the given string, and the rest of the string.
	func consumeIdentifier(s string) (letters, rest string) {
	i := 0
	for ; i < len(s); i++ {
	// must start with letter or hyphen or underscore
	if i == 0 {
	if !(isLetter(s[i]) \|\| s[i] == '-' \|\| s[i] == '_') {
	break
	}
	} else {
	if !(isLetter(s[i]) \|\| isNumber(s[i]) \|\| s[i] == '-' \|\| s[i] == '_') {
	break
	}
	}
	// CSS doesn't allow identifiers to start with two hyphens or a hyphen
	// followed by a digit, but we'll allow it.
	}
	return s[:i], s[i:]
	}

	// consumeNumber consumes and returns a (0-9)+ number at the beginning
	// of the given string, and the rest of the string.
	func consumeNumber(s string) (letters, rest string) {
	i := 0
	for ; i < len(s) && isNumber(s[i]); i++ {
	}
	return s[:i], s[i:]
	}

	func isAscii(s string) bool {
	for i := 0; i < len(s); i++ {
	if s[i] > 127 {
	return false
	}
	}
	return true
	}

	// elementSelector matches a node that has the given element name.
	type elementSelector struct {
	name string
	}

	func (s elementSelector) match(n html.Node) bool {
	if n.Type != html.ElementNode {
	return false
	}
	return n.Data == s.name
	}

	type attributeSelector struct {
	attribute, value string
	}

	// attributeSelector matches a node with an attribute that has a given value.
	func (s attributeSelector) match(n html.Node) bool {
	if n.Type != html.ElementNode {
	return false
	}
	for _, attr := range n.Attr {
	if attr.Key == s.attribute {
	return attr.Val == s.value
	}
	}
	return false
	}

	// idSelector matches an element that has the given id.
	type idSelector struct {
	id string
	}

	func (s idSelector) match(n html.Node) bool {
	if n.Type != html.ElementNode {
	return false
	}
	for _, attr := range n.Attr {
	if attr.Key == "id" {
	return attr.Val == s.id
	}
	}
	return false
	}

	// classSelector matches an element that has the given class set on it.
	type classSelector struct {
	class string
	}

	func (s classSelector) match(n html.Node) bool {
	if n.Type != html.ElementNode {
	return false
	}
	for _, attr := range n.Attr {
	if attr.Key == "class" {
	for _, f := range strings.Fields(attr.Val) {
	if f == s.class {
	return true
	}
	}
	break
	}
	}
	return false
	}

	// nthOfType implements the :nth-of-type() pseudoclass.
	type nthOfType struct {
	n int
	}

	func (s nthOfType) match(n html.Node) bool {
	if n.Type != html.ElementNode \|\| n.Parent == nil {
	return false
	}
	curChild := n.Parent.FirstChild
	i := 0
	for {
	if curChild.Type == html.ElementNode && curChild.Data == n.Data {
	i++
	if i == s.n {
	break
	}
	}
	if curChild.NextSibling == nil {
	break
	}
	curChild = curChild.NextSibling
	}
	if i != s.n {
	// there were fewer than n children of this element type.
	return false
	}
	return curChild == n
	}

	// nthChild implements the :nth-child() pseudoclass
	type nthChild struct {
	n int
	}

	func (s nthChild) match(n html.Node) bool {
	if n.Type != html.ElementNode \|\| n.Parent == nil {
	return false
	}
	curChild := n.Parent.FirstChild
	// Advance to next element node.
	for curChild.Type != html.ElementNode && curChild.NextSibling != nil {
	curChild = curChild.NextSibling
	}
	i := 1
	for ; i < s.n; i++ {
	if curChild.NextSibling == nil {
	break
	}
	curChild = curChild.NextSibling
	// Advance to next element node.
	for curChild.Type != html.ElementNode && curChild.NextSibling != nil {
	curChild = curChild.NextSibling
	}
	}
	if i != s.n {
	// there were fewer than n children.
	return false
	}
	return curChild == n
	}