From 8e6ceba2ac5ef5d29730ffd192fe28b1485765ae Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Fri, 28 Feb 2020 12:18:27 +1100
Subject: [PATCH] Array length and collect

---
 cmd/commands_test.go | 53 ++++++++++++++++++++++++++++++++++++++++----
 cmd/constant.go      |  1 +
 cmd/read.go          |  1 +
 cmd/utils.go         | 17 ++++++++++++--
 4 files changed, 66 insertions(+), 6 deletions(-)

diff --git a/cmd/commands_test.go b/cmd/commands_test.go
index f3d2f671..ce1f2bbe 100644
--- a/cmd/commands_test.go
+++ b/cmd/commands_test.go
@@ -225,6 +225,8 @@ func TestReadArrayLengthDeepMultipleCmd(t *testing.T) {
 	content := `holderA: 
 - things
 - whatever
+skipMe:
+- yep
 holderB: 
 - other things
 - cool
@@ -233,11 +235,54 @@ holderB:
 	defer test.RemoveTempYamlFile(filename)
 
 	cmd := getRootCommand()
-	result := test.RunCmd(cmd, fmt.Sprintf("read -l %s holder*", filename))
+	result := test.RunCmd(cmd, fmt.Sprintf("read -l -c %s holder*", filename))
 	if result.Error != nil {
 		t.Error(result.Error)
 	}
-	test.AssertResult(t, "4\n", result.Output)
+	test.AssertResult(t, "2\n", result.Output)
+}
+
+func TestReadCollectCmd(t *testing.T) {
+	content := `holderA: yep
+skipMe: not me
+holderB: me too
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -c %s holder*", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	expectedOutput := `- yep
+- me too
+`
+	test.AssertResult(t, expectedOutput, result.Output)
+}
+
+func TestReadCollectArrayCmd(t *testing.T) {
+	content := `- name: fred
+  value: 32
+- name: sam
+  value: 67
+- name: fernie
+  value: 103
+`
+	filename := test.WriteTempYamlFile(content)
+	defer test.RemoveTempYamlFile(filename)
+
+	cmd := getRootCommand()
+	result := test.RunCmd(cmd, fmt.Sprintf("read -c %s (name==f*)", filename))
+	if result.Error != nil {
+		t.Error(result.Error)
+	}
+	expectedOutput := `- name: fred
+  value: 32
+- name: fernie
+  value: 103
+`
+	test.AssertResult(t, expectedOutput, result.Output)
 }
 
 func TestReadArrayLengthDeepMultipleWithPathCmd(t *testing.T) {
@@ -302,11 +347,11 @@ holderB:
 	defer test.RemoveTempYamlFile(filename)
 
 	cmd := getRootCommand()
-	result := test.RunCmd(cmd, fmt.Sprintf("read -l %s holder*", filename))
+	result := test.RunCmd(cmd, fmt.Sprintf("read -l -c %s holder*", filename))
 	if result.Error != nil {
 		t.Error(result.Error)
 	}
-	test.AssertResult(t, "4\n", result.Output)
+	test.AssertResult(t, "2\n", result.Output)
 }
 
 func TestReadObjectLengthDeepMultipleWithPathsCmd(t *testing.T) {
diff --git a/cmd/constant.go b/cmd/constant.go
index b54500bb..69e3a8ea 100644
--- a/cmd/constant.go
+++ b/cmd/constant.go
@@ -8,6 +8,7 @@ import (
 var customTag = ""
 var printMode = "v"
 var printLength = false
+var collectIntoArray = false
 var writeInplace = false
 var writeScript = ""
 var sourceYamlFile = ""
diff --git a/cmd/read.go b/cmd/read.go
index d94b85e0..35f62afb 100644
--- a/cmd/read.go
+++ b/cmd/read.go
@@ -27,6 +27,7 @@ yq r -- things.yaml '--key-starting-with-dashes.blah'
 	cmdRead.PersistentFlags().StringVarP(&printMode, "printMode", "p", "v", "print mode (v (values, default), p (paths), pv (path and value pairs)")
 	cmdRead.PersistentFlags().StringVarP(&defaultValue, "defaultValue", "D", "", "default value printed when there are no results")
 	cmdRead.PersistentFlags().BoolVarP(&printLength, "length", "l", false, "print length of results")
+	cmdRead.PersistentFlags().BoolVarP(&collectIntoArray, "collect", "c", false, "collect results into array")
 	cmdRead.PersistentFlags().BoolVarP(&explodeAnchors, "explodeAnchors", "X", false, "explode anchors")
 	return cmdRead
 }
diff --git a/cmd/utils.go b/cmd/utils.go
index ce7d3ddb..bd5a88c7 100644
--- a/cmd/utils.go
+++ b/cmd/utils.go
@@ -176,6 +176,9 @@ func printResults(matchingNodes []*yqlib.NodeContext, writer io.Writer) error {
 		return nil
 	}
 	var errorWriting error
+
+	var arrayCollection = yaml.Node{Kind: yaml.SequenceNode}
+
 	for _, mappedDoc := range matchingNodes {
 		switch printMode {
 		case "p":
@@ -189,16 +192,26 @@ func printResults(matchingNodes []*yqlib.NodeContext, writer io.Writer) error {
 			parentNode.Content = make([]*yaml.Node, 2)
 			parentNode.Content[0] = &yaml.Node{Kind: yaml.ScalarNode, Value: lib.PathStackToString(mappedDoc.PathStack)}
 			parentNode.Content[1] = transformNode(mappedDoc.Node)
-			if err := printNode(&parentNode, bufferedWriter); err != nil {
+			if collectIntoArray {
+				arrayCollection.Content = append(arrayCollection.Content, &parentNode)
+			} else if err := printNode(&parentNode, bufferedWriter); err != nil {
 				return err
 			}
 		default:
-			if err := printNode(transformNode(mappedDoc.Node), bufferedWriter); err != nil {
+			if collectIntoArray {
+				arrayCollection.Content = append(arrayCollection.Content, mappedDoc.Node)
+			} else if err := printNode(transformNode(mappedDoc.Node), bufferedWriter); err != nil {
 				return err
 			}
 		}
 	}
 
+	if collectIntoArray {
+		if err := printNode(transformNode(&arrayCollection), bufferedWriter); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }