From cac3d77b68ee38dad6c834c156ec8ba2835bce26 Mon Sep 17 00:00:00 2001 From: Andrew Klopper Date: Mon, 8 May 2023 18:50:54 +0200 Subject: [PATCH 1/2] Add user-defined function support without any changes to pre-existing code --- .gitignore | 1 + userfn.go | 51 ++++++++++++++++++++++++++++++ userfn_test.go | 84 ++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 136 insertions(+) create mode 100644 userfn.go create mode 100644 userfn_test.go 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/userfn.go b/userfn.go new file mode 100644 index 0000000..1787cb6 --- /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, handler func([]interface{}) (interface{}, error), args string, variadic bool) 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..3a7f047 --- /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", 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 + }, "string|array[string],string", false) + 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", 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 + }, "expref,array", false) + + 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) + } +} From 1e424bf6fe118b5522b7f5adf430095cff055fed Mon Sep 17 00:00:00 2001 From: Andrew Klopper Date: Mon, 8 May 2023 19:25:02 +0200 Subject: [PATCH 2/2] Update documentation. --- README.md | 27 ++++++++++++++++++++++++++- userfn.go | 2 +- userfn_test.go | 8 ++++---- 3 files changed, 31 insertions(+), 6 deletions(-) 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 index 1787cb6..7aaefba 100644 --- a/userfn.go +++ b/userfn.go @@ -15,7 +15,7 @@ func NewExpressionEvaluator(intrArg interface{}, expArg interface{}) ExpressionE } } -func (jp *JMESPath) RegisterFunction(name string, handler func([]interface{}) (interface{}, error), args string, variadic bool) error { +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, ",") { diff --git a/userfn_test.go b/userfn_test.go index 3a7f047..e4eab7a 100644 --- a/userfn_test.go +++ b/userfn_test.go @@ -12,7 +12,7 @@ func TestUserDefinedFunctions(t *testing.T) { return } - err = searcher.RegisterFunction("icontains", func(args []interface{}) (interface{}, error) { + 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 @@ -24,7 +24,7 @@ func TestUserDefinedFunctions(t *testing.T) { } } return false, nil - }, "string|array[string],string", false) + }) if !assert.NoError(t, err) { return } @@ -46,7 +46,7 @@ func TestExpressionEvaluator(t *testing.T) { return } - err = searcher.RegisterFunction("my_map", func(args []interface{}) (interface{}, error) { + 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)) @@ -58,7 +58,7 @@ func TestExpressionEvaluator(t *testing.T) { mapped = append(mapped, current) } return mapped, nil - }, "expref,array", false) + }) if !assert.NoError(t, err) { return