Make WithContext work for methods on env struct (#602)

This makes the WithContext patcher work when passing in a struct for an
env that has methods that take a context.

This is a bit fiddly because we can't quite detect the difference of
function vs. method on a struct, so we have to check both the first and
the second param. Since it's highly unusual to pass the context as
anything other than the first parameter, this should work out just fine
in practice.

Fixes #600.

Author: daenney

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,

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
 }