collect

pkg/yqlib/treeops/data_tree_navigator_test.go

@@ -267,6 +267,34 @@ func TestDataTreeNavigatorCountWithFilter(t *testing.T) {
 	test.AssertResult(t, expected, resultsToString(results))
 }
 
+func TestDataTreeNavigatorCollectWithFilter(t *testing.T) {
+
+	nodes := readDoc(t, `f:
+  a: frog
+  b: dally
+  c: log`)
+
+	path, errPath := treeCreator.ParsePath("f(collect(. == *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: , Kind: SequenceNode, Anchor: 
+  - frog
+- log
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
 func TestDataTreeNavigatorCountWithFilter2(t *testing.T) {
 
 	nodes := readDoc(t, `f:
@@ -294,6 +322,34 @@ func TestDataTreeNavigatorCountWithFilter2(t *testing.T) {
 	test.AssertResult(t, expected, resultsToString(results))
 }
 
+func TestDataTreeNavigatorCollectWithFilter2(t *testing.T) {
+
+	nodes := readDoc(t, `f:
+  a: frog
+  b: dally
+  c: log`)
+
+	path, errPath := treeCreator.ParsePath("collect(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: , Kind: SequenceNode, Anchor: 
+  - frog
+- log
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
 func TestDataTreeNavigatorCountMultipleMatchesInside(t *testing.T) {
 
 	nodes := readDoc(t, `f:
@@ -321,6 +377,34 @@ func TestDataTreeNavigatorCountMultipleMatchesInside(t *testing.T) {
 	test.AssertResult(t, expected, resultsToString(results))
 }
 
+func TestDataTreeNavigatorCollectMultipleMatchesInside(t *testing.T) {
+
+	nodes := readDoc(t, `f:
+  a: [1,2]
+  b: dally
+  c: [3,4,5]`)
+
+	path, errPath := treeCreator.ParsePath("f(collect(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: , Kind: SequenceNode, Anchor: 
+  - [1, 2]
+- [3, 4, 5]
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
 func TestDataTreeNavigatorCountMultipleMatchesInsideSplat(t *testing.T) {
 
 	nodes := readDoc(t, `f:

pkg/yqlib/treeops/lib.go

@@ -38,6 +38,7 @@ var Assign = &OperationType{Type: "ASSIGN", NumArgs: 2, Precedence: 35, Handler:
 var DeleteChild = &OperationType{Type: "DELETE", NumArgs: 2, Precedence: 30, Handler: DeleteChildOperator}
 
 var Count = &OperationType{Type: "COUNT", NumArgs: 1, Precedence: 35, Handler: CountOperator}
+var Collect = &OperationType{Type: "COLLECT", NumArgs: 1, Precedence: 35, Handler: CollectOperator}
 
 // var Exists = &OperationType{Type: "Length", NumArgs: 2, Precedence: 35}
 // filters matches if they have the existing path

pkg/yqlib/treeops/operators.go

@@ -119,6 +119,37 @@ func CountOperator(d *dataTreeNavigator, matchMap *orderedmap.OrderedMap, pathNo
 	return results, nil
 }
 
+func CollectOperator(d *dataTreeNavigator, matchMap *orderedmap.OrderedMap, pathNode *PathTreeNode) (*orderedmap.OrderedMap, error) {
+	log.Debugf("-- collectOperation")
+
+	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
+		}
+
+		node := &yaml.Node{Kind: yaml.SequenceNode}
+
+		for childEl := childMatches.Front(); childEl != nil; childEl = childEl.Next() {
+			childCandidate := childEl.Value.(*CandidateNode)
+			node.Content = append(node.Content, childCandidate.Node)
+		}
+
+		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_tokeniser.go

@@ -78,6 +78,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(`([Cc][Oo][Ll][Ll][Ee][Cc][Tt])`), opToken(Collect, false))
 
 	lexer.Add([]byte(`\.\s*==\s*`), opToken(Equals, true))
 	lexer.Add([]byte(`\s*==\s*`), opToken(Equals, false))