Insert a node into an abstract syntax tree

Question

The ast module's documentation explains how to replace a node in the AST using the NodeTransformer class, but does not explain how to insert a new node into the tree.

For example, given this module:

import foo
import bar

class Baz(object):

    def spam(self):
        pass

I would like to add another import, and set a class variable on Baz.

How can I create and insert these nodes into the AST?

Answer 1

Python ASTs are essentially composed of nested lists, so new nodes can be inserted into these lists once they have been constructed.

First, get the AST that is to be changed:

>>> root = ast.parse(open('test.py').read())

>>> ast.dump(root)
"Module(body=[Import(names=[alias(name='foo', asname=None)]), Import(names=[alias(name='bar', asname=None)]), ClassDef(name='Baz', bases=[Name(id='object', ctx=Load())], body=[FunctionDef(name='spam', args=arguments(args=[Name(id='self', ctx=Param())], vararg=None, kwarg=None, defaults=[]), body=[Pass()], decorator_list=[])], decorator_list=[])])"

We can see that the outer Module has a body attribute that contains the top level elements of the module:

>>> root.body
[<_ast.Import object at 0x7f81685385d0>, <_ast.Import object at 0x7f8168538950>, <_ast.ClassDef object at 0x7f8168538b10>]

Construct an import node and insert:

>>> import_node = ast.Import(names=[ast.alias(name='quux', asname=None)])
>>> root.body.insert(2, import_node)

Like the root module node, the class definition node has a body attribute that contains its members:

>>> classdef = root.body[-1]
>>> ast.dump(classdef)
"ClassDef(name='Baz', bases=[Name(id='object', ctx=Load())], body=[FunctionDef(name='spam', args=arguments(args=[Name(id='self', ctx=Param())], vararg=None, kwarg=None, defaults=[]), body=[Pass()], decorator_list=[])], decorator_list=[])"

So we construct an assignment node and insert it:

>>> assign_node = ast.Assign(targets=[ast.Name(id='eggs', ctx=ast.Store())], value=ast.Str(s='ham')) 
>>> classdef.body.insert(0, assign_node)

To finish, fix up line numbers:

>>> ast.fix_missing_locations(root)
<_ast.Module object at 0x7f816812ef90>

We can verify that our nodes are in place by dumping the root node with ast.dump, or use the unparse^* tool from the CPython repository to generate source from the AST.

The Python3 unparse script^** can be found in the Tools directory of the CPython repository. In Python2 it was located in the Demo directory.

>>> from unparse import Unparser
>>> buf = StringIO()
>>> Unparser(root, buf)
<unparse.Unparser instance at 0x7f81685c6248>
>>> buf.seek(0)
>>> print(buf.read())

import foo
import bar
import quux

class Baz(object):
    eggs = 'ham'

    def spam(self):
        pass
>>> 

When constructing AST nodes, you can get an idea of what the node should look like by using ast.parse and ast.dump (observe that ast.parse wraps the statement in a module):

>>> root = ast.parse('import foo')
>>> ast.dump(root)
"Module(body=[Import(names=[alias(name='foo', asname=None)])])"

^* Credit to this answer for documenting the existence of the unparse script.

^** Use the version of the script from the git branch that corresponds to the Python version being used. For example, using the script from the 3.6 branch on 3.7 code may fail due to differences in the versions' respective grammars.