From a45a6e5b414f6cb838005d35d093fbcdc2e72d41 Mon Sep 17 00:00:00 2001
From: AuroraV <zzl_auroras@163.com>
Date: Sun, 2 Jul 2023 12:37:35 +0800
Subject: [PATCH] feat:customFuncs add ctx support

---
 cli/cli.go                   |  9 ++++---
 compiler.go                  |  4 +--
 execute.go                   |  2 +-
 func.go                      | 23 +++++++++--------
 option.go                    | 21 +++++++++-------
 option_function_test.go      |  3 ++-
 option_iter_function_test.go |  3 ++-
 option_test.go               | 49 ++++++++++++++++++------------------
 8 files changed, 61 insertions(+), 53 deletions(-)

diff --git a/cli/cli.go b/cli/cli.go
index a10ead55..83d08d52 100644
--- a/cli/cli.go
+++ b/cli/cli.go
@@ -2,6 +2,7 @@
 package cli
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -242,8 +243,8 @@ Usage:
 		gojq.WithFunction("debug", 0, 0, cli.funcDebug),
 		gojq.WithFunction("stderr", 0, 0, cli.funcStderr),
 		gojq.WithFunction("input_filename", 0, 0,
-			func(iter inputIter) func(any, []any) any {
-				return func(any, []any) any {
+			func(iter inputIter) func(context.Context, any, []any) any {
+				return func(context.Context, any, []any) any {
 					if fname := iter.Name(); fname != "" && (len(args) > 0 || !opts.InputNull) {
 						return fname
 					}
@@ -408,7 +409,7 @@ func (cli *cli) createMarshaler() marshaler {
 	return f
 }
 
-func (cli *cli) funcDebug(v any, _ []any) any {
+func (cli *cli) funcDebug(_ context.Context, v any, _ []any) any {
 	if err := newEncoder(false, 0).marshal([]any{"DEBUG:", v}, cli.errStream); err != nil {
 		return err
 	}
@@ -418,7 +419,7 @@ func (cli *cli) funcDebug(v any, _ []any) any {
 	return v
 }
 
-func (cli *cli) funcStderr(v any, _ []any) any {
+func (cli *cli) funcStderr(_ context.Context, v any, _ []any) any {
 	if err := newEncoder(false, 0).marshal(v, cli.errStream); err != nil {
 		return err
 	}
diff --git a/compiler.go b/compiler.go
index de5f9a10..56801d45 100644
--- a/compiler.go
+++ b/compiler.go
@@ -804,8 +804,8 @@ func (c *compiler) compileBreak(label string) error {
 	return nil
 }
 
-func funcBreak(label string) func(any, []any) any {
-	return func(v any, _ []any) any {
+func funcBreak(label string) func(context.Context, any, []any) any {
+	return func(_ context.Context, v any, _ []any) any {
 		return &breakError{label, v}
 	}
 }
diff --git a/execute.go b/execute.go
index dcf9d984..318537e4 100644
--- a/execute.go
+++ b/execute.go
@@ -189,7 +189,7 @@ loop:
 				for i := 0; i < argcnt; i++ {
 					args[i] = env.pop()
 				}
-				w := v[0].(func(any, []any) any)(x, args)
+				w := v[0].(func(context.Context, any, []any) any)(env.ctx, x, args)
 				if e, ok := w.(error); ok {
 					if er, ok := e.(*exitCodeError); !ok || er.value != nil || er.halt {
 						err = e
diff --git a/func.go b/func.go
index 6e8d1500..607e0897 100644
--- a/func.go
+++ b/func.go
@@ -1,6 +1,7 @@
 package gojq
 
 import (
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
@@ -33,7 +34,7 @@ const (
 type function struct {
 	argcount int
 	iter     bool
-	callback func(any, []any) any
+	callback func(context.Context, any, []any) any
 }
 
 func (fn function) accept(cnt int) bool {
@@ -202,7 +203,7 @@ func init() {
 
 func argFunc0(f func(any) any) function {
 	return function{
-		argcount0, false, func(v any, _ []any) any {
+		argcount0, false, func(_ context.Context, v any, _ []any) any {
 			return f(v)
 		},
 	}
@@ -210,7 +211,7 @@ func argFunc0(f func(any) any) function {
 
 func argFunc1(f func(_, _ any) any) function {
 	return function{
-		argcount1, false, func(v any, args []any) any {
+		argcount1, false, func(_ context.Context, v any, args []any) any {
 			return f(v, args[0])
 		},
 	}
@@ -218,7 +219,7 @@ func argFunc1(f func(_, _ any) any) function {
 
 func argFunc2(f func(_, _, _ any) any) function {
 	return function{
-		argcount2, false, func(v any, args []any) any {
+		argcount2, false, func(_ context.Context, v any, args []any) any {
 			return f(v, args[0], args[1])
 		},
 	}
@@ -226,7 +227,7 @@ func argFunc2(f func(_, _, _ any) any) function {
 
 func argFunc3(f func(_, _, _, _ any) any) function {
 	return function{
-		argcount3, false, func(v any, args []any) any {
+		argcount3, false, func(_ context.Context, v any, args []any) any {
 			return f(v, args[0], args[1], args[2])
 		},
 	}
@@ -718,7 +719,7 @@ func funcImplode(v any) any {
 	return sb.String()
 }
 
-func funcSplit(v any, args []any) any {
+func funcSplit(_ context.Context, v any, args []any) any {
 	s, ok := v.(string)
 	if !ok {
 		return &func0TypeError{"split", v}
@@ -809,7 +810,7 @@ func funcFormat(v, x any) any {
 	if f == nil {
 		return &formatNotFoundError{format}
 	}
-	return internalFuncs[f.Name].callback(v, nil)
+	return internalFuncs[f.Name].callback(context.Background(), v, nil)
 }
 
 var htmlEscaper = strings.NewReplacer(
@@ -1101,7 +1102,7 @@ func clampIndex(i, min, max int) int {
 	}
 }
 
-func funcFlatten(v any, args []any) any {
+func funcFlatten(_ context.Context, v any, args []any) any {
 	vs, ok := values(v)
 	if !ok {
 		return &func0TypeError{"flatten", v}
@@ -1145,7 +1146,7 @@ func (iter *rangeIter) Next() (any, bool) {
 	return v, true
 }
 
-func funcRange(_ any, xs []any) any {
+func funcRange(_ context.Context, _ any, xs []any) any {
 	for _, x := range xs {
 		switch x.(type) {
 		case int, float64, *big.Int:
@@ -2048,7 +2049,7 @@ func funcCapture(v any) any {
 	return w
 }
 
-func funcError(v any, args []any) any {
+func funcError(_ context.Context, v any, args []any) any {
 	if len(args) > 0 {
 		v = args[0]
 	}
@@ -2063,7 +2064,7 @@ func funcHalt(any) any {
 	return &exitCodeError{nil, 0, true}
 }
 
-func funcHaltError(v any, args []any) any {
+func funcHaltError(_ context.Context, v any, args []any) any {
 	code := 5
 	if len(args) > 0 {
 		var ok bool
diff --git a/option.go b/option.go
index f1a110fa..07d14cc6 100644
--- a/option.go
+++ b/option.go
@@ -1,6 +1,9 @@
 package gojq
 
-import "fmt"
+import (
+	"context"
+	"fmt"
+)
 
 // CompilerOption is a compiler option.
 type CompilerOption func(*compiler)
@@ -39,7 +42,7 @@ func WithVariables(variables []string) CompilerOption {
 // function. If you want to emit multiple values, call the empty function,
 // accept a filter for its argument, or call another built-in function, then
 // use LoadInitModules of the module loader.
-func WithFunction(name string, minarity, maxarity int, f func(any, []any) any) CompilerOption {
+func WithFunction(name string, minarity, maxarity int, f func(context.Context, any, []any) any) CompilerOption {
 	return withFunction(name, minarity, maxarity, false, f)
 }
 
@@ -48,15 +51,15 @@ func WithFunction(name string, minarity, maxarity int, f func(any, []any) any) C
 // returns an Iter to emit multiple values. You cannot define both iterator and
 // non-iterator functions of the same name (with possibly different arities).
 // See also [NewIter], which can be used to convert values or an error to an Iter.
-func WithIterFunction(name string, minarity, maxarity int, f func(any, []any) Iter) CompilerOption {
+func WithIterFunction(name string, minarity, maxarity int, f func(context.Context, any, []any) Iter) CompilerOption {
 	return withFunction(name, minarity, maxarity, true,
-		func(v any, args []any) any {
-			return f(v, args)
+		func(ctx context.Context, v any, args []any) any {
+			return f(ctx, v, args)
 		},
 	)
 }
 
-func withFunction(name string, minarity, maxarity int, iter bool, f func(any, []any) any) CompilerOption {
+func withFunction(name string, minarity, maxarity int, iter bool, f func(context.Context, any, []any) any) CompilerOption {
 	if !(0 <= minarity && minarity <= maxarity && maxarity <= 30) {
 		panic(fmt.Sprintf("invalid arity for %q: %d, %d", name, minarity, maxarity))
 	}
@@ -71,11 +74,11 @@ func withFunction(name string, minarity, maxarity int, iter bool, f func(any, []
 			}
 			c.customFuncs[name] = function{
 				argcount | fn.argcount, iter,
-				func(x any, xs []any) any {
+				func(ctx context.Context, x any, xs []any) any {
 					if argcount&(1<<len(xs)) != 0 {
-						return f(x, xs)
+						return f(ctx, x, xs)
 					}
-					return fn.callback(x, xs)
+					return fn.callback(ctx, x, xs)
 				},
 			}
 		} else {
diff --git a/option_function_test.go b/option_function_test.go
index fe13ff9a..1fdcff89 100644
--- a/option_function_test.go
+++ b/option_function_test.go
@@ -1,6 +1,7 @@
 package gojq_test
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"log"
@@ -31,7 +32,7 @@ func ExampleWithFunction() {
 	}
 	code, err := gojq.Compile(
 		query,
-		gojq.WithFunction("f", 0, 1, func(x any, xs []any) any {
+		gojq.WithFunction("f", 0, 1, func(_ context.Context, x any, xs []any) any {
 			if x, ok := toFloat(x); ok {
 				if len(xs) == 1 {
 					if y, ok := toFloat(xs[0]); ok {
diff --git a/option_iter_function_test.go b/option_iter_function_test.go
index a7ffc771..0dabe4e2 100644
--- a/option_iter_function_test.go
+++ b/option_iter_function_test.go
@@ -1,6 +1,7 @@
 package gojq_test
 
 import (
+	"context"
 	"fmt"
 	"log"
 
@@ -28,7 +29,7 @@ func ExampleWithIterFunction() {
 	}
 	code, err := gojq.Compile(
 		query,
-		gojq.WithIterFunction("f", 2, 2, func(_ any, xs []any) gojq.Iter {
+		gojq.WithIterFunction("f", 2, 2, func(_ context.Context, _ any, xs []any) gojq.Iter {
 			if x, ok := xs[0].(int); ok {
 				if y, ok := xs[1].(int); ok {
 					return &rangeIter{x, y}
diff --git a/option_test.go b/option_test.go
index a1bc581a..86752e76 100644
--- a/option_test.go
+++ b/option_test.go
@@ -1,6 +1,7 @@
 package gojq_test
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -297,13 +298,13 @@ func TestWithVariablesError2(t *testing.T) {
 
 func TestWithFunction(t *testing.T) {
 	options := []gojq.CompilerOption{
-		gojq.WithFunction("f", 0, 0, func(x any, _ []any) any {
+		gojq.WithFunction("f", 0, 0, func(_ context.Context, x any, _ []any) any {
 			if x, ok := x.(int); ok {
 				return x * 2
 			}
 			return fmt.Errorf("f cannot be applied to: %v", x)
 		}),
-		gojq.WithFunction("g", 1, 1, func(x any, xs []any) any {
+		gojq.WithFunction("g", 1, 1, func(_ context.Context, x any, xs []any) any {
 			if x, ok := x.(int); ok {
 				if y, ok := xs[0].(int); ok {
 					return x + y
@@ -364,13 +365,13 @@ func TestWithFunction(t *testing.T) {
 
 func TestWithFunctionDuplicateName(t *testing.T) {
 	options := []gojq.CompilerOption{
-		gojq.WithFunction("f", 0, 0, func(x any, _ []any) any {
+		gojq.WithFunction("f", 0, 0, func(_ context.Context, x any, _ []any) any {
 			if x, ok := x.(int); ok {
 				return x * 2
 			}
 			return fmt.Errorf("f cannot be applied to: %v", x)
 		}),
-		gojq.WithFunction("f", 1, 1, func(x any, xs []any) any {
+		gojq.WithFunction("f", 1, 1, func(_ context.Context, x any, xs []any) any {
 			if x, ok := x.(int); ok {
 				if y, ok := xs[0].(int); ok {
 					return x + y
@@ -378,13 +379,13 @@ func TestWithFunctionDuplicateName(t *testing.T) {
 			}
 			return fmt.Errorf("f cannot be applied to: %v, %v", x, xs)
 		}),
-		gojq.WithFunction("f", 0, 0, func(x any, _ []any) any {
+		gojq.WithFunction("f", 0, 0, func(_ context.Context, x any, _ []any) any {
 			if x, ok := x.(int); ok {
 				return x * 4
 			}
 			return fmt.Errorf("f cannot be applied to: %v", x)
 		}),
-		gojq.WithFunction("f", 2, 2, func(x any, xs []any) any {
+		gojq.WithFunction("f", 2, 2, func(_ context.Context, x any, xs []any) any {
 			if x, ok := x.(int); ok {
 				if y, ok := xs[0].(int); ok {
 					if z, ok := xs[1].(int); ok {
@@ -447,7 +448,7 @@ func TestWithFunctionDuplicateName(t *testing.T) {
 
 func TestWithFunctionMultipleArities(t *testing.T) {
 	options := []gojq.CompilerOption{
-		gojq.WithFunction("f", 0, 4, func(x any, xs []any) any {
+		gojq.WithFunction("f", 0, 4, func(_ context.Context, x any, xs []any) any {
 			if x, ok := x.(int); ok {
 				x *= 2
 				for _, y := range xs {
@@ -459,7 +460,7 @@ func TestWithFunctionMultipleArities(t *testing.T) {
 			}
 			return fmt.Errorf("f cannot be applied to: %v, %v", x, xs)
 		}),
-		gojq.WithFunction("f", 2, 3, func(x any, xs []any) any {
+		gojq.WithFunction("f", 2, 3, func(_ context.Context, x any, xs []any) any {
 			if x, ok := x.(int); ok {
 				for _, y := range xs {
 					if y, ok := y.(int); ok {
@@ -470,7 +471,7 @@ func TestWithFunctionMultipleArities(t *testing.T) {
 			}
 			return fmt.Errorf("f cannot be applied to: %v, %v", x, xs)
 		}),
-		gojq.WithFunction("g", 0, 30, func(x any, xs []any) any {
+		gojq.WithFunction("g", 0, 30, func(_ context.Context, x any, xs []any) any {
 			return nil
 		}),
 	}
@@ -543,7 +544,7 @@ func TestWithFunctionValueError(t *testing.T) {
 		t.Fatal(err)
 	}
 	code, err := gojq.Compile(query,
-		gojq.WithFunction("f", 0, 0, func(x any, _ []any) any {
+		gojq.WithFunction("f", 0, 0, func(_ context.Context, x any, _ []any) any {
 			return &valueError{expected}
 		}),
 	)
@@ -568,7 +569,7 @@ func TestWithFunctionCompileArgsError(t *testing.T) {
 		t.Fatal(err)
 	}
 	_, err = gojq.Compile(query,
-		gojq.WithFunction("f", 0, 1, func(any, []any) any {
+		gojq.WithFunction("f", 0, 1, func(context.Context, any, []any) any {
 			return 0
 		}),
 	)
@@ -591,7 +592,7 @@ func TestWithFunctionArityError(t *testing.T) {
 				}
 			}()
 			t.Fatal(gojq.Compile(query,
-				gojq.WithFunction("f", tc.min, tc.max, func(any, []any) any {
+				gojq.WithFunction("f", tc.min, tc.max, func(context.Context, any, []any) any {
 					return 0
 				}),
 			))
@@ -605,10 +606,10 @@ func TestWithIterFunction(t *testing.T) {
 		t.Fatal(err)
 	}
 	code, err := gojq.Compile(query,
-		gojq.WithIterFunction("f", 0, 0, func(any, []any) gojq.Iter {
+		gojq.WithIterFunction("f", 0, 0, func(context.Context, any, []any) gojq.Iter {
 			return gojq.NewIter(1, 2, 3)
 		}),
-		gojq.WithIterFunction("g", 2, 2, func(_ any, xs []any) gojq.Iter {
+		gojq.WithIterFunction("g", 2, 2, func(_ context.Context, _ any, xs []any) gojq.Iter {
 			if x, ok := xs[0].(int); ok {
 				if y, ok := xs[1].(int); ok {
 					return &rangeIter{x, y}
@@ -616,7 +617,7 @@ func TestWithIterFunction(t *testing.T) {
 			}
 			return gojq.NewIter(fmt.Errorf("g cannot be applied to: %v", xs))
 		}),
-		gojq.WithIterFunction("h", 0, 0, func(any, []any) gojq.Iter {
+		gojq.WithIterFunction("h", 0, 0, func(context.Context, any, []any) gojq.Iter {
 			return gojq.NewIter()
 		}),
 	)
@@ -646,7 +647,7 @@ func TestWithIterFunctionError(t *testing.T) {
 		t.Fatal(err)
 	}
 	code, err := gojq.Compile(query,
-		gojq.WithIterFunction("f", 0, 0, func(any, []any) gojq.Iter {
+		gojq.WithIterFunction("f", 0, 0, func(context.Context, any, []any) gojq.Iter {
 			return gojq.NewIter(1, errors.New("error"), 3)
 		}),
 	)
@@ -687,7 +688,7 @@ func TestWithIterFunctionPathIndexing(t *testing.T) {
 		t.Fatal(err)
 	}
 	code, err := gojq.Compile(query,
-		gojq.WithIterFunction("f", 0, 0, func(any, []any) gojq.Iter {
+		gojq.WithIterFunction("f", 0, 0, func(context.Context, any, []any) gojq.Iter {
 			return gojq.NewIter(0, 1, 2)
 		}),
 	)
@@ -712,7 +713,7 @@ func TestWithIterFunctionPathInputValue(t *testing.T) {
 		t.Fatal(err)
 	}
 	code, err := gojq.Compile(query,
-		gojq.WithIterFunction("f", 0, 0, func(v any, _ []any) gojq.Iter {
+		gojq.WithIterFunction("f", 0, 0, func(_ context.Context, v any, _ []any) gojq.Iter {
 			return gojq.NewIter(v, v, v)
 		}),
 	)
@@ -737,7 +738,7 @@ func TestWithIterFunctionPathEmpty(t *testing.T) {
 		t.Fatal(err)
 	}
 	code, err := gojq.Compile(query,
-		gojq.WithIterFunction("f", 0, 0, func(any, []any) gojq.Iter {
+		gojq.WithIterFunction("f", 0, 0, func(context.Context, any, []any) gojq.Iter {
 			return gojq.NewIter()
 		}),
 	)
@@ -762,7 +763,7 @@ func TestWithIterFunctionInvalidPathError(t *testing.T) {
 		t.Fatal(err)
 	}
 	code, err := gojq.Compile(query,
-		gojq.WithIterFunction("f", 0, 0, func(any, []any) gojq.Iter {
+		gojq.WithIterFunction("f", 0, 0, func(context.Context, any, []any) gojq.Iter {
 			return gojq.NewIter(map[string]any{"x": 1})
 		}),
 	)
@@ -789,7 +790,7 @@ func TestWithIterFunctionPathError(t *testing.T) {
 		t.Fatal(err)
 	}
 	code, err := gojq.Compile(query,
-		gojq.WithIterFunction("f", 0, 0, func(any, []any) gojq.Iter {
+		gojq.WithIterFunction("f", 0, 0, func(context.Context, any, []any) gojq.Iter {
 			return gojq.NewIter(errors.New("error"))
 		}),
 	)
@@ -822,10 +823,10 @@ func TestWithIterFunctionDefineError(t *testing.T) {
 		}
 	}()
 	t.Fatal(gojq.Compile(query,
-		gojq.WithFunction("f", 0, 0, func(any, []any) any {
+		gojq.WithFunction("f", 0, 0, func(context.Context, any, []any) any {
 			return 0
 		}),
-		gojq.WithIterFunction("f", 0, 0, func(any, []any) gojq.Iter {
+		gojq.WithIterFunction("f", 0, 0, func(context.Context, any, []any) gojq.Iter {
 			return gojq.NewIter()
 		}),
 	))
@@ -848,7 +849,7 @@ func TestWithFunctionWithModuleLoader(t *testing.T) {
 		t.Fatal(err)
 	}
 	code, err := gojq.Compile(query,
-		gojq.WithFunction("f", 0, 0, func(x any, _ []any) any {
+		gojq.WithFunction("f", 0, 0, func(_ context.Context, x any, _ []any) any {
 			if x, ok := x.(int); ok {
 				return x * 2
 			}