Add user-defined function support without any changes to pre-existing code

.gitignore

@@ -2,3 +2,4 @@
 jmespath-fuzz.zip
 cpu.out
 go-jmespath.test
+.idea

README.md

@@ -2,7 +2,7 @@
 
 [![Build Status](https://img.shields.io/travis/jmespath/go-jmespath.svg)](https://travis-ci.org/jmespath/go-jmespath)
 
-
+NOTE: This is a fork of [go-jmespath](https://github.com/jmespath/go-jmespath) with support for user-defined functions
 
 go-jmespath is a GO implementation of JMESPath,
 which is a query language for JSON.  It will take a JSON
@@ -71,6 +71,31 @@ you are going to run multiple searches with it:
 	result = "bar"
 ```
 
+## User-defined Functions
+
+User-defined functions are added to precompiled queries as follows:
+
+```go
+precompiled, err := Compile("icontains(@, 'Bar')")
+err = precompiled.RegisterFunction("icontains", "string|array[string],string", false, func(args []interface{}) (interface{}, error) {
+    needle := strings.ToLower(args[1].(string))
+    if haystack, ok := args[0].(string); ok {
+        return strings.Contains(strings.ToLower(haystack), needle), nil
+    }
+    array, _ := toArrayStr(args[0])
+    for _, el := range array {
+        if strings.ToLower(el) == needle {
+            return true, nil
+        }
+    }
+    return false, nil
+})
+result, err = searcher.Search([]interface{}{"foo", "BAR", "baz"})
+```
+
+Support for JMESPath expression arguments (as used by `map()`, for example) is provided through the `NewExpressionEvaluator()` function.
+See the [test cases](userfn_test.go) for an example.
+
 ## More Resources
 
 The example above only show a small amount of what

userfn.go

@@ -0,0 +1,51 @@
+package jmespath
+
+import (
+	"fmt"
+	"strings"
+)
+
+type ExpressionEvaluator func(value interface{}) (interface{}, error)
+
+func NewExpressionEvaluator(intrArg interface{}, expArg interface{}) ExpressionEvaluator {
+	intr := intrArg.(*treeInterpreter)
+	node := expArg.(expRef).ref
+	return func(value interface{}) (interface{}, error) {
+		return intr.Execute(node, value)
+	}
+}
+
+func (jp *JMESPath) RegisterFunction(name string, args string, variadic bool, handler func([]interface{}) (interface{}, error)) error {
+	hasExpRef := false
+	var arguments []argSpec
+	for _, arg := range strings.Split(args, ",") {
+		var argTypes []jpType
+		for _, argType := range strings.Split(arg, "|") {
+			switch t := jpType(argType); t {
+			case jpExpref:
+				hasExpRef = true
+				fallthrough
+			case jpNumber, jpString, jpArray, jpObject, jpArrayNumber, jpArrayString, jpAny:
+				argTypes = append(argTypes, t)
+			default:
+				return fmt.Errorf("unknown argument type: %s", argType)
+			}
+		}
+		arguments = append(arguments, argSpec{
+			types: argTypes,
+		})
+	}
+	if variadic {
+		if len(arguments) == 0 {
+			return fmt.Errorf("variadic functions require at least one argument")
+		}
+		arguments[len(arguments)-1].variadic = true
+	}
+	jp.intr.fCall.functionTable[name] = functionEntry{
+		name:      name,
+		arguments: arguments,
+		handler:   handler,
+		hasExpRef: hasExpRef,
+	}
+	return nil
+}

userfn_test.go

@@ -0,0 +1,84 @@
+package jmespath
+
+import (
+	"github.com/jmespath/go-jmespath/internal/testify/assert"
+	"strings"
+	"testing"
+)
+
+func TestUserDefinedFunctions(t *testing.T) {
+	searcher, err := Compile("icontains(@, 'Bar')")
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	err = searcher.RegisterFunction("icontains", "string|array[string],string", false, func(args []interface{}) (interface{}, error) {
+		needle := strings.ToLower(args[1].(string))
+		if haystack, ok := args[0].(string); ok {
+			return strings.Contains(strings.ToLower(haystack), needle), nil
+		}
+		array, _ := toArrayStr(args[0])
+		for _, el := range array {
+			if strings.ToLower(el) == needle {
+				return true, nil
+			}
+		}
+		return false, nil
+	})
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	actual, err := searcher.Search("fooBARbaz")
+	if assert.NoError(t, err) {
+		assert.Equal(t, true, actual)
+	}
+
+	actual, err = searcher.Search([]interface{}{"foo", "BAR", "baz"})
+	if assert.NoError(t, err) {
+		assert.Equal(t, true, actual)
+	}
+}
+
+func TestExpressionEvaluator(t *testing.T) {
+	searcher, err := Compile("my_map(&id, @)")
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	err = searcher.RegisterFunction("my_map", "expref,array", false, func(args []interface{}) (interface{}, error) {
+		evaluator := NewExpressionEvaluator(args[0], args[1])
+		arr := args[2].([]interface{})
+		mapped := make([]interface{}, 0, len(arr))
+		for _, value := range arr {
+			current, err := evaluator(value)
+			if err != nil {
+				return nil, err
+			}
+			mapped = append(mapped, current)
+		}
+		return mapped, nil
+	})
+
+	if !assert.NoError(t, err) {
+		return
+	}
+
+	actual, err := searcher.Search([]interface{}{
+		map[string]interface{}{
+			"id":    1,
+			"value": "a",
+		},
+		map[string]interface{}{
+			"id":    2,
+			"value": "b",
+		},
+		map[string]interface{}{
+			"id":    3,
+			"value": "c",
+		},
+	})
+	if assert.NoError(t, err) {
+		assert.Equal(t, []interface{}{1, 2, 3}, actual)
+	}
+}