New coll.Index function

Signed-off-by: Dave Henderson <[email protected]>

coll/index.go

@@ -0,0 +1,160 @@
+package coll
+
+import (
+	"fmt"
+	"reflect"
+)
+
+// much of the code here is taken from the Go source code, in particular from
+// text/template/exec.go and text/template/funcs.go
+
+// Index returns the result of indexing the given map, slice, or array by the
+// given index arguments. This is similar to the `index` template function, but
+// will return an error if the key is not found. Note that the argument order is
+// different from the template function definition found in `funcs/coll.go` to
+// allow for variadic indexes.
+func Index(v interface{}, keys ...interface{}) (interface{}, error) {
+	item := reflect.ValueOf(v)
+	item = indirectInterface(item)
+	if !item.IsValid() {
+		return nil, fmt.Errorf("index of untyped nil")
+	}
+
+	indexes := make([]reflect.Value, len(keys))
+	for i, k := range keys {
+		indexes[i] = reflect.ValueOf(k)
+	}
+
+	for _, index := range indexes {
+		index = indirectInterface(index)
+		var isNil bool
+		if item, isNil = indirect(item); isNil {
+			return nil, fmt.Errorf("index of nil pointer")
+		}
+		switch item.Kind() {
+		case reflect.Array, reflect.Slice, reflect.String:
+			x, err := indexArg(index, item.Len())
+			if err != nil {
+				return nil, err
+			}
+
+			item = item.Index(x)
+		case reflect.Map:
+			x, err := prepareArg(index, item.Type().Key())
+			if err != nil {
+				return nil, err
+			}
+
+			if v := item.MapIndex(x); v.IsValid() {
+				item = v
+			} else {
+				// the map doesn't contain the key, so return an error
+				return nil, fmt.Errorf("map has no key %v", x.Interface())
+			}
+		case reflect.Invalid:
+			// the loop holds invariant: item.IsValid()
+			panic("unreachable")
+		default:
+			return nil, fmt.Errorf("can't index item of type %s", item.Type())
+		}
+	}
+
+	return item.Interface(), nil
+}
+
+// indexArg checks if a reflect.Value can be used as an index, and converts it to int if possible.
+func indexArg(index reflect.Value, cap int) (int, error) {
+	var x int64
+	switch index.Kind() {
+	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+		x = index.Int()
+	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
+		x = int64(index.Uint())
+	case reflect.Invalid:
+		return 0, fmt.Errorf("cannot index slice/array with nil")
+	default:
+		return 0, fmt.Errorf("cannot index slice/array with type %s", index.Type())
+	}
+
+	// note - this has been modified from the original to check for x == cap as
+	// well. IMO the original (> only) is a bug.
+	if x < 0 || int(x) < 0 || int(x) >= cap {
+		return 0, fmt.Errorf("index out of range: %d", x)
+	}
+
+	return int(x), nil
+}
+
+// prepareArg checks if value can be used as an argument of type argType, and
+// converts an invalid value to appropriate zero if possible.
+func prepareArg(value reflect.Value, argType reflect.Type) (reflect.Value, error) {
+	if !value.IsValid() {
+		if !canBeNil(argType) {
+			return reflect.Value{}, fmt.Errorf("value is nil; should be of type %s", argType)
+		}
+
+		value = reflect.Zero(argType)
+	}
+
+	if value.Type().AssignableTo(argType) {
+		return value, nil
+	}
+
+	if intLike(value.Kind()) && intLike(argType.Kind()) && value.Type().ConvertibleTo(argType) {
+		value = value.Convert(argType)
+
+		return value, nil
+	}
+
+	return reflect.Value{}, fmt.Errorf("value has type %s; should be %s", value.Type(), argType)
+}
+
+var reflectValueType = reflect.TypeOf((*reflect.Value)(nil)).Elem()
+
+// canBeNil reports whether an untyped nil can be assigned to the type. See reflect.Zero.
+func canBeNil(typ reflect.Type) bool {
+	switch typ.Kind() {
+	case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice:
+		return true
+	case reflect.Struct:
+		return typ == reflectValueType
+	}
+
+	return false
+}
+
+func intLike(typ reflect.Kind) bool {
+	switch typ {
+	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+		return true
+	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
+		return true
+	}
+	return false
+}
+
+// indirect returns the item at the end of indirection, and a bool to indicate
+// if it's nil. If the returned bool is true, the returned value's kind will be
+// either a pointer or interface.
+func indirect(v reflect.Value) (rv reflect.Value, isNil bool) {
+	for ; v.Kind() == reflect.Pointer || v.Kind() == reflect.Interface; v = v.Elem() {
+		if v.IsNil() {
+			return v, true
+		}
+	}
+	return v, false
+}
+
+// indirectInterface returns the concrete value in an interface value,
+// or else the zero reflect.Value.
+// That is, if v represents the interface value x, the result is the same as reflect.ValueOf(x):
+// the fact that x was an interface value is forgotten.
+func indirectInterface(v reflect.Value) reflect.Value {
+	if v.Kind() != reflect.Interface {
+		return v
+	}
+	if v.IsNil() {
+		return reflect.Value{}
+	}
+	return v.Elem()
+}

coll/index_test.go

