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):
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):
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.