From d7716551cfb0f5de6b09d216274d5d57aee95655 Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Fri, 9 Oct 2020 16:38:07 +1100
Subject: [PATCH] arrays

---
 pkg/yqlib/treeops/data_tree_navigator_test.go | 54 +++++++++++++++++++
 pkg/yqlib/treeops/traverse.go                 | 36 ++++++++++++-
 2 files changed, 88 insertions(+), 2 deletions(-)

diff --git a/pkg/yqlib/treeops/data_tree_navigator_test.go b/pkg/yqlib/treeops/data_tree_navigator_test.go
index c3c3165c..35e04742 100644
--- a/pkg/yqlib/treeops/data_tree_navigator_test.go
+++ b/pkg/yqlib/treeops/data_tree_navigator_test.go
@@ -54,6 +54,60 @@ func TestDataTreeNavigatorSimple(t *testing.T) {
 	test.AssertResult(t, expected, resultsToString(results))
 }
 
+func TestDataTreeNavigatorArraySimple(t *testing.T) {
+
+	nodes := readDoc(t, `- b: apple`)
+
+	path, errPath := treeCreator.ParsePath("[0]")
+	if errPath != nil {
+		t.Error(errPath)
+	}
+	results, errNav := treeNavigator.GetMatchingNodes(nodes, path)
+
+	if errNav != nil {
+		t.Error(errNav)
+	}
+
+	expected := `
+-- Node --
+  Document 0, path: [0]
+  Tag: !!map, Kind: MappingNode, Anchor: 
+  b: apple
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
+func TestDataTreeNavigatorArraySplat(t *testing.T) {
+
+	nodes := readDoc(t, `- b: apple
+- c: banana`)
+
+	path, errPath := treeCreator.ParsePath("[*]")
+	if errPath != nil {
+		t.Error(errPath)
+	}
+	results, errNav := treeNavigator.GetMatchingNodes(nodes, path)
+
+	if errNav != nil {
+		t.Error(errNav)
+	}
+
+	expected := `
+-- Node --
+  Document 0, path: [0]
+  Tag: !!map, Kind: MappingNode, Anchor: 
+  b: apple
+
+-- Node --
+  Document 0, path: [1]
+  Tag: !!map, Kind: MappingNode, Anchor: 
+  c: banana
+`
+
+	test.AssertResult(t, expected, resultsToString(results))
+}
+
 func TestDataTreeNavigatorSimpleDeep(t *testing.T) {
 
 	nodes := readDoc(t, `a: 
diff --git a/pkg/yqlib/treeops/traverse.go b/pkg/yqlib/treeops/traverse.go
index 45200e78..e980e731 100644
--- a/pkg/yqlib/treeops/traverse.go
+++ b/pkg/yqlib/treeops/traverse.go
@@ -47,6 +47,37 @@ func (t *traverser) traverseMap(candidate *CandidateNode, pathNode *PathElement)
 		}
 	}
 	return newMatches, nil
+}
+
+func (t *traverser) traverseArray(candidate *CandidateNode, pathNode *PathElement) ([]*CandidateNode, error) {
+	log.Debug("pathNode Value %v", pathNode.Value)
+	if pathNode.Value == "[*]" {
+
+		var contents = candidate.Node.Content
+		var newMatches = make([]*CandidateNode, len(contents))
+
+		for index := 0; index < len(contents); index = index + 1 {
+			newMatches[index] = &CandidateNode{
+				Document: candidate.Document,
+				Path:     append(candidate.Path, index),
+				Node:     contents[index],
+			}
+		}
+		return newMatches, nil
+
+	}
+
+	index := pathNode.Value.(int64)
+	if int64(len(candidate.Node.Content)) < index {
+		// handle auto append here
+		return make([]*CandidateNode, 0), nil
+	}
+
+	return []*CandidateNode{&CandidateNode{
+		Node:     candidate.Node.Content[index],
+		Document: candidate.Document,
+		Path:     append(candidate.Path, index),
+	}}, nil
 
 }
 
@@ -58,8 +89,9 @@ func (t *traverser) Traverse(matchingNode *CandidateNode, pathNode *PathElement)
 		log.Debug("its a map with %v entries", len(value.Content)/2)
 		return t.traverseMap(matchingNode, pathNode)
 
-	// case yaml.SequenceNode:
-	// 	log.Debug("its a sequence of %v things!", len(value.Content))
+	case yaml.SequenceNode:
+		log.Debug("its a sequence of %v things!", len(value.Content))
+		return t.traverseArray(matchingNode, pathNode)
 
 	// 	switch head := head.(type) {
 	// 	case int64: