From a3b58326b226ce08ad0d02ce46dc71649e432ede Mon Sep 17 00:00:00 2001
From: Simon Frost <sdwfrost@microsoft.com>
Date: Tue, 24 Mar 2026 07:00:08 -0700
Subject: [PATCH] Support keyword arguments in @formula function calls
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add parse! methods for :kw and :parameters Expr nodes so that
keyword arguments in function calls pass through instead of
erroring. This enables syntax like:

    @formula(y ~ s(x; k=10, bs=:cr))
    @formula(y ~ s(x, k=10, bs=:cr))

Kwargs are preserved in FunctionTerm.exorig and can be extracted
by downstream packages via the new kwarg_exprs() and has_kwargs()
helper functions.

This is fully backward compatible — existing formulas without
kwargs work identically. The 9 pre-existing broken tests are
unchanged.

Motivation: packages like GAM.jl need to pass configuration
options (basis type, dimension, etc.) to smooth term constructors
within formulas. Without this change, they must define custom
formula macros (@gam_formula) instead of using @formula.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 src/StatsModels.jl |  3 +++
 src/formula.jl     | 12 ++++++++++++
 src/terms.jl       | 40 +++++++++++++++++++++++++++++++++++++++
 test/formula.jl    | 47 ++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 102 insertions(+)

diff --git a/src/StatsModels.jl b/src/StatsModels.jl
index 79996c21..8a7effc2 100644
--- a/src/StatsModels.jl
+++ b/src/StatsModels.jl
@@ -51,6 +51,9 @@ export
     FunctionTerm,
     MatrixTerm,
 
+    kwarg_exprs,
+    has_kwargs,
+
     lag, lead, # Reexported from ShiftedArrays
 
     term,
diff --git a/src/formula.jl b/src/formula.jl
index 12dd8f64..2884b2b7 100644
--- a/src/formula.jl
+++ b/src/formula.jl
@@ -65,6 +65,17 @@ end
 
 function parse!(ex::Expr, protected::Bool=false)
     catch_dollar(ex)
+
+    # Handle keyword argument nodes that appear inside function calls.
+    # These are generated by Julia's parser for syntax like f(x; k=10) or f(x, k=10).
+    # We quote them so they pass through as literal values in FunctionTerm.args,
+    # allowing downstream packages to extract kwargs from FunctionTerm.exorig.
+    if Meta.isexpr(ex, :kw)
+        return Meta.quot(ex)
+    elseif Meta.isexpr(ex, :parameters)
+        return Meta.quot(ex)
+    end
+
     check_call(ex)
 
     if ex.args[1] ∈ SPECIALS && !protected
@@ -88,3 +99,4 @@ end
 parse!(::Nothing, protected) = :(nothing)
 parse!(s::Symbol, protected) = :(Term($(Meta.quot(s))))
 parse!(n::Number, protected) = :(ConstantTerm($n))
+parse!(q::QuoteNode, protected) = Meta.quot(q)
diff --git a/src/terms.jl b/src/terms.jl
index 3617e28b..ac1effc2 100644
--- a/src/terms.jl
+++ b/src/terms.jl
@@ -120,6 +120,46 @@ width(::FunctionTerm) = 1
 
 Base.:(==)(a::FunctionTerm, b::FunctionTerm) = a.f == b.f && a.args == b.args && a.exorig == b.exorig
 
+"""
+    kwarg_exprs(ft::FunctionTerm)
+
+Extract keyword argument expressions from a `FunctionTerm`'s original expression.
+Returns a vector of `Expr(:kw, name, value)` nodes, or an empty vector if there
+are no keyword arguments.
+
+This is useful for downstream packages that extend `@formula` syntax with
+keyword arguments in function calls (e.g., `@formula(y ~ s(x; k=10, bs=:cr))`).
+
+# Example
+
+```julia
+f = @formula(y ~ foo(x; a=1, b=2))
+ft = f.rhs
+kwarg_exprs(ft)  # [:(a = 1), :(b = 2)]
+```
+"""
+function kwarg_exprs(ft::FunctionTerm)
+    kws = Expr[]
+    for arg in ft.exorig.args[2:end]
+        if arg isa Expr && Meta.isexpr(arg, :parameters)
+            for kw in arg.args
+                kw isa Expr && Meta.isexpr(kw, :kw) && push!(kws, kw)
+            end
+        elseif arg isa Expr && Meta.isexpr(arg, :kw)
+            push!(kws, arg)
+        end
+    end
+    return kws
+end
+
+"""
+    has_kwargs(ft::FunctionTerm)
+
+Return `true` if the `FunctionTerm` was constructed from an expression containing
+keyword arguments.
+"""
+has_kwargs(ft::FunctionTerm) = !isempty(kwarg_exprs(ft))
+
 """
     InteractionTerm{Ts} <: AbstractTerm
 
diff --git a/test/formula.jl b/test/formula.jl
index 4e1a4b93..bb76dad7 100644
--- a/test/formula.jl
+++ b/test/formula.jl
@@ -138,4 +138,51 @@
     f = @formula(foo ~ bar)
     @test f == deepcopy(f)
 
+    # Keyword arguments in function calls
+    @testset "function call kwargs" begin
+        myf(args...; kwargs...) = sum(args)
+
+        # Semicolon kwargs syntax
+        f = @formula(y ~ myf(x; k=10, bs=:cr))
+        ft = f.rhs
+        @test ft isa FunctionTerm
+        @test has_kwargs(ft)
+        kws = kwarg_exprs(ft)
+        @test length(kws) == 2
+        @test kws[1].args[1] == :k
+        @test kws[1].args[2] == 10
+        @test kws[2].args[1] == :bs
+
+        # Comma kwargs syntax
+        f2 = @formula(y ~ myf(x, k=10, bs=:cr))
+        @test has_kwargs(f2.rhs)
+        @test length(kwarg_exprs(f2.rhs)) == 2
+
+        # No kwargs — backward compat
+        f3 = @formula(y ~ myf(x, 10))
+        @test !has_kwargs(f3.rhs)
+        @test isempty(kwarg_exprs(f3.rhs))
+
+        # Mixed positional + kwargs
+        f4 = @formula(y ~ myf(x, 10; bs=:cr))
+        @test has_kwargs(f4.rhs)
+        kws4 = kwarg_exprs(f4.rhs)
+        @test length(kws4) == 1
+        @test kws4[1].args[1] == :bs
+
+        # Multiple terms, some with kwargs
+        f5 = @formula(y ~ myf(x; k=10) + log(z) + w)
+        for t in f5.rhs
+            if t isa FunctionTerm && t.f === myf
+                @test has_kwargs(t)
+            elseif t isa FunctionTerm && t.f === log
+                @test !has_kwargs(t)
+            end
+        end
+
+        # Standard formulas still work
+        f6 = @formula(y ~ 1 + x * z)
+        @test f6 isa FormulaTerm
+    end
+
 end