Macro to automatically add method with default_rng() for functions using RNGs

General Usage
,
1

Does anybody have a macro so I can write something like

@with_default_rng function my_func(rng::AbstractRNG, x, y, z)
    out =  # do something with rng, x, y, z
    return out
end

and it will automatically add the method

my_func(x, y, z) = my_func(default_rng(), x, y, z)

because it’s super inconvenient to have to add extra methods for every function using randomness to use the default_rng().

2

Or maybe a more general macro that lets you define a single default argument in any position.
So that

@default_arg foo(x = a, y, z) = ...

expands into

foo(y, z) = foo(a, y, z)
foo(x, y, z) = ...

macro default_arg(ex)
    (ex.head === :function || ex.head === :(=)) || throw(ArgumentError("Invalid expression"))
    call = ex.args[1]
    call1 = copy(call)
    call2 = copy(call1)
    Nargs = length(call.args)
    for i in 2:Nargs
         arg = call.args[i]
         if Meta.isexpr(arg, :kw, 2)
             call.args[i] = arg.args[1]
             deleteat!(call1.args, i)
             call2.args[i] = arg.args[2]
             break
        end
    end
    q = quote
        $ex
        $call1 = $call2
    end
    esc(q)
end

Yields:


julia> @macroexpand @default_arg function my_func(rng = default_rng(), x, y, z)
           out = rand(rng)*x + y/z
       end
quote
    #= REPL[33]:17 =#
    function my_func(rng, x, y, z)
        #= REPL[39]:1 =#
        #= REPL[39]:2 =#
        out = rand(rng) * x + y / z
    end
    #= REPL[33]:18 =#
    my_func(x, y, z) = my_func(default_rng(), x, y, z)
end

julia> @default_arg function my_func(rng = default_rng(), x, y, z)
           out = rand(rng)*x + y/z
       end
my_func (generic function with 2 methods)

julia> my_func(1.2, 2.3, 4.5)
1.0098752663124524
3 Likes