Getting parent of AST node in Python

Question:

I’m working with Abstract Syntax Trees in Python 3. The ast library gives many ways to get children of the node (you can use iter_child_nodes() or walk()) but no ways to get parent of one. Also, every node has links to its children, but it hasn’t links to its parent.

How I can get the parent of AST node if I don’t want to write some plugin to ast library?

What is the most correct way to do this?

Asked By: VeLKerr

||

Answers:

You might create some hash table associating AST nodes to AST nodes and scan (recursively) your topmost AST tree to register in that hash table the parent of each node.

It wouldn’t be really be a plugin, but you can always write a function which adds a weakref to parent in every child.

Answered By: Dmitry Rubanovich

Here’s some actual code:

for node in ast.walk(root):
    for child in ast.iter_child_nodes(node):
        child.parent = node

There’s no need for a hash table, you can just put an attribute directly on the node.

Answered By: Alex Hall

You can also use ast.NodeTransformer to achieve this:

Code:

import ast


class Parentage(ast.NodeTransformer):
    # current parent (module)
    parent = None

    def visit(self, node):
        # set parent attribute for this node
        node.parent = self.parent
        # This node becomes the new parent
        self.parent = node
        # Do any work required by super class 
        node = super().visit(node)
        # If we have a valid node (ie. node not being removed)
        if isinstance(node, ast.AST):
            # update the parent, since this may have been transformed 
            # to a different node by super
            self.parent = node.parent
        return node

Usage:

module = Parentage().visit(ast.parse('def _(): ...'))
assert module.parent is None
assert module.body[0].parent is module

Later on when you want to edit the tree in some other way, you can subclass:

class SomeRefactoring(Parentage):
    def visit_XXX(node):
        self.generic_visit(node)
        f'do some work on {node.parent} here if you want'
        return node

Note:

Its worth noting that some nodes can have multiple parents. For example:

module = ast.parse("warnings.warn('Dinosaurs!')")
func = module.body[0].value.func
name, ctx = ast.iter_child_nodes(func)
assert ctx is next(ast.iter_child_nodes(name))

Which shows that the same ast.Load node ctx has two parents – func and name. The parent will be set by the last position that the node appears in in the tree.

Answered By: codeMonkey