Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Question: Possible to do y ~ 1 + poly(x,deg=5) #262

Open
behinger opened this issue Aug 11, 2022 · 2 comments
Open

Question: Possible to do y ~ 1 + poly(x,deg=5) #262

behinger opened this issue Aug 11, 2022 · 2 comments

Comments

@behinger
Copy link

Hi!
hope I didn't miss it in the docs somewhere, is it possible to extend the formula syntax with keyword-value pairs?
E.g. poly(x,deg=5) instead of the boring poly(x,5)

Cheers, Benedikt

@ararslan
Copy link
Member

It's not currently possible but

diff --git a/src/formula.jl b/src/formula.jl
index 6cab6e1..6d2ef57 100644
--- a/src/formula.jl
+++ b/src/formula.jl
@@ -188,13 +188,19 @@ function parse!(ex::Expr, rewrites::Vector)
 
     # parse a copy of non-special calls
     ex_parsed = ex.args[1] ∉ SPECIALS ? deepcopy(ex) : ex
-    
+
     # iterate over children, checking for special rules
     child_idx = 2
     while child_idx <= length(ex_parsed.args)
-        @debug "  ($(ex_parsed.args[1])) i=$child_idx: $(ex_parsed.args[child_idx])"
+        child = ex_parsed.args[child_idx]
+        @debug "  ($(ex_parsed.args[1])) i=$child_idx: $child"
+        if Meta.isexpr(child, :parameters) || Meta.isexpr(child, :kw)
+            @debug "  not descending into keywords"
+            child_idx += 1
+            continue
+        end
         # depth first: parse each child first
-        parse!(ex_parsed.args[child_idx], rewrites)
+        parse!(child, rewrites)
         # find first rewrite rule that applies
         rule = filterfirst(r->applies(ex_parsed, child_idx, r), rewrites)
         # re-write according to that rule and update the child to position rewrite nex_parsedt
@@ -228,7 +234,7 @@ function capture_call_ex!(ex::Expr, ex_parsed::Expr)
                f_anon_ex,
                tuple(symbols...),
                Meta.quot(deepcopy(ex)),
-               :[$(ex_parsed.args[2:end]...)]]
+               :[$(map(e -> e isa Expr ? Meta.quot(e) : e, ex_parsed.args[2:end])...)]]
     return ex
 end
 
@@ -244,6 +250,10 @@ function terms!(ex::Expr)
     elseif is_call(ex, :capture_call)
         # final argument of capture_call holds parsed terms
         ex.args[end].args .= terms!.(ex.args[end].args)
+    elseif Meta.isexpr(ex, :parameters)
+        ex.args[2:end] .= terms!.(ex.args[2:end])
+    elseif Meta.isexpr(ex, :kw)
+        ex.args[2] = terms!(ex.args[2])
     end
     return ex
 end
diff --git a/src/terms.jl b/src/terms.jl
index b49a902..4b5df99 100644
--- a/src/terms.jl
+++ b/src/terms.jl
@@ -327,8 +327,19 @@ capture_call(args...) = FunctionTerm(args...)
 
 extract_symbols(x) = Symbol[]
 extract_symbols(x::Symbol) = [x]
-extract_symbols(ex::Expr) =
-    is_call(ex) ? mapreduce(extract_symbols, union, ex.args[2:end]) : Symbol[]
+function extract_symbols(ex::Expr)
+    if is_call(ex)
+        mapreduce(extract_symbols, union, ex.args[2:end])
+    elseif Meta.isexpr(ex, :parameters, 1)  # `f(; x)`
+        extract_symbols(ex.args[1])
+    elseif Meta.isexpr(ex, :parameters)  # `f(; x=a)`
+        mapreduce(extract_symbols, union, ex.args[1].args)
+    elseif Meta.isexpr(ex, :kw)  # `f(x=a)` or `f(; x=a)`
+        extract_symbols(last(ex.args))
+    else
+        Symbol[]
+    end
+end
 
 ################################################################################
 # showing terms

will get you part of the way there if you wanted to have a go at implementing it.

@kleinschmidt
Copy link
Member

I think with #183, it would be more straightforward to capture kwargs as well (but would require adding a separate field to FunctionTerm to store the kwargs

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

3 participants