VeLKerr
VeLKerr

Reputation: 3177

Getting parent of AST node in Python

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?

Upvotes: 21

Views: 8358

Answers (4)

astromancer
astromancer

Reputation: 611

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.

Upvotes: 4

Alex Hall
Alex Hall

Reputation: 36063

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.

Upvotes: 36

Dmitry Rubanovich
Dmitry Rubanovich

Reputation: 2627

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

Upvotes: 0

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.

Upvotes: 1

Related Questions