CodyKochmann
12/27/2017 - 1:33 PM

use this example to learn how to modify the body of a function by using ast

use this example to learn how to modify the body of a function by using ast

# 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