@replacerand macro to replace calls to rand() by a given function e.g. x->0.5

ReplaceRand.jl

"""
A macro `@replacerand` that replaces the value returned by a call to `rand` with something else.
"""

using Cassette

Cassette.@context ReplaceRand

Cassette.overdub(ctx::ReplaceRand, fn::typeof(rand), args...) = ctx.metadata(args...)

replacerand(f, replacement=(x...) -> 0.5) = Cassette.overdub(ReplaceRand(metadata=replacement), f)

macro replacerand(replacement, ex)
  :(replacerand(() -> $(esc(ex)), $(esc(replacement))))
end

# define what rand should be replaced by
replacement() = 0.5
replacement(n) = zeros(n) .+ 0.5
replacement(_, n) = zeros(n) .+ 0.5

@replacerand replacement rand()

function foo()
    return rand()
end

@replacerand replacement foo()

module Bar
    baz() = rand()
end

using .Bar

@replacerand replacement Bar.baz()

@replacerand replacement rand(2)
@replacerand replacement rand((2, 3))

using Random
rng = Random.MersenneTwister()

@replacerand replacement rand(rng, 2)
@replacerand replacement rand(rng, (2, 3))

Brilliant - thanks very much @femtomc