Instrumenting Python Code

Question

I want to instrument every if condition encountered inside a python function automatically. For example if my code hits:

if x > 2:
    #do something

I want to record/ print "x > 2" to the console.

I found this tool "equip" https://github.com/neuroo/equip/ but it only instuments at function boundaries. I think I need to be able to instrument at the bytecode instruction level to get those conditionals out.

Currently, I have to manually stdout/append after every if condition like

if x > 2:    
    list.append("x>2")

Answer 1

This is a simple example to get you started. Suppose you have a module like this:

def foo(x):
    if x > 100:
        print('big')
    else:
        print('small')


if __name__ == '__main__':
    foo(5)
    foo(500)

If you exec foo, you get:

small
big

Now you want to print the test clause of every if statement if the test is True. Lets start by importing foo:

>>> import foo

Then get the source code:

>>> source = inspect.getsource(foo)

Lets parse the source in order to get an abstract syntax tree:

>>> tree = ast.parse(source)

The next step is defining a NodeTransformer that will modify the tree:

>>> class IfTransformer(ast.NodeTransformer):
         def visit_If(self, node):
             new_node = ast.Expr(value=ast.Call(
                  func=ast.Name(id='print', ctx=ast.Load()), 
                  args=[ast.Str(s=astunparse.unparse(node.test))], 
                  keywords=[]
             ))
             node.body.insert(0, new_node)
             return ast.fix_missing_locations(node)

In order to modify the tree we make our IfTransformer visit all nodes:

>>> IfTransformer().visit(tree)

Then you can compile and execute your new source:

>>> exec(compile(tree, 'foo.py', 'exec'))
(__name__ == '__main__')

small
(x > 100)

big

For each if clause where the test is True, you got the test printed. A smart guy like you can figure everything out from here.

Check out this video from Pycon 2011.