Macro to Extend Function with Additional Arguements (MacroTools.jl)?

My goal is to create a macro that takes an arbitrary function and adds additional arguments and changes the docstring.

For example, say I define the following function:

"This is my function"
function f(a, b; x=1.0, y=1.0)
    return a + b + x + y
end

And now assume I’d also like to define the following augmented function (where the extra arguments c and d have been added:

"This is my augmented function"
function f(a, b, c, d; x=1.0, y=1.0)
    return f(a, b; x=x, y=y) + c + d
end

I would like to create a macro that can define an augmented function automatically. Ideally, I would just write:

@augment begin
    "This is my function"
    function f(a, b; x=1.0, y=1.0)
        return a + b + x + y
    end
end

and this would define the original function and augmented function to use.

So far I have tried using MacroTools.jl as follows:

using MacroTools

macro augment(expr)
    # Evaluate original function expression
    func = eval(expr)

    # Split function definition into dictionaries
    func_dict = MacroTools.splitdef(expr)
    new_dict = deepcopy(func_dict)

    # Change new function bodies
    new_dict[:body] = :($func($(func_dict[:args]...); $(func_dict[:kwargs]...)) + c + d)

    # Change new function Arguments
    push!(new_dict[:args], :c, :d)

    # Convert back to expressions and evaluate
    eval(MacroTools.combinedef(new_dict))
    return nothing
end

The above code successfully augments functions with c and d arguments, however, there are two problems. Firstly, the keyword arguments in the augmented function are always set to their default value, such that calling f(1, 2, 3, 4; x=5, y=6) acutally calls f(1, 2, 3, 4; x=1.0, y=1.0). Secondly, the @augment macro above does not handle the docstring conversion. As soon as I enter a doctring to the function, the MacroTools.splitdef(expr) fails.

This task is proving quite difficult to me so any help is appreciated.

No fix to your problem, just comments:

  • There is usually no need to use eval inside a macro. What eval does in your example is that it takes an Expr and “compiles it” when the macro is run. However, macros are mostly just meant as syntax transformations that happen before compilation.
  • As such macros should return Exprs which are then implicitly evaled in the module where the macro was called from. Note that you will need to wrap the returned expression into an esc(...) block to undo macro hygiene.
  • To debug macro calls you can use @macroexpand or @macroexpand1 to see what the result of the syntax transformation. Note that this only works on the returned Expr given by a macro and won’t help in your case as you return nothing. Use like
@macroexpand @augment begin
    "This is my function"
    function f(a, b; x=1.0, y=1.0)
        return a + b + x + y
    end
end

Can I ask what speaks against using an ordinary function here that would save you all of the headache with macros? Something like this could work

add_c_d(f, c, d) = f + c + d

or maybe

call_and_add_c_d(fn, args, kwargs, c, d) = fn(args...; kwargs...) + c + d
1 Like

Cheers, thanks. I’ve now changed my solution to return the expr’s instead.

I’m doing this just to learn macro functionality mainly.

I found a solution using Expronicon.jl:

using ExproniconLite

macro augment(expr)
    # Split function definition into dictionaries
    func_expr = JLFunction(expr)
    new_expr = deepcopy(func_expr)

    # Map the keyword values to the keyword symbols for help calling the function later
    call_expr = deepcopy(func_expr)
    map(x -> x.args[2] = x.args[1], call_expr.kwargs)

    # Change new function body
    new_expr.body = :(return $(call_expr.name)($(call_expr.args...); $(call_expr.kwargs...)) + c + d) # THE KW ARGS DON'T HAVE THE SYMBOL DEFINITION AT THE STARTK

    # Change the doc
    new_expr.doc = "This is my augmented function"

    # Change new function Arguments
    push!(new_expr.args, :c, :d)

    # Return both function expr's
    return quote
        $(esc(expr))
        $(esc(codegen_ast(new_expr)))
    end
end