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 } }