From 1a97b2704143a6e8a426d9fd1466744b68a780f2 Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Wed, 26 Feb 2020 10:44:58 +1100
Subject: [PATCH] wip - array length

---
 cmd/constant.go |  2 ++
 cmd/read.go     |  2 ++
 cmd/utils.go    | 40 ++++++++++++++++++++++++++++++++++++----
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/cmd/constant.go b/cmd/constant.go
index 0df9397..0303dde 100644
--- a/cmd/constant.go
+++ b/cmd/constant.go
@@ -7,6 +7,8 @@ import (
 
 var customTag = ""
 var printMode = "v"
+var resultsAsArray = false
+var printLength = false
 var writeInplace = false
 var writeScript = ""
 var sourceYamlFile = ""
diff --git a/cmd/read.go b/cmd/read.go
index bd4f297..c0ed1a9 100644
--- a/cmd/read.go
+++ b/cmd/read.go
@@ -26,6 +26,8 @@ yq r -- things.yaml '--key-starting-with-dashes.blah'
 	cmdRead.PersistentFlags().StringVarP(&docIndex, "doc", "d", "0", "process document index number (0 based, * for all documents)")
 	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(&resultsAsArray, "asArray", "a", false, "print results as array")
+	cmdRead.PersistentFlags().BoolVarP(&printLength, "length", "l", false, "print length of results")
 	cmdRead.PersistentFlags().BoolVarP(&explodeAnchors, "explodeAnchors", "X", false, "explode anchors")
 	return cmdRead
 }
diff --git a/cmd/utils.go b/cmd/utils.go
index 257edd9..c06e738 100644
--- a/cmd/utils.go
+++ b/cmd/utils.go
@@ -77,7 +77,22 @@ func appendDocument(originalMatchingNodes []*yqlib.NodeContext, dataBucket yaml.
 	return append(originalMatchingNodes, matchingNodes...), nil
 }
 
-func printValue(node *yaml.Node, writer io.Writer) error {
+func lengthOf(node *yaml.Node) int {
+	switch node.Kind {
+	case yaml.ScalarNode:
+		return len(node.Value)
+	case yaml.MappingNode:
+		return len(node.Content) / 2
+	default:
+		return len(node.Content)
+	}
+}
+
+func printValue(node *yaml.Node, writer io.Writer, shouldPrintLength bool) error {
+	if shouldPrintLength {
+		return writeString(writer, fmt.Sprintf("%v\n", lengthOf(node)))
+	}
+
 	if node.Kind == yaml.ScalarNode {
 		_, errorWriting := writer.Write([]byte(node.Value + "\n"))
 		return errorWriting
@@ -159,6 +174,9 @@ func printResults(matchingNodes []*yqlib.NodeContext, writer io.Writer) error {
 		}
 		return nil
 	}
+
+	var arrayNode = yaml.Node{Kind: yaml.SequenceNode}
+
 	var errorWriting error
 	for _, mappedDoc := range matchingNodes {
 		switch printMode {
@@ -172,17 +190,31 @@ func printResults(matchingNodes []*yqlib.NodeContext, writer io.Writer) error {
 			var parentNode = yaml.Node{Kind: yaml.MappingNode}
 			parentNode.Content = make([]*yaml.Node, 2)
 			parentNode.Content[0] = &yaml.Node{Kind: yaml.ScalarNode, Value: lib.PathStackToString(mappedDoc.PathStack)}
-			parentNode.Content[1] = mappedDoc.Node
-			if err := printValue(&parentNode, bufferedWriter); err != nil {
+			if printLength {
+				parentNode.Content[1] = &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprintf("%v", lengthOf(mappedDoc.Node))}
+			} else {
+				parentNode.Content[1] = mappedDoc.Node
+			}
+			if resultsAsArray {
+				arrayNode.Content = append(arrayNode.Content, &parentNode)
+			} else if err := printValue(&parentNode, bufferedWriter, false); err != nil {
 				return err
 			}
 		default:
-			if err := printValue(mappedDoc.Node, bufferedWriter); err != nil {
+			if resultsAsArray || printLength {
+				arrayNode.Content = append(arrayNode.Content, mappedDoc.Node)
+			} else if err := printValue(mappedDoc.Node, bufferedWriter, false); err != nil {
 				return err
 			}
 		}
 	}
 
+	if resultsAsArray || (printMode == "v" && printLength) {
+		if err := printValue(&arrayNode, bufferedWriter, printLength); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }