Examples of working with ASTs

Working versions of these examples are in the examples directory of the source repository.

Wrapping integers

In Python code, 1/3 would normally be evaluated to a floating-point number, that can never be exactly one third. Mathematical software, like SymPy or Sage, often wants to use exact fractions instead. One way to make 1/3 produce an exact fraction is to wrap the integer literals 1 and 3 in a class:

class IntegerWrapper(ast.NodeTransformer):
    """Wraps all integers in a call to Integer()"""
    def visit_Num(self, node):
        if isinstance(node.n, int):
            return ast.Call(func=ast.Name(id='Integer', ctx=ast.Load()),
                            args=[node], keywords=[])
        return node

tree = ast.parse("1/3")
tree = IntegerWrapper().visit(tree)
# Add lineno & col_offset to the nodes we created
ast.fix_missing_locations(tree)

# The tree is now equivalent to Integer(1)/Integer(3)
# We would also need to define the Integer class and its __truediv__ method.

See wrap_integers.py for a working demonstration.

Simple test framework

These two manipulations let you write test scripts as a simple series of assert statements. First, we need to run the statements one by one, so execution doesn’t stop at the first test failure:

tree = ast.parse(code)
lines = [None] + code.splitlines()  # None at [0] so we can index lines from 1
test_namespace = {}

for node in tree.body:
    wrapper = ast.Module(body=[node])
    try:
        co = compile(wrapper, "<ast>", 'exec')
        exec(co, test_namespace)
    except AssertionError as e:
        print("Assertion failed on line", node.lineno, ":")
        print(lines[node.lineno])
        # If the error has a message, show it.
        if e.args:
            print(e)
        print()

Next, we transform assert a == b into a function call assert_equal(a, b), which can give more information about the failure. We could turn many other assertions into similar function calls.

class AssertCmpTransformer(ast.NodeTransformer):
    def visit_Assert(self, node):
        if isinstance(node.test, ast.Compare) and \
                len(node.test.ops) == 1 and \
                isinstance(node.test.ops[0], ast.Eq):
            call = ast.Call(func=ast.Name(id='assert_equal', ctx=ast.Load()),
                            args=[node.test.left, node.test.comparators[0]],
                            keywords=[])
            # Wrap the call in an Expr node, because the return value isn't used.
            newnode = ast.Expr(value=call)
            ast.copy_location(newnode, node)
            ast.fix_missing_locations(newnode)
            return newnode

        # Remember to return the original node if we don't want to change it.
        return node

See test_framework/run.py for a working demonstration of both parts.

Real projects

• pytest uses the AST to produce useful error messages when assertions fail.
• astsearch lets you search through Python code based on semantics rather than text, e.g. to find every += 1 in your code.
• astpath is a more powerful search tool using XPath expressions on Python code.
• bellybutton is a linter designed to be readily customised.