From 51c878b0c80316111842d53541ff20488f25debc Mon Sep 17 00:00:00 2001 From: Ville Vesilehto Date: Tue, 16 Dec 2025 12:55:41 +0200 Subject: [PATCH] fix(patcher): ctx into nested custom funcs and Env WithContext patcher now looks up function types from the Functions table and Env methods when the callee type is interface{}. This fixes context injection for custom functions and Env methods nested as arguments inside method calls with unknown callee types (e.g., Now2().After(Date2())). Also improved the regression test to actually verify context is passed to both functions, and added a test for Env methods. Signed-off-by: Ville Vesilehto --- expr.go | 11 ++++++-- patcher/with_context.go | 25 ++++++++++++++++- test/issues/823/issue_test.go | 53 ++++++++++++++++++++++++++++++++--- 3 files changed, 81 insertions(+), 8 deletions(-) diff --git a/expr.go b/expr.go index 280605e67..76fbd426f 100644 --- a/expr.go +++ b/expr.go @@ -195,9 +195,14 @@ 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, - }) + return func(c *conf.Config) { + c.Visitors = append(c.Visitors, patcher.WithContext{ + Name: name, + Functions: c.Functions, + Env: &c.Env, + NtCache: &c.NtCache, + }) + } } // Timezone sets default timezone for date() and now() builtin functions. diff --git a/patcher/with_context.go b/patcher/with_context.go index f9861a2c2..2043041ba 100644 --- a/patcher/with_context.go +++ b/patcher/with_context.go @@ -4,11 +4,16 @@ import ( "reflect" "github.com/expr-lang/expr/ast" + "github.com/expr-lang/expr/checker/nature" + "github.com/expr-lang/expr/conf" ) // WithContext adds WithContext.Name argument to all functions calls with a context.Context argument. type WithContext struct { - Name string + Name string + Functions conf.FunctionsTable // Optional: used to look up function types when callee type is unknown. + Env *nature.Nature // Optional: used to look up method types when callee type is unknown. + NtCache *nature.Cache // Optional: cache for nature lookups. } // Visit adds WithContext.Name argument to all functions calls with a context.Context argument. @@ -19,6 +24,24 @@ func (w WithContext) Visit(node *ast.Node) { if fn == nil { return } + // If callee type is interface{} (unknown), look up the function type from + // the Functions table or Env. This handles cases where the checker returns early + // without visiting nested call arguments (e.g., Date2() in Now2().After(Date2())) + // because the outer call's type is unknown due to missing context arguments. + if fn.Kind() == reflect.Interface { + if ident, ok := call.Callee.(*ast.IdentifierNode); ok { + if w.Functions != nil { + if f, ok := w.Functions[ident.Value]; ok { + fn = f.Type() + } + } + if fn.Kind() == reflect.Interface && w.Env != nil { + if m, ok := w.Env.MethodByName(w.NtCache, ident.Value); ok { + fn = m.Type + } + } + } + } if fn.Kind() != reflect.Func { return } diff --git a/test/issues/823/issue_test.go b/test/issues/823/issue_test.go index 672cc1f1f..221267de4 100644 --- a/test/issues/823/issue_test.go +++ b/test/issues/823/issue_test.go @@ -2,7 +2,6 @@ package issue_test import ( "context" - "fmt" "testing" "time" @@ -14,26 +13,72 @@ type env struct { Ctx context.Context `expr:"ctx"` } +// TestIssue823 verifies that WithContext injects context into nested custom +// function calls. The bug was that date2() nested as an argument to After() +// didn't receive the context because its callee type was unknown. func TestIssue823(t *testing.T) { + now2Called := false + date2Called := false + p, err := expr.Compile( "now2().After(date2())", expr.Env(env{}), expr.WithContext("ctx"), expr.Function( "now2", - func(params ...any) (any, error) { return time.Now(), nil }, + func(params ...any) (any, error) { + require.Len(t, params, 1, "now2 should receive context") + _, ok := params[0].(context.Context) + require.True(t, ok, "now2 first param should be context.Context") + now2Called = true + return time.Now(), nil + }, new(func(context.Context) time.Time), ), expr.Function( "date2", - func(params ...any) (any, error) { return time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC), nil }, + func(params ...any) (any, error) { + require.Len(t, params, 1, "date2 should receive context") + _, ok := params[0].(context.Context) + require.True(t, ok, "date2 first param should be context.Context") + date2Called = true + return time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC), nil + }, new(func(context.Context) time.Time), ), ) - fmt.Printf("Compile result err: %v\n", err) require.NoError(t, err) r, err := expr.Run(p, &env{Ctx: context.Background()}) require.NoError(t, err) require.True(t, r.(bool)) + require.True(t, now2Called, "now2 should have been called") + require.True(t, date2Called, "date2 should have been called") +} + +// envWithMethods tests that Env methods with context.Context work correctly +// when nested in method chains (similar to TestIssue823 but with Env methods). +type envWithMethods struct { + Ctx context.Context `expr:"ctx"` +} + +func (e *envWithMethods) Now2(ctx context.Context) time.Time { + return time.Now() +} + +func (e *envWithMethods) Date2(ctx context.Context) time.Time { + return time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC) +} + +func TestIssue823_EnvMethods(t *testing.T) { + p, err := expr.Compile( + "Now2().After(Date2())", + expr.Env(&envWithMethods{}), + expr.WithContext("ctx"), + ) + require.NoError(t, err) + + r, err := expr.Run(p, &envWithMethods{Ctx: context.Background()}) + require.NoError(t, err) + require.True(t, r.(bool)) }