Merge pull request #46 from itchyny/add-custom-funcs

README.md

@@ -118,6 +118,7 @@ func main() {
 - [`gojq.WithModuleLoader`](https://pkg.go.dev/github.com/itchyny/gojq?tab=doc#WithModuleLoader) allows to load modules. By default, the module feature is disabled. If you want to load modules from the filesystem, use [`gojq.NewModuleLoader`](https://pkg.go.dev/github.com/itchyny/gojq?tab=doc#NewModuleLoader).
 - [`gojq.WithEnvironLoader`](https://pkg.go.dev/github.com/itchyny/gojq?tab=doc#WithEnvironLoader) allows to configure the environment variables referenced by `env` and `$ENV`. By default, OS environment variables are not accessible due to security reason. You can use `gojq.WithEnvironLoader(os.Environ)` if you want.
 - [`gojq.WithVariables`](https://pkg.go.dev/github.com/itchyny/gojq?tab=doc#WithVariables) allows to configure the variables which can be used in the query. Pass the values of the variables to [`code.Run`](https://pkg.go.dev/github.com/itchyny/gojq?tab=doc#Code.Run) in the same order.
+- [`gojq.WithFunction`](https://pkg.go.dev/github.com/itchyny/gojq?tab=doc#WithFunction) allows to add a custom internal function.
 - [`gojq.WithInputIter`](https://pkg.go.dev/github.com/itchyny/gojq?tab=doc#WithInputIter) allows to use `input` and `inputs` functions. By default, these functions are disabled.
 
 ## Bug Tracker

compiler.go

@@ -14,6 +14,7 @@ type compiler struct {
 	moduleLoader  ModuleLoader
 	environLoader func() []string
 	variables     []string
+	customFuncs   map[string]function
 	inputIter     Iter
 	codes         []*code
 	codeinfos     []codeinfo
@@ -195,7 +196,8 @@ func (c *compiler) compileImport(i *Import) error {
 
 func (c *compiler) compileModule(q *Query, alias string) error {
 	cc := &compiler{
-		moduleLoader: c.moduleLoader, environLoader: c.environLoader, variables: c.variables, inputIter: c.inputIter,
+		moduleLoader: c.moduleLoader, environLoader: c.environLoader,
+		variables: c.variables, customFuncs: c.customFuncs, inputIter: c.inputIter,
 		codeoffset: c.pc(), scopes: c.scopes, scopecnt: c.scopecnt,
 	}
 	defer cc.newScopeDepth()()
@@ -282,7 +284,8 @@ func (c *compiler) compileFuncDef(e *FuncDef, builtin bool) error {
 	pc, argsorder := c.pc(), getArgsOrder(e.Args)
 	c.funcs = append(c.funcs, &funcinfo{e.Name, pc, e.Args, argsorder})
 	cc := &compiler{
-		moduleLoader: c.moduleLoader, environLoader: c.environLoader, inputIter: c.inputIter,
+		moduleLoader: c.moduleLoader, environLoader: c.environLoader,
+		customFuncs: c.customFuncs, inputIter: c.inputIter,
 		codeoffset: pc, scopecnt: c.scopecnt, funcs: c.funcs,
 	}
 	scope := cc.newScope()
@@ -904,6 +907,13 @@ func (c *compiler) compileFunc(e *Func) error {
 		case "stderr":
 			c.append(&code{op: opdebug, v: "STDERR:"})
 			return nil
+		case "builtins":
+			return c.compileCallInternal(
+				[3]interface{}{c.funcBuiltins, 0, e.Name},
+				e.Args,
+				nil,
+				false,
+			)
 		case "input":
 			if c.inputIter == nil {
 				return &inputNotAllowedError{}
@@ -925,9 +935,59 @@ func (c *compiler) compileFunc(e *Func) error {
 			return c.compileCall(e.Name, e.Args)
 		}
 	}
+	if fn, ok := c.customFuncs[e.Name]; ok && fn.accept(len(e.Args)) {
+		return c.compileCallInternal(
+			[3]interface{}{fn.callback, len(e.Args), e.Name},
+			e.Args,
+			nil,
+			false,
+		)
+	}
 	return &funcNotFoundError{e}
 }
 
+func (c *compiler) funcBuiltins(interface{}, []interface{}) interface{} {
+	type funcNameArity struct {
+		name  string
+		arity int
+	}
+	var xs []*funcNameArity
+	for _, fds := range builtinFuncDefs {
+		for _, fd := range fds {
+			if fd.Name[0] != '_' {
+				xs = append(xs, &funcNameArity{fd.Name, len(fd.Args)})
+			}
+		}
+	}
+	for name, fn := range internalFuncs {
+		if name[0] != '_' {
+			for i, cnt := 0, fn.argcount; cnt > 0; i, cnt = i+1, cnt>>1 {
+				if cnt&1 > 0 {
+					xs = append(xs, &funcNameArity{name, i})
+				}
+			}
+		}
+	}
+	for name, fn := range c.customFuncs {
+		if name[0] != '_' {
+			for i, cnt := 0, fn.argcount; cnt > 0; i, cnt = i+1, cnt>>1 {
+				if cnt&1 > 0 {
+					xs = append(xs, &funcNameArity{name, i})
+				}
+			}
+		}
+	}
+	sort.Slice(xs, func(i, j int) bool {
+		return xs[i].name < xs[j].name ||
+			xs[i].name == xs[j].name && xs[i].arity < xs[j].arity
+	})
+	ys := make([]interface{}, len(xs))
+	for i, x := range xs {
+		ys[i] = x.name + "/" + fmt.Sprint(x.arity)
+	}
+	return ys
+}
+
 func (c *compiler) funcInput(interface{}, []interface{}) interface{} {
 	v, ok := c.inputIter.Next()
 	if !ok {

func.go

@@ -45,6 +45,7 @@ func init() {
 		"debug":          argFunc0(nil),
 		"stderr":         argFunc0(nil),
 		"env":            argFunc0(nil),
+		"builtins":       argFunc0(nil),
 		"input":          argFunc0(nil),
 		"modulemeta":     argFunc0(nil),
 		"length":         argFunc0(funcLength),
@@ -174,7 +175,6 @@ func init() {
 		"error":          {argcount0 | argcount1, funcError},
 		"halt":           argFunc0(funcHalt),
 		"halt_error":     {argcount0 | argcount1, funcHaltError},
-		"builtins":       argFunc0(funcBuiltins),
 		"_type_error":    argFunc1(internalfuncTypeError),
 	}
 }
@@ -1617,32 +1617,6 @@ func funcHaltError(v interface{}, args []interface{}) interface{} {
 	return &exitCodeError{v, code, true}
 }
 
-func funcBuiltins(interface{}) interface{} {
-	var xs []string
-	for name, fn := range internalFuncs {
-		if name[0] != '_' {
-			for i, cnt := 0, fn.argcount; cnt > 0; i, cnt = i+1, cnt>>1 {
-				if cnt&1 > 0 {
-					xs = append(xs, name+"/"+fmt.Sprint(i))
-				}
-			}
-		}
-	}
-	for _, fds := range builtinFuncDefs {
-		for _, fd := range fds {
-			if fd.Name[0] != '_' {
-				xs = append(xs, fd.Name+"/"+fmt.Sprint(len(fd.Args)))
-			}
-		}
-	}
-	sort.Strings(xs)
-	ys := make([]interface{}, len(xs))
-	for i, x := range xs {
-		ys[i] = x
-	}
-	return ys
-}
-
 func internalfuncTypeError(v, x interface{}) interface{} {
 	if x, ok := x.(string); ok {
 		return &funcTypeError{x, v}

option.go

@@ -1,5 +1,7 @@
 package gojq
 
+import "fmt"
+
 // CompilerOption ...
 type CompilerOption func(*compiler)
 
@@ -29,6 +31,36 @@ func WithVariables(variables []string) CompilerOption {
 	}
 }
 
+// WithFunction is a compiler option for adding a custom internal function.
+// Specify the minimum and maximum count of the function arguments. These
+// values should satisfy 0 <= minarity <= maxarity <= 30, otherwise panics.
+// On handling numbers, you should take account to int, float64 and *big.Int.
+func WithFunction(name string, minarity int, maxarity int,
+	f func(interface{}, []interface{}) interface{}) CompilerOption {
+	if !(0 <= minarity && minarity <= maxarity && maxarity <= 30) {
+		panic(fmt.Sprintf("invalid arity for %q: %d, %d", name, minarity, maxarity))
+	}
+	argcount := 1<<(maxarity+1) - 1<<minarity
+	return func(c *compiler) {
+		if c.customFuncs == nil {
+			c.customFuncs = make(map[string]function)
+		}
+		if fn, ok := c.customFuncs[name]; ok {
+			c.customFuncs[name] = function{
+				argcount | fn.argcount,
+				func(x interface{}, xs []interface{}) interface{} {
+					if argcount&(1<<len(xs)) != 0 {
+						return f(x, xs)
+					}
+					return fn.callback(x, xs)
+				},
+			}
+		} else {
+			c.customFuncs[name] = function{argcount, f}
+		}
+	}
+}
+
 // WithInputIter is a compiler option for input iterator used by input(s)/0.
 // Note that input and inputs functions are not allowed by default. We have
 // to distinguish the query input and the values for input(s) functions. For

option_function_test.go

@@ -0,0 +1,71 @@
+package gojq_test
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"math/big"
+	"strconv"
+
+	"github.com/itchyny/gojq"
+)
+
+func toFloat(x interface{}) (float64, bool) {
+	switch x := x.(type) {
+	case int:
+		return float64(x), true
+	case float64:
+		return x, true
+	case *big.Int:
+		f, err := strconv.ParseFloat(x.String(), 64)
+		return f, err == nil
+	default:
+		return 0.0, false
+	}
+}
+
+func ExampleWithFunction() {
+	query, err := gojq.Parse(".[] | f | f(3)")
+	if err != nil {
+		log.Fatalln(err)
+	}
+	code, err := gojq.Compile(
+		query,
+		gojq.WithFunction("f", 0, 1, func(x interface{}, xs []interface{}) interface{} {
+			if x, ok := toFloat(x); ok {
+				if len(xs) == 1 {
+					if y, ok := toFloat(xs[0]); ok {
+						x *= y
+					} else {
+						return fmt.Errorf("f cannot be applied to: %v, %v", x, xs)
+					}
+				} else {
+					x += 2
+				}
+				return x
+			}
+			return fmt.Errorf("f cannot be applied to: %v, %v", x, xs)
+		}),
+	)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	input := []interface{}{0, 1, 2.5, json.Number("10000000000000000000000000000000000000000")}
+	iter := code.Run(input)
+	for {
+		v, ok := iter.Next()
+		if !ok {
+			break
+		}
+		if err, ok := v.(error); ok {
+			log.Fatalln(err)
+		}
+		fmt.Printf("%#v\n", v)
+	}
+
+	// Output:
+	// 6
+	// 9
+	// 13.5
+	// 3e+40
+}