Add expr.WithContext()

expr.go

@@ -11,6 +11,7 @@ import (
 	"github.com/antonmedv/expr/file"
 	"github.com/antonmedv/expr/optimizer"
 	"github.com/antonmedv/expr/parser"
+	"github.com/antonmedv/expr/patcher"
 	"github.com/antonmedv/expr/vm"
 )
 
@@ -152,6 +153,13 @@ func EnableBuiltin(name string) Option {
 	}
 }
 
+// WithContext passes context to all functions calls with a context.Context argument.
+func WithContext(name string) Option {
+	return Patch(patcher.WithContext{
+		Name: name,
+	})
+}
+
 // Compile parses and compiles given input expression to bytecode program.
 func Compile(input string, ops ...Option) (*vm.Program, error) {
 	config := conf.CreateNew()

expr_test.go

@@ -1,6 +1,7 @@
 package expr_test
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"reflect"
@@ -488,6 +489,47 @@ func ExamplePatch() {
 	// Output : Hello, you, world!
 }
 
+func ExampleWithContext() {
+	env := map[string]any{
+		"fn": func(ctx context.Context, _, _ int) int {
+			// An infinite loop that can be canceled by context.
+			for {
+				select {
+				case <-ctx.Done():
+					return 42
+				}
+			}
+		},
+		"ctx": context.TODO(), // Context should be passed as a variable.
+	}
+
+	program, err := expr.Compile(`fn(1, 2)`,
+		expr.Env(env),
+		expr.WithContext("ctx"), // Pass context variable name.
+	)
+	if err != nil {
+		fmt.Printf("%v", err)
+		return
+	}
+
+	// Cancel context after 100 milliseconds.
+	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*100)
+	defer cancel()
+
+	// After program is compiled, context can be passed to Run.
+	env["ctx"] = ctx
+
+	// Run will return 42 after 100 milliseconds.
+	output, err := expr.Run(program, env)
+	if err != nil {
+		fmt.Printf("%v", err)
+		return
+	}
+
+	fmt.Printf("%v", output)
+	// Output: 42
+}
+
 func TestExpr_readme_example(t *testing.T) {
 	env := map[string]any{
 		"greet":   "Hello, %v!",

patcher/with_context.go

@@ -0,0 +1,36 @@
+package patcher
+
+import (
+	"reflect"
+
+	"github.com/antonmedv/expr/ast"
+)
+
+// WithContext adds WithContext.Name argument to all functions calls with a context.Context argument.
+type WithContext struct {
+	Name string
+}
+
+// Visit adds WithContext.Name argument to all functions calls with a context.Context argument.
+func (w WithContext) Visit(node *ast.Node) {
+	switch (*node).(type) {
+	case *ast.CallNode:
+		call := (*node).(*ast.CallNode)
+		fn := call.Callee.Type()
+		if fn.Kind() != reflect.Func {
+			return
+		}
+		if fn.NumIn() == 0 {
+			return
+		}
+		if fn.In(0).String() != "context.Context" {
+			return
+		}
+		ast.Patch(node, &ast.CallNode{
+			Callee: call.Callee,
+			Arguments: append([]ast.Node{
+				&ast.IdentifierNode{Value: w.Name},
+			}, call.Arguments...),
+		})
+	}
+}

patcher/with_context_test.go

@@ -0,0 +1,32 @@
+package patcher_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/antonmedv/expr"
+	"github.com/antonmedv/expr/patcher"
+)
+
+func TestWithContext(t *testing.T) {
+	env := map[string]any{
+		"fn": func(ctx context.Context, a int) int {
+			return ctx.Value("value").(int) + a
+		},
+		"ctx": context.TODO(),
+	}
+
+	withContext := patcher.WithContext{Name: "ctx"}
+
+	program, err := expr.Compile(`fn(40)`, expr.Env(env), expr.Patch(withContext))
+	require.NoError(t, err)
+
+	ctx := context.WithValue(context.Background(), "value", 2)
+	env["ctx"] = ctx
+
+	output, err := expr.Run(program, env)
+	require.NoError(t, err)
+	require.Equal(t, 42, output)
+}