# by: Cody Kochmann
# this code works in both python 2 and 3
import ast
import inspect
def insert_line_into_function(fn, line):
''' this function essentially lets you safely add source code to the body
of a function without penalizing its speed with a decorator '''
# store fn's name for later
___fn_name = fn.__name__
# turn the function into an ast object
fn_ast = ast.parse(inspect.getsource(fn))
# insert the line's ast form into the functions ast body
fn_ast.body[0].body.insert(0, ast.parse(line).body[0])
# compile the new fn and load it into memory
exec(compile(fn_ast, '<string>', 'exec'))
# return the newly loaded function
return locals()[___fn_name]
# this is the function we are gonna modify with ast
def my_adder(a,b,c):
d = a+b
if c%2:
return a+b
else:
return a+d
# heres the code we're gonna inject into my_adder
code_to_inject = 'assert type(a) == int, "a needs to be an int"'
# currently my_adder accepts floats in the first argument
print(my_adder(1,2,3))
print(my_adder(1.0,2,3))
# this my_adder should now throw an assertion for non ints in the first arg
my_adder = insert_line_into_function(my_adder, code_to_inject)
print(my_adder(1,2,3))
print(my_adder(1.0,2,3)) # this line should now raise an AssertionError