Merge pull request #3106 from Shell32-Natsu/array-index

support array index in PathGetter

kyaml/yaml/fns.go

@@ -6,6 +6,7 @@ package yaml
 import (
 	"fmt"
 	"regexp"
+	"strconv"
 	"strings"
 
 	"github.com/davecgh/go-spew/spew"
@@ -123,6 +124,34 @@ func (e ElementSetter) Filter(rn *RNode) (*RNode, error) {
 	return NewRNode(e.Element), nil
 }
 
+// GetElementByIndex will return a Filter which can be applied to a sequence
+// node to get the element specified by the index
+func GetElementByIndex(index int) ElementIndexer {
+	return ElementIndexer{Index: index}
+}
+
+// ElementIndexer picks the element with a specified index. Index starts from
+// 0 to len(list) - 1. a hyphen ("-") means the last index.
+type ElementIndexer struct {
+	Index int
+}
+
+// Filter implements Filter
+func (i ElementIndexer) Filter(rn *RNode) (*RNode, error) {
+	// rn.Elements will return error if rn is not a sequence node.
+	elems, err := rn.Elements()
+	if err != nil {
+		return nil, err
+	}
+	if i.Index < 0 {
+		return elems[len(elems)-1], nil
+	}
+	if i.Index >= len(elems) {
+		return nil, nil
+	}
+	return elems[i.Index], nil
+}
+
 // Clear returns a FieldClearer
 func Clear(name string) FieldClearer {
 	return FieldClearer{Name: name}
@@ -364,12 +393,15 @@ type PathGetter struct {
 	// Each path part may be one of:
 	// * FieldMatcher -- e.g. "spec"
 	// * Map Key -- e.g. "app.k8s.io/version"
-	// * List Entry -- e.g. "[name=nginx]" or "[=-jar]"
+	// * List Entry -- e.g. "[name=nginx]" or "[=-jar]" or "0" or "-"
 	//
 	// Map Keys and Fields are equivalent.
 	// See FieldMatcher for more on Fields and Map Keys.
 	//
-	// List Entries are specified as map entry to match [fieldName=fieldValue].
+	// List Entries can be specified as map entry to match [fieldName=fieldValue]
+	// or a postional index like 0 to get the element. - (unquoted hyphen) is
+	// special and means the last element.
+	//
 	// See Elem for more on List Entries.
 	//
 	// Examples:
@@ -382,6 +414,8 @@ type PathGetter struct {
 	// * The leaf Node (final path) will be created with a Kind matching Create
 	// * Intermediary Nodes will be created as either a MappingNodes or
 	//   SequenceNodes as appropriate for each's Path location.
+	// * If a list item is specified by a index (an offset or "-"), this item will
+	//   not be created even Create is set.
 	Create yaml.Kind `yaml:"create,omitempty"`
 
 	// Style is the style to apply to created value Nodes.
@@ -402,12 +436,12 @@ func (l PathGetter) Filter(rn *RNode) (*RNode, error) {
 		if len(l.Path) > i+1 {
 			nextPart = l.Path[i+1]
 		}
-		if IsListIndex(part) {
-			match, err = l.doElem(match, part)
-		} else {
-			fieldPath = append(fieldPath, part)
-			match, err = l.doField(match, part, l.getKind(nextPart))
+		var fltr Filter
+		fltr, err = l.getFilter(part, nextPart, &fieldPath)
+		if err != nil {
+			return nil, err
 		}
+		match, err = match.Pipe(fltr)
 		if IsMissingOrError(match, err) {
 			return nil, err
 		}
@@ -416,14 +450,36 @@ func (l PathGetter) Filter(rn *RNode) (*RNode, error) {
 	return match, nil
 }
 
-func (l PathGetter) doElem(rn *RNode, part string) (*RNode, error) {
+func (l PathGetter) getFilter(part, nextPart string, fieldPath *[]string) (Filter, error) {
+	idx, err := strconv.Atoi(part)
+	switch {
+	case err == nil:
+		// part is a number
+		if idx < 0 {
+			return nil, fmt.Errorf("array index %d cannot be negative", idx)
+		}
+		return GetElementByIndex(idx), nil
+	case part == "-":
+		// part is a hyphen
+		return GetElementByIndex(-1), nil
+	case IsListIndex(part):
+		// part is surrounded by brackets
+		return l.elemFilter(part)
+	default:
+		// mapping node
+		*fieldPath = append(*fieldPath, part)
+		return l.fieldFilter(part, l.getKind(nextPart))
+	}
+}
+
+func (l PathGetter) elemFilter(part string) (Filter, error) {
 	var match *RNode
 	name, value, err := SplitIndexNameValue(part)
 	if err != nil {
 		return nil, errors.Wrap(err)
 	}
 	if !IsCreate(l.Create) {
-		return rn.Pipe(MatchElement(name, value))
+		return MatchElement(name, value), nil
 	}
 
 	var elem *RNode
@@ -443,15 +499,15 @@ func (l PathGetter) doElem(rn *RNode, part string) (*RNode, error) {
 		})
 	}
 	// Append the Node
-	return rn.Pipe(ElementMatcher{FieldName: name, FieldValue: value, Create: elem})
+	return ElementMatcher{FieldName: name, FieldValue: value, Create: elem}, nil
 }
 
-func (l PathGetter) doField(
-	rn *RNode, name string, kind yaml.Kind) (*RNode, error) {
+func (l PathGetter) fieldFilter(
+	name string, kind yaml.Kind) (Filter, error) {
 	if !IsCreate(l.Create) {
-		return rn.Pipe(Get(name))
+		return Get(name), nil
 	}
-	return rn.Pipe(FieldMatcher{Name: name, Create: &RNode{value: &yaml.Node{Kind: kind, Style: l.Style}}})
+	return FieldMatcher{Name: name, Create: &RNode{value: &yaml.Node{Kind: kind, Style: l.Style}}}, nil
 }
 
 func (l PathGetter) getKind(nextPart string) yaml.Kind {

kyaml/yaml/fns_test.go

@@ -57,6 +57,27 @@ func TestAppend(t *testing.T) {
 	assert.Nil(t, rn)
 }
 
+func TestGetElementByIndex(t *testing.T) {
+	node, err := Parse(`
+- 0
+- 1
+- 2
+`)
+	assert.NoError(t, err)
+
+	rn, err := node.Pipe(GetElementByIndex(0))
+	assert.NoError(t, err)
+	assert.Equal(t, "0\n", assertNoErrorString(t)(rn.String()))
+
+	rn, err = node.Pipe(GetElementByIndex(2))
+	assert.NoError(t, err)
+	assert.Equal(t, "2\n", assertNoErrorString(t)(rn.String()))
+
+	rn, err = node.Pipe(GetElementByIndex(-1))
+	assert.NoError(t, err)
+	assert.Equal(t, "2\n", assertNoErrorString(t)(rn.String()))
+}
+
 func TestGetElementByKey(t *testing.T) {
 	node, err := Parse(`
 - b: c
@@ -340,6 +361,18 @@ j: k
   p: q
 `, assertNoErrorString(t)(rn.String()))
 
+	rn, err = node.Pipe(Lookup("a", "b", "0"))
+	assert.NoError(t, err)
+	assert.Equal(t, s, assertNoErrorString(t)(node.String()))
+	assert.Equal(t, `f: g
+`, assertNoErrorString(t)(rn.String()))
+
+	rn, err = node.Pipe(Lookup("a", "b", "-", "h"))
+	assert.NoError(t, err)
+	assert.Equal(t, s, assertNoErrorString(t)(node.String()))
+	assert.Equal(t, `i
+`, assertNoErrorString(t)(rn.String()))
+
 	rn, err = node.Pipe(Lookup("l"))
 	assert.NoError(t, err)
 	assert.Equal(t, s, assertNoErrorString(t)(node.String()))
@@ -384,6 +417,16 @@ j: k
 	assert.NoError(t, err)
 	assert.Equal(t, s, assertNoErrorString(t)(node.String()))
 	assert.Nil(t, rn)
+
+	rn, err = node.Pipe(Lookup("a", "b", "-1"))
+	assert.Errorf(t, err, "array index -1 cannot be negative")
+	assert.Equal(t, s, assertNoErrorString(t)(node.String()))
+	assert.Nil(t, rn)
+
+	rn, err = node.Pipe(Lookup("a", "b", "99"))
+	assert.NoError(t, err)
+	assert.Equal(t, s, assertNoErrorString(t)(node.String()))
+	assert.Nil(t, rn)
 }
 
 func TestSetField_Fn(t *testing.T) {