Support keyword arguments in @formula function calls

src/StatsModels.jl

@@ -51,6 +51,9 @@ export
     FunctionTerm,
     MatrixTerm,
 
+    kwarg_exprs,
+    has_kwargs,
+
     lag, lead, # Reexported from ShiftedArrays
 
     term,

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)

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
 

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