Repository: expr-lang/expr (Go)

Issue:
WithContext fails on env struct with method resulting in "not enough arguments to call"
When using a struct to pass the environment and using a method on that struct which takes a `context.Context` as `ctx` on the first argument, calling that function in expr fails despite the use of the `withContext` patcher:

```go
type testEnvContext struct {
	Context context.Context `expr:"ctx"`
}

func (testEnvContext) Fn(ctx context.Context, a int) int {
	return ctx.Value("value").(int) + a
}

func TestWithContext_env_struct(t *testing.T) {
	withContext := patcher.WithContext{Name: "ctx"}

	program, err := expr.Compile(`Fn(40)`, expr.Env(testEnvContext{}), expr.Patch(withContext))
	require.NoError(t, err)

	ctx := context.WithValue(context.Background(), "value", 2)
	env := testEnvContext{
		Context: ctx,
	}

	output, err := expr.Run(program, env)
	require.NoError(t, err)
	require.Equal(t, 42, output)
}
```

```
--- FAIL: TestWithContext_env_struct (0.00s)
    /expr/patcher/with_context_test.go:77: 
        	Error Trace:	/expr/patcher/with_context_test.go:77
        	Error:      	Received unexpected error:
        	            	not enough arguments to call Fn (1:1)
        	            	 | Fn(40)
        	            	 | ^
        	Test:       	TestWithContext_env_struct
FAIL
FAIL	github.com/expr-lang/expr/patcher	0.004s
FAIL
```

Using the example in the docs with the custom visitor does work https://github.com/expr-lang/expr/blob/85b178c3dfa2ded602d1319f13729f5c5041345a/docs/visitor-and-patch.md#L86-L97


Discussion:
Looks like a bug.
I think the issue is here: https://github.com/expr-lang/expr/blob/85b178c3dfa2ded602d1319f13729f5c5041345a/patcher/with_context.go#L28

If I add a Println for it, I get: `patcher_test.testEnvContext` In the case of the method on a struct, it needs `fn.In(1).String()`, then I get `context.Context`.
Yes, looks like this is the culprit. 
If you happen to know how to resolve the difference I'm happy to send a PR. Just not sure how to figure it out since `reflect.Kind()` is `reflect.Func` in both cases.
Right now we can just check first and second IN arguments. 

Later I’m planning to implement more advance type system #563 called Nature 🙃 



Patch:
diff --git a/patcher/with_context.go b/patcher/with_context.go
index 55b604261..f9861a2c2 100644
--- a/patcher/with_context.go
+++ b/patcher/with_context.go
@@ -22,11 +22,18 @@ func (w WithContext) Visit(node *ast.Node) {
 		if fn.Kind() != reflect.Func {
 			return
 		}
-		if fn.NumIn() == 0 {
-			return
-		}
-		if fn.In(0).String() != "context.Context" {
+		switch fn.NumIn() {
+		case 0:
 			return
+		case 1:
+			if fn.In(0).String() != "context.Context" {
+				return
+			}
+		default:
+			if fn.In(0).String() != "context.Context" &&
+				fn.In(1).String() != "context.Context" {
+				return
+			}
 		}
 		ast.Patch(node, &ast.CallNode{
 			Callee: call.Callee,


Test patch:
diff --git a/patcher/with_context_test.go b/patcher/with_context_test.go
index afad4e6f0..5ce64191f 100644
--- a/patcher/with_context_test.go
+++ b/patcher/with_context_test.go
@@ -62,6 +62,30 @@ func TestWithContext_with_env_Function(t *testing.T) {
 	require.Equal(t, 42, output)
 }
 
+type testEnvContext struct {
+	Context context.Context `expr:"ctx"`
+}
+
+func (testEnvContext) Fn(ctx context.Context, a int) int {
+	return ctx.Value("value").(int) + a
+}
+
+func TestWithContext_env_struct(t *testing.T) {
+	withContext := patcher.WithContext{Name: "ctx"}
+
+	program, err := expr.Compile(`Fn(40)`, expr.Env(testEnvContext{}), expr.Patch(withContext))
+	require.NoError(t, err)
+
+	ctx := context.WithValue(context.Background(), "value", 2)
+	env := testEnvContext{
+		Context: ctx,
+	}
+
+	output, err := expr.Run(program, env)
+	require.NoError(t, err)
+	require.Equal(t, 42, output)
+}
+
 type TestFoo struct {
 	contextValue int
 }