diff --git a/.gitignore b/.gitignore index 5091fb0..cf83e55 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ jmespath-fuzz.zip cpu.out go-jmespath.test +.idea diff --git a/README.md b/README.md index 110ad79..e8a93c1 100644 --- a/README.md +++ b/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 diff --git a/userfn.go b/userfn.go new file mode 100644 index 0000000..7aaefba --- /dev/null +++ b/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 +} diff --git a/userfn_test.go b/userfn_test.go new file mode 100644 index 0000000..e4eab7a --- /dev/null +++ b/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) + } +}