property: Add/Remove/Toggle Class functions

property.go

@@ -2,10 +2,14 @@ package goquery
 
 import (
 	"bytes"
+	"regexp"
+	"strings"
 
 	"code.google.com/p/go.net/html"
 )
 
+var rxClassTrim = regexp.MustCompile("[\t\r\n]")
+
 // Attr gets the specified attribute's value for the first element in the
 // Selection. To get the value for each element individually, use a looping
 // construct such as Each or Map method.
@@ -58,6 +62,83 @@ func (s *Selection) Html() (ret string, e error) {
 	return
 }
 
+// Add the given class(es) to each element in the set of matched elements.
+func (s *Selection) AddClass(class string) *Selection {
+	rclasses := getClassesSlice(class)
+
+	for _, n := range s.Nodes {
+		classes, attr := getClassesAndAttr(n, true)
+		for _, rcl := range rclasses {
+			if strings.Index(classes, " "+rcl+" ") == -1 {
+				classes += rcl + " "
+			}
+		}
+
+		setClasses(n, attr, classes)
+	}
+
+	return s
+}
+
+// HasClass determines whether any of the matched elements are assigned the
+// given class.
+func (s *Selection) HasClass(class string) bool {
+	class = " " + class + " "
+	for _, n := range s.Nodes {
+		classes, _ := getClassesAndAttr(n, false)
+		if strings.Index(classes, class) > -1 {
+			return true
+		}
+	}
+	return false
+}
+
+// Remove the given class(es) from each element in the set of matched elements.
+func (s *Selection) RemoveClass(class string) *Selection {
+	rclasses := getClassesSlice(class)
+
+	for _, n := range s.Nodes {
+		classes, attr := getClassesAndAttr(n, true)
+		for _, rcl := range rclasses {
+			classes = strings.Replace(classes, rcl, "", -1)
+		}
+
+		setClasses(n, attr, classes)
+	}
+
+	return s
+}
+
+// Remove all classes from each element in the set of matched elements.
+func (s *Selection) RemoveClasses() *Selection {
+	for _, n := range s.Nodes {
+		_, attr := getClassesAndAttr(n, false)
+		setClasses(n, attr, "")
+	}
+
+	return s
+}
+
+// Add or remove the given class(es) for each element in the set of matched elements.
+func (s *Selection) ToggleClass(class string) *Selection {
+	tcls := getClassesSlice(class)
+
+	for _, n := range s.Nodes {
+		classes, attr := getClassesAndAttr(n, true)
+		for _, tcl := range tcls {
+			if strings.Index(classes, tcl) != -1 {
+				classes = strings.Replace(classes, tcl, "", -1)
+			} else {
+				classes += tcl + " "
+			}
+		}
+
+		setClasses(n, attr, classes)
+	}
+
+	return s
+}
+
 // Get the specified node's text content.
 func getNodeText(node *html.Node) string {
 	if node.Type == html.TextNode {
@@ -74,18 +155,70 @@ func getNodeText(node *html.Node) string {
 	return ""
 }
 
-// Private function to get the specified attribute's value from a node.
-func getAttributeValue(attrName string, n *html.Node) (val string, exists bool) {
+func getAttribute(attrName string, n *html.Node) (attr *html.Attribute, exists bool) {
 	if n == nil {
 		return
 	}
 
-	for _, a := range n.Attr {
+	for i, a := range n.Attr {
 		if a.Key == attrName {
-			val = a.Val
+			attr = &n.Attr[i]
 			exists = true
 			return
 		}
 	}
+
 	return
 }
+
+// Private function to get the specified attribute's value from a node.
+func getAttributeValue(attrName string, n *html.Node) (val string, exists bool) {
+	if a, ok := getAttribute(attrName, n); ok {
+		val = a.Val
+		exists = true
+	}
+	return
+}
+
+// Get and normalize the "class" attribute from the node.
+func getClassesAndAttr(n *html.Node, create bool) (classes string, attr *html.Attribute) {
+	// Applies only to element nodes
+	if n.Type == html.ElementNode {
+		attr, _ = getAttribute("class", n)
+		if attr == nil && create {
+			n.Attr = append(n.Attr, html.Attribute{
+				Key: "class",
+				Val: "",
+			})
+			attr, _ = getAttribute("class", n)
+		}
+	}
+
+	if attr == nil {
+		classes = " "
+	} else {
+		classes = rxClassTrim.ReplaceAllString(" "+attr.Val+" ", " ")
+	}
+
+	return
+}
+
+func getClassesSlice(classes string) []string {
+	return strings.Split(rxClassTrim.ReplaceAllString(" "+classes+" ", " "), " ")
+}
+
+func setClasses(n *html.Node, attr *html.Attribute, classes string) {
+	classes = strings.TrimSpace(classes)
+
+	if classes == "" {
+		for i, a := range n.Attr {
+			if a.Key == "class" {
+				n.Attr[i], n.Attr[len(n.Attr)-1], n.Attr =
+					n.Attr[len(n.Attr)-1], html.Attribute{}, n.Attr[:len(n.Attr)-1]
+				return
+			}
+		}
+	} else {
+		attr.Val = classes
+	}
+}

property_test.go

@@ -83,3 +83,98 @@ func TestNbsp(t *testing.T) {
 		t.Errorf("Html: expected a non-breaking space at index 4, got %d", ix)
 	}
 }
+
+func TestAddClass(t *testing.T) {
+	sel := Doc2Clone().Find("#main")
+	sel.AddClass("main main main")
+
+	// Make sure that class was only added once
+	if a, ok := sel.Attr("class"); !ok || a != "main" {
+		t.Error("Expected #main to have class main")
+	}
+}
+
+func TestAddEmptyClass(t *testing.T) {
+	sel := Doc2Clone().Find("#main")
+	sel.AddClass("")
+
+	// Make sure that class was only added once
+	if a, ok := sel.Attr("class"); ok {
+		t.Errorf("Expected #main to not to have a class, have: %s", a)
+	}
+}
+
+func TestAddClasses(t *testing.T) {
+	sel := Doc2Clone().Find("#main")
+	sel.AddClass("a b")
+
+	// Make sure that class was only added once
+	if !sel.HasClass("a") || !sel.HasClass("b") {
+		t.Errorf("#main does not have classes")
+	}
+}
+
+func TestHasClass(t *testing.T) {
+	sel := Doc().Find("div")
+	if !sel.HasClass("span12") {
+		t.Error("Expected at least one div to have class span12.")
+	}
+}
+
+func TestHasClassNone(t *testing.T) {
+	sel := Doc().Find("h2")
+	if sel.HasClass("toto") {
+		t.Error("Expected h1 to have no class.")
+	}
+}
+
+func TestHasClassNotFirst(t *testing.T) {
+	sel := Doc().Find(".alert")
+	if !sel.HasClass("alert-error") {
+		t.Error("Expected .alert to also have class .alert-error.")
+	}
+}
+
+func TestRemoveClass(t *testing.T) {
+	sel := Doc2Clone().Find("#nf1")
+	sel.RemoveClass("one row")
+
+	if !sel.HasClass("even") || sel.HasClass("one") || sel.HasClass("row") {
+		classes, _ := sel.Attr("class")
+		t.Error("Expected #nf1 to have class even, has ", classes)
+	}
+}
+
+func TestRemoveAllClasses(t *testing.T) {
+	sel := Doc2Clone().Find("#nf1")
+	sel.RemoveClasses()
+
+	if a, ok := sel.Attr("class"); ok {
+		t.Error("All classes were not removed, has ", a)
+	}
+
+	sel = Doc2Clone().Find("#main")
+	sel.RemoveClasses()
+	if a, ok := sel.Attr("class"); ok {
+		t.Error("All classes were not removed, has ", a)
+	}
+}
+
+func TestToggleClass(t *testing.T) {
+	sel := Doc2Clone().Find("#nf1")
+
+	sel.ToggleClass("one")
+	if sel.HasClass("one") {
+		t.Error("Expected #nf1 to not have class one")
+	}
+
+	sel.ToggleClass("one")
+	if !sel.HasClass("one") {
+		t.Error("Expected #nf1 to have class one")
+	}
+
+	sel.ToggleClass("one even row")
+	if a, ok := sel.Attr("class"); ok {
+		t.Error("Expected #nf1 to have no classes, have ", a)
+	}
+}

query.go

@@ -1,15 +1,10 @@
 package goquery
 
 import (
-	"regexp"
-	"strings"
-
 	"code.google.com/p/cascadia"
 	"code.google.com/p/go.net/html"
 )
 
-var rxClassTrim = regexp.MustCompile("[\t\r\n]")
-
 // Is checks the current matched set of elements against a selector and
 // returns true if at least one of these elements matches.
 func (s *Selection) Is(selector string) bool {
@@ -43,24 +38,6 @@ func (s *Selection) IsNodes(nodes ...*html.Node) bool {
 	return s.FilterNodes(nodes...).Length() > 0
 }
 
-// HasClass determines whether any of the matched elements are assigned the
-// given class.
-func (s *Selection) HasClass(class string) bool {
-	class = " " + class + " "
-	for _, n := range s.Nodes {
-		// Applies only to element nodes
-		if n.Type == html.ElementNode {
-			if elClass, ok := getAttributeValue("class", n); ok {
-				elClass = rxClassTrim.ReplaceAllString(" "+elClass+" ", " ")
-				if strings.Index(elClass, class) > -1 {
-					return true
-				}
-			}
-		}
-	}
-	return false
-}
-
 // Contains returns true if the specified Node is within,
 // at any depth, one of the nodes in the Selection object.
 // It is NOT inclusive, to behave like jQuery's implementation, and

query_test.go

@@ -72,27 +72,6 @@ func TestIsNodes(t *testing.T) {
 	}
 }
 
-func TestHasClass(t *testing.T) {
-	sel := Doc().Find("div")
-	if !sel.HasClass("span12") {
-		t.Error("Expected at least one div to have class span12.")
-	}
-}
-
-func TestHasClassNone(t *testing.T) {
-	sel := Doc().Find("h2")
-	if sel.HasClass("toto") {
-		t.Error("Expected h1 to have no class.")
-	}
-}
-
-func TestHasClassNotFirst(t *testing.T) {
-	sel := Doc().Find(".alert")
-	if !sel.HasClass("alert-error") {
-		t.Error("Expected .alert to also have class .alert-error.")
-	}
-}
-
 func TestDocContains(t *testing.T) {
 	sel := Doc().Find("h1")
 	if !Doc().Contains(sel.Nodes[0]) {