gex/gen/domquery/node.go

package domquery

import (
	"strings"
)

type Node struct {
	token    string
	Children []*Node
	Parent   *Node
}

func (node *Node) TagName() string {
	return getTagName(node.token)
}

func (node *Node) OuterHTML() string {
	str := node.token
	str += node.InnerHTML()
	str += getCloseNode(node).token
	return str
}

func (node *Node) SetOuterHTML(str string) {
	closeTag := getCloseNode(node)
	closeTag.token = ""

	//use the token to insert str
	node.token = str
	node.Children = []*Node{}

	// rebuild parent
	tokList := GetTokenList(node.Parent.InnerHTML())
	node.Parent.Children = []*Node{}
	node.Parent = BuildTree(tokList, node.Parent)
}

func (node *Node) InnerHTML() string {

	str := ""

	for _, child := range node.Children {
		str += child.token
		str += child.InnerHTML()
	}

	return str
}

func (node *Node) SetInnerHTML(str string) {
	node.Children = []*Node{}
	tokList := GetTokenList(str)
	BuildTree(tokList, node)
}

func (node *Node) QuerySelector(str string) *Node {
	found := &Node{}

	for _, child := range node.Children {
		if matchSelector(child, str) {
			return child
		}
		found := child.QuerySelector(str)

		if found.token != "" {
			return found
		}
	}

	return found
}

func (node *Node) QuerySelectorAll(str string) []*Node {
	found := []*Node{}

	for _, child := range node.Children {
		if getTagName(child.token) == str && (getTokType(child.token) == "open" || getTokType(child.token) == "selfclosing") {
			found = append(found, child)
		}
		if str == "*" && (getTokType(child.token) == "open" || getTokType(child.token) == "selfclosing") {
			found = append(found, child)
		}
		found = append(found, child.QuerySelectorAll(str)...)
	}

	return found
}

func (node *Node) GetAttribute(str string) string {

	for key, val := range node.AttributeList() {
		if key == str {
			return val
		}
	}
	return ""
}

func (node *Node) HasAttribute(str string) bool {

	for key := range node.AttributeList() {
		if key == str {
			return true
		}
	}
	return false
}

func (node *Node) SetAttribute(attr string, val string) {
	attributes := node.AttributeList()
	attributes[attr] = val

	attrStr := ""
	for key, val := range attributes {
		attrStr += key
		if val != "" {
			attrStr += "=" + "\"" + val + "\""
		}
		attrStr += " "

	}
	attrStr = strings.Trim(attrStr, " ")
	node.token = "<" + getTagName(node.token) + " " + attrStr + ">"
}

func (node *Node) RemoveAttribute(attr string) {
	attributes := node.AttributeList()

	attrStr := ""
	for key, val := range attributes {
		if key == attr {
			continue
		}
		attrStr += key
		if val != "" {
			attrStr += "=" + "\"" + val + "\""
		}
		attrStr += " "

	}
	attrStr = strings.Trim(attrStr, " ")
	node.token = "<" + getTagName(node.token) + " " + attrStr + ">"
}

func (node *Node) AttributeList() map[string]string {
	attributesList := make(map[string]string)

	if getTokType(node.token) != "open" && getTokType(node.token) != "selfclosing" {
		return attributesList
	}

	attrStr := strings.Replace(node.token, "<"+getTagName(node.token), "", 1)
	attrStr = strings.Trim(attrStr, "<> ")

	tok := ""
	state := "key"
	strKey := ""

	for i := 0; i < len(attrStr); i++ {
		chr := string(attrStr[i])

		if state == "key" {
			if i == len(attrStr)-1 {
				tok += chr
				strKey = strings.Trim(tok, " ")
				attributesList[strKey] = ""
				continue
			}
			if chr == "=" {
				strKey = strings.Trim(tok, " ")
				state = "qual"
				tok = ""
				continue
			}
			if chr == " " {
				strKey = strings.Trim(tok, " ")
				if strKey != "" {
					attributesList[strKey] = ""
				}
				strKey = ""
				tok = ""
				continue
			}
		}

		if state == "qual" {
			if chr == "\"" {
				state = "val"
				tok = ""
				continue
			}
		}

		if state == "val" {
			if chr == "\"" {
				attributesList[strKey] = strings.Trim(tok, " ")
				state = "key"
				tok = ""
				continue
			}
		}

		tok += chr
	}

	return attributesList
}

func (node *Node) ClassList() []string {
	classStr := node.GetAttribute("class")
	classList := strings.Split(classStr, " ")

	return classList
}

func (node *Node) ClassListAdd(classStr string) {
	classList := append(node.ClassList(), classStr)

	str := strings.Join(classList, " ")
	str = strings.Trim(str, " ")
	node.SetAttribute("class", str)
}

func (node *Node) ClassListRemove(classStr string) {
	classList := []string{}

	for _, class := range node.ClassList() {
		if class != classStr {
			classList = append(classList, class)
		}
	}

	str := strings.Join(classList, " ")

	str = strings.Trim(str, " ")
	node.SetAttribute("class", str)
}

func (node *Node) Remove() {
	closeNode := getCloseNode(node)

	i := 0
	for _, n := range node.Parent.Children {
		if n == closeNode {
			node.Parent.Children = append(node.Parent.Children[:i], node.Parent.Children[i+1:]...)
			break
		}
		i++
	}

	i = 0
	for _, n := range node.Parent.Children {
		if n == node {
			node.Parent.Children = append(node.Parent.Children[:i], node.Parent.Children[i+1:]...)
			break
		}
		i++
	}
}

func (node *Node) CreateElement(tagName string) *Node {
	newNode := &Node{}
	newNode.token = "<" + tagName + ">"
	return newNode
}

func (node *Node) Append(newNode *Node) {
	newNode.Parent = node
	node.Children = append(node.Children, newNode)
	if getTokType(newNode.token) != "selfclosing" {
		closeNode := &Node{}
		closeNode.token = "</" + getTagName(newNode.token) + ">"
		node.Children = append(node.Children, closeNode)
	}

}