@@ -0,0 +1,50 @@
+package coll
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIndex(t *testing.T) {
+	out, err := Index(map[string]interface{}{
+		"foo": "bar", "baz": "qux",
+	}, "foo")
+	assert.NoError(t, err)
+	assert.Equal(t, "bar", out)
+
+	out, err = Index(map[string]interface{}{
+		"foo": "bar", "baz": "qux", "quux": "corge",
+	}, "foo", 2)
+	assert.NoError(t, err)
+	assert.Equal(t, byte('r'), out)
+
+	out, err = Index([]interface{}{"foo", "bar", "baz"}, 2)
+	assert.NoError(t, err)
+	assert.Equal(t, "baz", out)
+
+	out, err = Index([]interface{}{"foo", "bar", "baz"}, 2, 2)
+	assert.NoError(t, err)
+	assert.Equal(t, byte('z'), out)
+
+	// error cases
+	out, err = Index([]interface{}{"foo", "bar", "baz"}, 0, 1, 2)
+	assert.Error(t, err)
+	assert.Nil(t, out)
+
+	out, err = Index(nil, 0)
+	assert.Error(t, err)
+	assert.Nil(t, out)
+
+	out, err = Index("foo", nil)
+	assert.Error(t, err)
+	assert.Nil(t, out)
+
+	out, err = Index(map[interface{}]string{nil: "foo", 2: "bar"}, "baz")
+	assert.Error(t, err)
+	assert.Nil(t, out)
+
+	out, err = Index([]int{}, 0)
+	assert.Error(t, err)
+	assert.Nil(t, out)
+}

docs-src/content/functions/coll.yml

@@ -114,6 +114,34 @@ funcs:
         $ gomplate -i '{{ $o := data.JSON (getenv "DATA") -}}
         {{ if (has $o "foo") }}{{ $o.foo }}{{ else }}THERE IS NO FOO{{ end }}'
         THERE IS NO FOO
+  - name: coll.Index
+    description: |
+      Returns the result of indexing the given map, slice, or array by the given
+      key or index. This is similar to the built-in `index` function, but the
+      arguments are ordered differently for pipeline compatibility. Also this
+      function is more strict, and will return an error when trying to access a
+      non-existent map key.
+
+      Multiple indexes may be given, for nested indexing.
+    pipeline: true
+    arguments:
+      - name: indexes...
+        required: true
+        description: The key or index 
+      - name: in
+        required: true
+        description: The map, slice, or array to index
+    examples:
+      - |
+        $ gomplate -i '{{ coll.Index "foo" (dict "foo" "bar") }}'
+        bar
+      - |
+        $ gomplate -i '{{ $m := json `{"foo": {"bar": "baz"}}` -}}
+          {{ coll.Index "foo" "bar" $m }}'
+        baz
+      - |
+        $ gomplate -i '{{ coll.Slice "foo" "bar" "baz" | coll.Index 1 }}'
+        bar
   - name: coll.JSONPath
     alias: jsonpath
     description: |

docs/content/functions/coll.md

@@ -165,6 +165,48 @@ $ gomplate -i '{{ $o := data.JSON (getenv "DATA") -}}
 THERE IS NO FOO
 ```
 
+## `coll.Index`
+
+Returns the result of indexing the given map, slice, or array by the given
+key or index. This is similar to the built-in `index` function, but the
+arguments are ordered differently for pipeline compatibility. Also this
+function is more strict, and will return an error when trying to access a
+non-existent map key.
+
+Multiple indexes may be given, for nested indexing.
+
+### Usage
+
+```go
+coll.Index indexes... in
+```
+```go
+in | coll.Index indexes...
+```
+
+### Arguments
+
+| name | description |
+|------|-------------|
+| `indexes...` | _(required)_ The key or index |
+| `in` | _(required)_ The map, slice, or array to index |
+
+### Examples
+
+```console
+$ gomplate -i '{{ coll.Index "foo" (dict "foo" "bar") }}'
+bar
+```
+```console
+$ gomplate -i '{{ $m := json `{"foo": {"bar": "baz"}}` -}}
+  {{ coll.Index "foo" "bar" $m }}'
+baz
+```
+```console
+$ gomplate -i '{{ coll.Slice "foo" "bar" "baz" | coll.Index 1 }}'
+bar
+```
+
 ## `coll.JSONPath`
 
 **Alias:** `jsonpath`

funcs/coll.go

@@ -2,6 +2,7 @@ package funcs
 
 import (
 	"context"
+	"fmt"
 	"reflect"
 
 	"github.com/hairyhenderson/gomplate/v3/conv"
@@ -79,6 +80,22 @@ func (CollFuncs) Has(in interface{}, key string) bool {
 	return coll.Has(in, key)
 }
 
+// Index returns the result of indexing the last argument with the preceding
+// index keys. This is similar to the `index` built-in template function, but
+// the arguments are ordered differently for pipeline compatibility. Also, this
+// function is more strict, and will return an error when the value doesn't
+// contain the given key.
+func (CollFuncs) Index(args ...interface{}) (interface{}, error) {
+	if len(args) < 2 {
+		return nil, fmt.Errorf("wrong number of args: wanted at least 2, got %d", len(args))
+	}
+
+	item := args[len(args)-1]
+	indexes := args[:len(args)-1]
+
+	return coll.Index(item, indexes...)
+}
+
 // Dict -
 func (CollFuncs) Dict(in ...interface{}) (map[string]interface{}, error) {
 	return coll.Dict(in...)