macros that cleanly Injecting code into a function's AST in julia

cleaninject.jl

#cleaninject.jl

#cleanly injects some code into the AST of a function.

function code_to_inject()
  println("this code is injected")
end

function code_to_inject(a,b)
  println("injected code handles $a and $b")
end

macro inject_code_prepend(f)
  #make sure this macro precedes a function definition.
  isa(f, Expr) || error("checkable macro must precede a function definition")
  (f.head == :function) || error("checkable macro must precede a function definition")

  #be lazy and let the parser do the hard work.
  b2 = parse("code_to_inject()")

  #inject the generated code into the AST.
  unshift!(f.args[2].args, b2)
  #return the escaped function to the parser so that it generates the new function.
  return Expr(:escape, f)
end

macro inject_code_append(f)
  #make sure this macro precedes a function definition.
  isa(f, Expr) || error("checkable macro must precede a function definition")
  (f.head == :function) || error("checkable macro must precede a function definition")

  #be lazy and let the parser do the hard work.
  b2 = parse("code_to_inject()")

  #inject the generated code into the AST.
  push!(f.args[2].args, b2)
  #return the escaped function to the parser so that it generates the new function.
  return Expr(:escape, f)
end

macro inject_code_with_args(f)
  #make sure this macro precedes a function definition.
  isa(f, Expr) || error("checkable macro must precede a function definition")
  (f.head == :function) || error("checkable macro must precede a function definition")

  #be lazy and let the parser do the hard work.
  b2 = parse(string("code_to_inject(", join(f.args[1].args[2:end], ","), ")"))

  #inject the generated code into the AST.
  unshift!(f.args[2].args, b2)
  #return the escaped function to the parser so that it generates the new function.
  return Expr(:escape, f)
end

################################################################################
# RESULTS

#=

julia> @inject_code_prepend function p()
       println("victim function")
       end
p (generic function with 1 method)

julia> p()
this code is injected
victim function

julia> @inject_code_append function p()
       println("victim function")
       end
p (generic function with 1 method)

julia> p()
victim function
this code is injected

julia> @inject_code_with_args function p(a, b)
       println("victim called with $a and $b")
       end
p (generic function with 2 methods)

julia> p(1, 2)
injected code handles 1 and 2
victim called with 1 and 2

=#