Playing with ASTs in Python

Earlier this week, we had company hack day, to take a short break from the hustle of the roadmap (and there’s lots!!) to work on something else, and my team worked on a CLI tool. We used click to build some cool stuff, and one of the features I built wraps a Python lambda handler with IOpipe’s library – which meant I got to play with AST things!

AST whosawhatsit?

AST stands for Abstract Syntax Tree – it’s common around the modern web not because of tooling like Babel, as well as your friendly neighborhood code linter. Any system that seeks to read a program and understand or reason about it probably uses AST in the process. There’s a nice picture over on the wikipedia article on ASTs that maybe makes them more understandable.

ASTs in Python

The AST module is in the standard library (woot!) which makes it very easy to use. The docs aren’t terribly great, so there are some “missing docs” provided out there on the net. In particular, the extra docs are helpful for referencing the various types available in the ast module, however I had great luck with writing some Python and then inspecting the AST to see what the AST broke it down to.

Don’t believe me though, since it’s in the standard library, you can pop open your handy REPL and try:

>>> import ast
>>> ast.parse('x=2')
<_ast.Module object at 0x7f373f7efe10>
>>> ast.dump(ast.parse('x=2'))
"Module(body=[Assign(targets=[Name(id='x', ctx=Store())], value=Num(n=2))])"

And you can tell how you can see the tree that the ast is understanding. The body of the statement x=2 only has one node, an assignment, and that assignment has its id (x) and value (2).

Let’s say you want to walk the tree, that’s easily done:

for node in ast.iter_child_nodes(ast.parse(some_code)):

From there, I can look for which node is the handler, and I can add a decorator to its properties:

if isinstance(node, ast.FunctionDef):
    if node.name == handler_name:
        node.decorator_list.append(
            ast.Name(id='iopipe.decorator', ctx=ast.Load()))

Boom! Decorated! Inject some other code nodes into the AST, and I’ve got the new tree I want, but how to get it back out? That, you’ll need a bit of help for. I used astor, after trying codegen and it didn’t love imports at the top of the file, perhaps something to do with Python 3? (shrug)

In general, I was very pleased with how far I got in a few hours, having never played with this section of Python before, and using ASTs myself!