From 6ef04e1e777a2e2cabc7f5aa3f6e03a0b1677339 Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Fri, 28 Feb 2020 11:37:27 +1100
Subject: [PATCH] wip

---
 cmd/commands_test.go | 24 ++++++++++++------------
 cmd/utils.go         | 41 ++++++++++++++++++++++++++---------------
 2 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/cmd/commands_test.go b/cmd/commands_test.go
index fd1f5b10..f3d2f671 100644
--- a/cmd/commands_test.go
+++ b/cmd/commands_test.go
@@ -256,7 +256,7 @@ holderB:
 	if result.Error != nil {
 		t.Error(result.Error)
 	}
-	test.AssertResult(t, "holderA: 2\nholderB: 2", result.Output)
+	test.AssertResult(t, "holderA: 2\nholderB: 2\n", result.Output)
 }
 
 func TestReadObjectLengthCmd(t *testing.T) {
@@ -276,8 +276,8 @@ dog: bark
 
 func TestReadObjectLengthDeepCmd(t *testing.T) {
 	content := `holder: 
-	cat: meow
-	dog: bark
+  cat: meow
+  dog: bark
 `
 	filename := test.WriteTempYamlFile(content)
 	defer test.RemoveTempYamlFile(filename)
@@ -292,11 +292,11 @@ func TestReadObjectLengthDeepCmd(t *testing.T) {
 
 func TestReadObjectLengthDeepMultipleCmd(t *testing.T) {
 	content := `holderA: 
-	cat: meow
-	dog: bark
+  cat: meow
+  dog: bark
 holderB: 
-	elephant: meow
-	zebra: bark
+  elephant: meow
+  zebra: bark
 `
 	filename := test.WriteTempYamlFile(content)
 	defer test.RemoveTempYamlFile(filename)
@@ -311,11 +311,11 @@ holderB:
 
 func TestReadObjectLengthDeepMultipleWithPathsCmd(t *testing.T) {
 	content := `holderA: 
-	cat: meow
-	dog: bark
+  cat: meow
+  dog: bark
 holderB: 
-	elephant: meow
-	zebra: bark
+  elephant: meow
+  zebra: bark
 `
 	filename := test.WriteTempYamlFile(content)
 	defer test.RemoveTempYamlFile(filename)
@@ -338,7 +338,7 @@ func TestReadScalarLengthCmd(t *testing.T) {
 	if result.Error != nil {
 		t.Error(result.Error)
 	}
-	test.AssertResult(t, "2\n", result.Output)
+	test.AssertResult(t, "4\n", result.Output)
 }
 
 func TestReadDeepSplatCmd(t *testing.T) {
diff --git a/cmd/utils.go b/cmd/utils.go
index 16acca54..ce7d3ddb 100644
--- a/cmd/utils.go
+++ b/cmd/utils.go
@@ -77,6 +77,30 @@ func appendDocument(originalMatchingNodes []*yqlib.NodeContext, dataBucket yaml.
 	return append(originalMatchingNodes, matchingNodes...), nil
 }
 
+func lengthOf(node *yaml.Node) int {
+	kindToCheck := node.Kind
+	if node.Kind == yaml.DocumentNode && len(node.Content) == 1 {
+		log.Debugf("length of document node, calculating length of child")
+		kindToCheck = node.Content[0].Kind
+	}
+	switch kindToCheck {
+	case yaml.ScalarNode:
+		return len(node.Value)
+	case yaml.MappingNode:
+		return len(node.Content) / 2
+	default:
+		return len(node.Content)
+	}
+}
+
+// transforms node before printing, if required
+func transformNode(node *yaml.Node) *yaml.Node {
+	if printLength {
+		return &yaml.Node{Kind: yaml.ScalarNode, Value: fmt.Sprintf("%v", lengthOf(node))}
+	}
+	return node
+}
+
 func printNode(node *yaml.Node, writer io.Writer) error {
 	var encoder yqlib.Encoder
 	if outputToJSON {
@@ -128,19 +152,6 @@ func explode(matchingNodes []*yqlib.NodeContext) error {
 	return nil
 }
 
-func convertToLength(node *yaml.Node) *yaml.Node {
-
-		switch kindToCheck {
-	case yaml.ScalarNode:
-		return len(node.Value)
-	case yaml.MappingNode:
-		return len(node.Content) / 2
-	default:
-		return len(node.Content)
-	}
-
-}
-
 func printResults(matchingNodes []*yqlib.NodeContext, writer io.Writer) error {
 	if prettyPrint {
 		setStyle(matchingNodes, 0)
@@ -177,12 +188,12 @@ 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
+			parentNode.Content[1] = transformNode(mappedDoc.Node)
 			if err := printNode(&parentNode, bufferedWriter); err != nil {
 				return err
 			}
 		default:
-			if err := printNode(mappedDoc.Node, bufferedWriter); err != nil {
+			if err := printNode(transformNode(mappedDoc.Node), bufferedWriter); err != nil {
 				return err
 			}
 		}