added single count operator

pkg/yqlib/treeops/data_tree_navigator_test.go

@@ -240,6 +240,197 @@ func TestDataTreeNavigatorDeleteViaSelf(t *testing.T) {
 	test.AssertResult(t, expected, resultsToString(results))
 }
 
+func TestDataTreeNavigatorCountWithFilter(t *testing.T) {
+
+	nodes := readDoc(t, `f:
+  a: frog
+  b: dally
+  c: log`)
+
+	path, errPath := treeCreator.ParsePath("f(count(. == *og))")
+	if errPath != nil {
+		t.Error(errPath)
+	}
+	results, errNav := treeNavigator.GetMatchingNodes(nodes, path)
+
+	if errNav != nil {
+		t.Error(errNav)
+	}
+
+	expected := `
+-- Node --
+  Document 0, path: [f]
+  Tag: !!int, Kind: ScalarNode, Anchor: 
+  2
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
+func TestDataTreeNavigatorCountWithFilter2(t *testing.T) {
+
+	nodes := readDoc(t, `f:
+  a: frog
+  b: dally
+  c: log`)
+
+	path, errPath := treeCreator.ParsePath("count(f(. == *og))")
+	if errPath != nil {
+		t.Error(errPath)
+	}
+	results, errNav := treeNavigator.GetMatchingNodes(nodes, path)
+
+	if errNav != nil {
+		t.Error(errNav)
+	}
+
+	expected := `
+-- Node --
+  Document 0, path: []
+  Tag: !!int, Kind: ScalarNode, Anchor: 
+  2
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
+func TestDataTreeNavigatorCountMultipleMatchesInside(t *testing.T) {
+
+	nodes := readDoc(t, `f:
+  a: [1,2]
+  b: dally
+  c: [3,4,5]`)
+
+	path, errPath := treeCreator.ParsePath("f(count(a or c))")
+	if errPath != nil {
+		t.Error(errPath)
+	}
+	results, errNav := treeNavigator.GetMatchingNodes(nodes, path)
+
+	if errNav != nil {
+		t.Error(errNav)
+	}
+
+	expected := `
+-- Node --
+  Document 0, path: [f]
+  Tag: !!int, Kind: ScalarNode, Anchor: 
+  2
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
+func TestDataTreeNavigatorCountMultipleMatchesInsideSplat(t *testing.T) {
+
+	nodes := readDoc(t, `f:
+  a: [1,2,3]
+  b: [1,2,3,4]
+  c: [1,2,3,4,5]`)
+
+	path, errPath := treeCreator.ParsePath("f(count( (a or c)*))")
+	if errPath != nil {
+		t.Error(errPath)
+	}
+	results, errNav := treeNavigator.GetMatchingNodes(nodes, path)
+
+	if errNav != nil {
+		t.Error(errNav)
+	}
+
+	expected := `
+-- Node --
+  Document 0, path: [f]
+  Tag: !!int, Kind: ScalarNode, Anchor: 
+  8
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
+func TestDataTreeNavigatorCountMultipleMatchesOutside(t *testing.T) {
+
+	nodes := readDoc(t, `f:
+  a: [1,2,3]
+  b: [1,2,3,4]
+  c: [1,2,3,4,5]`)
+
+	path, errPath := treeCreator.ParsePath("f(a or c)(count(*))")
+	if errPath != nil {
+		t.Error(errPath)
+	}
+	results, errNav := treeNavigator.GetMatchingNodes(nodes, path)
+
+	if errNav != nil {
+		t.Error(errNav)
+	}
+
+	expected := `
+-- Node --
+  Document 0, path: [f a]
+  Tag: !!int, Kind: ScalarNode, Anchor: 
+  3
+-- Node --
+  Document 0, path: [f c]
+  Tag: !!int, Kind: ScalarNode, Anchor: 
+  5
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
+func TestDataTreeNavigatorCountOfResults(t *testing.T) {
+
+	nodes := readDoc(t, `- apple
+- sdfsd
+- apple`)
+
+	path, errPath := treeCreator.ParsePath("count(*)")
+	if errPath != nil {
+		t.Error(errPath)
+	}
+	results, errNav := treeNavigator.GetMatchingNodes(nodes, path)
+
+	if errNav != nil {
+		t.Error(errNav)
+	}
+
+	expected := `
+-- Node --
+  Document 0, path: []
+  Tag: !!int, Kind: ScalarNode, Anchor: 
+  3
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
+func TestDataTreeNavigatorCountNoMatches(t *testing.T) {
+
+	nodes := readDoc(t, `- apple
+- sdfsd
+- apple`)
+
+	path, errPath := treeCreator.ParsePath("count(5)")
+	if errPath != nil {
+		t.Error(errPath)
+	}
+	results, errNav := treeNavigator.GetMatchingNodes(nodes, path)
+
+	if errNav != nil {
+		t.Error(errNav)
+	}
+
+	expected := `
+-- Node --
+  Document 0, path: []
+  Tag: !!int, Kind: ScalarNode, Anchor: 
+  0
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
 func TestDataTreeNavigatorDeleteAndWrite(t *testing.T) {
 
 	nodes := readDoc(t, `a: 

pkg/yqlib/treeops/lib.go

@@ -47,7 +47,10 @@ var Equals = &OperationType{Type: "EQUALS", NumArgs: 2, Precedence: 30, Handler:
 var Assign = &OperationType{Type: "ASSIGN", NumArgs: 2, Precedence: 35, Handler: AssignOperator}
 var DeleteChild = &OperationType{Type: "DELETE", NumArgs: 2, Precedence: 30, Handler: DeleteChildOperator}
 
-// var Length = &OperationType{Type: "Length", NumArgs: 2, Precedence: 35}
+var Count = &OperationType{Type: "COUNT", NumArgs: 1, Precedence: 35, Handler: CountOperator}
+
+// var Exists = &OperationType{Type: "Length", NumArgs: 2, Precedence: 35}
+// filters matches if they have the existing path
 
 type PathElement struct {
 	PathElementType PathElementType

pkg/yqlib/treeops/operators.go

@@ -1,6 +1,8 @@
 package treeops
 
 import (
+	"fmt"
+
 	"github.com/elliotchance/orderedmap"
 	"gopkg.in/yaml.v3"
 )
@@ -93,6 +95,30 @@ func EqualsOperator(d *dataTreeNavigator, matchMap *orderedmap.OrderedMap, pathN
 	return results, nil
 }
 
+func CountOperator(d *dataTreeNavigator, matchMap *orderedmap.OrderedMap, pathNode *PathTreeNode) (*orderedmap.OrderedMap, error) {
+	log.Debugf("-- countOperation")
+	var results = orderedmap.NewOrderedMap()
+
+	for el := matchMap.Front(); el != nil; el = el.Next() {
+		candidate := el.Value.(*CandidateNode)
+		elMap := orderedmap.NewOrderedMap()
+		elMap.Set(el.Key, el.Value)
+		childMatches, errChild := d.getMatchingNodes(elMap, pathNode.Rhs)
+
+		if errChild != nil {
+			return nil, errChild
+		}
+
+		length := childMatches.Len()
+		node := &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprintf("%v", length), Tag: "!!int"}
+		lengthCand := &CandidateNode{Node: node, Document: candidate.Document, Path: candidate.Path}
+		results.Set(candidate.getKey(), lengthCand)
+
+	}
+
+	return results, nil
+}
+
 func findMatchingChildren(d *dataTreeNavigator, results *orderedmap.OrderedMap, candidate *CandidateNode, lhs *PathTreeNode, valuePattern string) error {
 	var children *orderedmap.OrderedMap
 	var err error

pkg/yqlib/treeops/path_postfix_test.go

@@ -47,6 +47,22 @@ Operation - TRAVERSE
 	test.AssertResultComplex(t, expectedOutput, actual)
 }
 
+func TestPostFixLength(t *testing.T) {
+	var infix = "len(a)"
+	var expectedOutput = `PathKey - 'a'
+--------
+Operation - Length
+--------
+`
+
+	actual, err := testExpression(infix)
+	if err != nil {
+		t.Error(err)
+	}
+
+	test.AssertResultComplex(t, expectedOutput, actual)
+}
+
 func TestPostFixSimpleExample(t *testing.T) {
 	var infix = "a"
 	var expectedOutput = `PathKey - 'a'

pkg/yqlib/treeops/path_tokeniser.go

@@ -36,7 +36,7 @@ func pathToken(wrapped bool) lex.Action {
 func opToken(op *OperationType, againstSelf bool) lex.Action {
 	return func(s *lex.Scanner, m *machines.Match) (interface{}, error) {
 		value := string(m.Bytes)
-		return &Token{PathElementType: Operation, OperationType: op, Value: value, StringValue: value, PrefixSelf: againstSelf}, nil
+		return &Token{PathElementType: Operation, OperationType: op, Value: op.Type, StringValue: value, PrefixSelf: againstSelf}, nil
 	}
 }
 
@@ -77,6 +77,7 @@ func initLexer() (*lex.Lexer, error) {
 
 	lexer.Add([]byte(`([Oo][Rr])`), opToken(Or, false))
 	lexer.Add([]byte(`([Aa][Nn][Dd])`), opToken(And, false))
+	lexer.Add([]byte(`([Cc][Oo][Uu][Nn][Tt])`), opToken(Count, false))
 
 	lexer.Add([]byte(`\.\s*==\s*`), opToken(Equals, true))
 	lexer.Add([]byte(`\s*==\s*`), opToken(Equals, false))

pkg/yqlib/treeops/path_tokeniser_test.go

@@ -10,48 +10,50 @@ var tokeniserTests = []struct {
 	path           string
 	expectedTokens []interface{}
 }{ // TODO: Ensure ALL documented examples have tests! sheesh
-
+	{"len(.)", append(make([]interface{}, 0), "LENGTH", "(", "SELF", ")")},
-	{"a OR (b OR c)", append(make([]interface{}, 0), "a", "OR", "(", "b", "OR", "c", ")")},
+	{"\"len\"(.)", append(make([]interface{}, 0), "len", ".", "(", "SELF", ")")},
-	{"a .- (b OR c)", append(make([]interface{}, 0), "a", " .- ", "(", "b", "OR", "c", ")")},
+	// {"a OR (b OR c)", append(make([]interface{}, 0), "a", "OR", "(", "b", "OR", "c", ")")},
-	{"(animal==3)", append(make([]interface{}, 0), "(", "animal", "==", int64(3), ")")},
+	// {"a OR (b OR c)", append(make([]interface{}, 0), "a", "OR", "(", "b", "OR", "c", ")")},
-	{"(animal==f3)", append(make([]interface{}, 0), "(", "animal", "==", "f3", ")")},
+	// {"a .- (b OR c)", append(make([]interface{}, 0), "a", " .- ", "(", "b", "OR", "c", ")")},
-	{"apples.BANANAS", append(make([]interface{}, 0), "apples", ".", "BANANAS")},
+	// {"(animal==3)", append(make([]interface{}, 0), "(", "animal", "==", int64(3), ")")},
-	{"appl*.BANA*", append(make([]interface{}, 0), "appl*", ".", "BANA*")},
+	// {"(animal==f3)", append(make([]interface{}, 0), "(", "animal", "==", "f3", ")")},
-	{"a.b.**", append(make([]interface{}, 0), "a", ".", "b", ".", "**")},
+	// {"apples.BANANAS", append(make([]interface{}, 0), "apples", ".", "BANANAS")},
-	{"a.\"=\".frog", append(make([]interface{}, 0), "a", ".", "=", ".", "frog")},
+	// {"appl*.BANA*", append(make([]interface{}, 0), "appl*", ".", "BANA*")},
-	{"a.b.*", append(make([]interface{}, 0), "a", ".", "b", ".", "*")},
+	// {"a.b.**", append(make([]interface{}, 0), "a", ".", "b", ".", "**")},
-	{"a.b.thin*", append(make([]interface{}, 0), "a", ".", "b", ".", "thin*")},
+	// {"a.\"=\".frog", append(make([]interface{}, 0), "a", ".", "=", ".", "frog")},
-	{"a.b[0]", append(make([]interface{}, 0), "a", ".", "b", ".", int64(0))},
+	// {"a.b.*", append(make([]interface{}, 0), "a", ".", "b", ".", "*")},
-	{"a.b.[0]", append(make([]interface{}, 0), "a", ".", "b", ".", int64(0))},
+	// {"a.b.thin*", append(make([]interface{}, 0), "a", ".", "b", ".", "thin*")},
-	{"a.b[*]", append(make([]interface{}, 0), "a", ".", "b", ".", "[*]")},
+	// {"a.b[0]", append(make([]interface{}, 0), "a", ".", "b", ".", int64(0))},
-	{"a.b.[*]", append(make([]interface{}, 0), "a", ".", "b", ".", "[*]")},
+	// {"a.b.[0]", append(make([]interface{}, 0), "a", ".", "b", ".", int64(0))},
-	{"a.b[+]", append(make([]interface{}, 0), "a", ".", "b", ".", "[+]")},
+	// {"a.b[*]", append(make([]interface{}, 0), "a", ".", "b", ".", "[*]")},
-	{"a.b.[+]", append(make([]interface{}, 0), "a", ".", "b", ".", "[+]")},
+	// {"a.b.[*]", append(make([]interface{}, 0), "a", ".", "b", ".", "[*]")},
-	{"a.b[-12]", append(make([]interface{}, 0), "a", ".", "b", ".", int64(-12))},
+	// {"a.b[+]", append(make([]interface{}, 0), "a", ".", "b", ".", "[+]")},
-	{"a.b.0", append(make([]interface{}, 0), "a", ".", "b", ".", int64(0))},
+	// {"a.b.[+]", append(make([]interface{}, 0), "a", ".", "b", ".", "[+]")},
-	// {"a.b.-12", append(make([]interface{}, 0), "a", ".", "b", ".", int64(-12))},
+	// {"a.b[-12]", append(make([]interface{}, 0), "a", ".", "b", ".", int64(-12))},
-	{"a", append(make([]interface{}, 0), "a")},
+	// {"a.b.0", append(make([]interface{}, 0), "a", ".", "b", ".", int64(0))},
-	{"\"a.b\".c", append(make([]interface{}, 0), "a.b", ".", "c")},
+	// // {"a.b.-12", append(make([]interface{}, 0), "a", ".", "b", ".", int64(-12))},
-	{`b."foo.bar"`, append(make([]interface{}, 0), "b", ".", "foo.bar")},
+	// {"a", append(make([]interface{}, 0), "a")},
-	{"animals(.==cat)", append(make([]interface{}, 0), "animals", ".", "(", "SELF", ".==", "cat", ")")},
+	// {"\"a.b\".c", append(make([]interface{}, 0), "a.b", ".", "c")},
-	{"animals.(.==cat)", append(make([]interface{}, 0), "animals", ".", "(", "SELF", ".==", "cat", ")")},
+	// {`b."foo.bar"`, append(make([]interface{}, 0), "b", ".", "foo.bar")},
-	{"animals(. == cat)", append(make([]interface{}, 0), "animals", ".", "(", "SELF", ". == ", "cat", ")")},
+	// {"animals(.==cat)", append(make([]interface{}, 0), "animals", ".", "(", "SELF", ".==", "cat", ")")},
-	{"animals(.==c*)", append(make([]interface{}, 0), "animals", ".", "(", "SELF", ".==", "c*", ")")},
+	// {"animals.(.==cat)", append(make([]interface{}, 0), "animals", ".", "(", "SELF", ".==", "cat", ")")},
-	{"animals(a.b==c*)", append(make([]interface{}, 0), "animals", ".", "(", "a", ".", "b", "==", "c*", ")")},
+	// {"animals(. == cat)", append(make([]interface{}, 0), "animals", ".", "(", "SELF", ". == ", "cat", ")")},
-	{"animals.(a.b==c*)", append(make([]interface{}, 0), "animals", ".", "(", "a", ".", "b", "==", "c*", ")")},
+	// {"animals(.==c*)", append(make([]interface{}, 0), "animals", ".", "(", "SELF", ".==", "c*", ")")},
-	{"(a.b==c*).animals", append(make([]interface{}, 0), "(", "a", ".", "b", "==", "c*", ")", ".", "animals")},
+	// {"animals(a.b==c*)", append(make([]interface{}, 0), "animals", ".", "(", "a", ".", "b", "==", "c*", ")")},
-	{"(a.b==c*)animals", append(make([]interface{}, 0), "(", "a", ".", "b", "==", "c*", ")", ".", "animals")},
+	// {"animals.(a.b==c*)", append(make([]interface{}, 0), "animals", ".", "(", "a", ".", "b", "==", "c*", ")")},
-	{"[1].a.d", append(make([]interface{}, 0), int64(1), ".", "a", ".", "d")},
+	// {"(a.b==c*).animals", append(make([]interface{}, 0), "(", "a", ".", "b", "==", "c*", ")", ".", "animals")},
-	{"[1]a.d", append(make([]interface{}, 0), int64(1), ".", "a", ".", "d")},
+	// {"(a.b==c*)animals", append(make([]interface{}, 0), "(", "a", ".", "b", "==", "c*", ")", ".", "animals")},
-	{"a[0]c", append(make([]interface{}, 0), "a", ".", int64(0), ".", "c")},
+	// {"[1].a.d", append(make([]interface{}, 0), int64(1), ".", "a", ".", "d")},
-	{"a.[0].c", append(make([]interface{}, 0), "a", ".", int64(0), ".", "c")},
+	// {"[1]a.d", append(make([]interface{}, 0), int64(1), ".", "a", ".", "d")},
-	{"[0]", append(make([]interface{}, 0), int64(0))},
+	// {"a[0]c", append(make([]interface{}, 0), "a", ".", int64(0), ".", "c")},
-	{"0", append(make([]interface{}, 0), int64(0))},
+	// {"a.[0].c", append(make([]interface{}, 0), "a", ".", int64(0), ".", "c")},
-	{"a.b[+]c", append(make([]interface{}, 0), "a", ".", "b", ".", "[+]", ".", "c")},
+	// {"[0]", append(make([]interface{}, 0), int64(0))},
-	{"a.cool(s.d.f == cool)", append(make([]interface{}, 0), "a", ".", "cool", ".", "(", "s", ".", "d", ".", "f", " == ", "cool", ")")},
+	// {"0", append(make([]interface{}, 0), int64(0))},
-	{"a.cool.(s.d.f==cool OR t.b.h==frog).caterpillar", append(make([]interface{}, 0), "a", ".", "cool", ".", "(", "s", ".", "d", ".", "f", "==", "cool", "OR", "t", ".", "b", ".", "h", "==", "frog", ")", ".", "caterpillar")},
+	// {"a.b[+]c", append(make([]interface{}, 0), "a", ".", "b", ".", "[+]", ".", "c")},
-	{"a.cool(s.d.f==cool and t.b.h==frog)*", append(make([]interface{}, 0), "a", ".", "cool", ".", "(", "s", ".", "d", ".", "f", "==", "cool", "and", "t", ".", "b", ".", "h", "==", "frog", ")", ".", "*")},
+	// {"a.cool(s.d.f == cool)", append(make([]interface{}, 0), "a", ".", "cool", ".", "(", "s", ".", "d", ".", "f", " == ", "cool", ")")},
-	{"a.cool(s.d.f==cool and t.b.h==frog).th*", append(make([]interface{}, 0), "a", ".", "cool", ".", "(", "s", ".", "d", ".", "f", "==", "cool", "and", "t", ".", "b", ".", "h", "==", "frog", ")", ".", "th*")},
+	// {"a.cool.(s.d.f==cool OR t.b.h==frog).caterpillar", append(make([]interface{}, 0), "a", ".", "cool", ".", "(", "s", ".", "d", ".", "f", "==", "cool", "OR", "t", ".", "b", ".", "h", "==", "frog", ")", ".", "caterpillar")},
+	// {"a.cool(s.d.f==cool and t.b.h==frog)*", append(make([]interface{}, 0), "a", ".", "cool", ".", "(", "s", ".", "d", ".", "f", "==", "cool", "and", "t", ".", "b", ".", "h", "==", "frog", ")", ".", "*")},
+	// {"a.cool(s.d.f==cool and t.b.h==frog).th*", append(make([]interface{}, 0), "a", ".", "cool", ".", "(", "s", ".", "d", ".", "f", "==", "cool", "and", "t", ".", "b", ".", "h", "==", "frog", ")", ".", "th*")},
 }
 
 var tokeniser = NewPathTokeniser()

pkg/yqlib/treeops/path_tree.go

@@ -42,10 +42,20 @@ func (p *pathTreeCreator) CreatePathTree(postFixPath []*PathElement) (*PathTreeN
 	for _, pathElement := range postFixPath {
 		var newNode = PathTreeNode{PathElement: pathElement}
 		if pathElement.PathElementType == Operation {
-			remaining, lhs, rhs := stack[:len(stack)-2], stack[len(stack)-2], stack[len(stack)-1]
+			numArgs := pathElement.OperationType.NumArgs
-			newNode.Lhs = lhs
+			if numArgs == 0 {
-			newNode.Rhs = rhs
+				remaining := stack[:len(stack)-1]
-			stack = remaining
+				stack = remaining
+			} else if numArgs == 1 {
+				remaining, rhs := stack[:len(stack)-1], stack[len(stack)-1]
+				newNode.Rhs = rhs
+				stack = remaining
+			} else {
+				remaining, lhs, rhs := stack[:len(stack)-2], stack[len(stack)-2], stack[len(stack)-1]
+				newNode.Lhs = lhs
+				newNode.Rhs = rhs
+				stack = remaining
+			}
 		}
 		stack = append(stack, &newNode)
 	}