Create Formulas with Macros

Hi all,

I’m practicing working with macros and ran into an idea that I want to attempt to implement. Essentially, I want to create a macro @form that takes a formula expression. The expression would then be applied to a named tuple and return a new named tuple based on the formula values. Without a macro, I can achieve this with the following:

# create named tuple
nt = (; a =1, b=2)

# function of interest 
formulafun(nt) =(:a, nt[:a] + nt[:b] +1)

function foo(nt, formulafun)

    operation = formulafun(nt)
    
     # return new named tuple with key = :a and value = 3
    updated_val = NamedTuple{(operation[1],)}((operation[2],))
   
    # merge with original tuple to 'update' the ':a' key
    return merge(nt, updated_val)
end;


#foo(nt, formulafun)

I’d like to mimic foo with something like @form :a ~ :a + :b +1 (inspired by StatsModels.jl) but I’m definitely stuck. Has anyone tried something similar? I’d appreciate any help!

macro form(nt, exp)
    quote
        exp.args[3]
    end 
end

@form :a ~ :a + :b +1
1 Like

With Macrotools.jl

using MacroTools

macro form(nt, expr)
    @capture(expr, :lhs_ ~ rhs_)
    new_rhs = MacroTools.postwalk(rhs) do subex
        if @capture(subex, :s_)
            return :($nt[$(QuoteNode(s))])
        else
            return subex
        end
    end
    return :( merge($nt, (; $lhs = $new_rhs)) )
end

val = (; a =1, b=2)

@form(val, :a ~ :a + :b + 1) # returns (a = 4, b = 2)
2 Likes

Wow, thanks a lot!

After referencing the Macrotools.jl docs, I have some remaining questions regarding the following lines:

new_rhs = MacroTools.postwalk(rhs) do subex
        if @capture(subex, :s_)
            return :($nt[$(QuoteNode(s))])
        else
            return subex
        end
    end
end

Is the idea that for each element of the rhs subexpression, i.e. :(:a + :b + 1), if the subexpression is a symbol then return the value associated with each key associated value, otherwise, return the element? If this is the case, does the do block syntax then reduce the subexpression to :(1 + 2 + 1) which is then evaluated via $new_rhs?

Not exactly: if the subexpression represents a symbol, then it returns an expression of the form nt[:s], where nt is the first argument of the macro (purportedly the named tuple), and s is the symbol. So if the right hand side is :(:a + :b + 1), and the first argument was val, the whole expression will be replaced by :(val[:a] + val[:b] + 1).

You can see the result with @macroexpand @form(val, :a ~ :a + :b + 1)

This seems just a minor detail, but it is important, because the macro does not “know” what expressions that are passed as arguments “mean”, it just takes them and creates another expression. It’s only when it is evaluated afterwards that val[:a] is replaced by 1, etc.

By the way, the macro looks nicer if the line

return :($nt[$(QuoteNode(s))])

is replaced by

return :($nt.$s)

Got it! Thank you so much for taking the time to help me through this.

1 Like