Get depth of an abstract syntax tree in Python

Question

I'm trying to work with abstract syntax tree in Python 3.7, I use the library ast from Python library. I would like to know the depth of my tree created with ast.parse. I haven't found any built-in function in the library for this.

I have tried this with this toy example:

py_ast = ast.parse('''pd.read_json(elevations)''')

def depth_ast(tree, depth=0):
    branches_depth = []
    if isinstance(tree, ast.AST):
        for i, field in enumerate(tree._fields):
            child = getattr(tree, field)
            depth += 1
            branch_depth = depth_ast(child, depth)
            branches_depth.append(branch_depth)
    elif isinstance(tree, list):  # multiple cardinality
        if tree == []:
            return depth
        else:
            for i, arg in enumerate(tree):
                depth += 1
                branch_depth = depth_ast(arg, depth)
                branches_depth.append(branch_depth)
            return max(branches_depth)
    else:  # primitives
        return depth
    try:
        return max(branches_depth)
    except:
        return depth

print(depth_ast(py_ast))

I get a depth of 8 for this snippet of code whereas it should be 6 (Module -> Expr -> Call -> Attribute -> Name -> 'pd') and I can't understand why.

Answer 1

There are two related concepts here:

• The depth of a node in a tree, which is the distance between that node and the root of the tree, and
• The height of a node in a tree (more accurately, the height of subtree rooted at that node), which is the maximum distance from that node to any descendant.

The depth of a tree is the maximum depth of any node; the height of a tree is the height of the root. It should be clear that these are the same value.

It's certainly possible to walk the tree recursively, passing the current depth through the recursion. This could even be useful if you needed to annotate each ast node with its depth, for some reason. But if you just want to compute the height of the tree, there's a much more natural recursive solution:

# This is a generic recursive algorithm
def height(node):
    return 1 + max((height(child) for child in children(node)),
                   default = 0)

That relies on the existence of a function which iterates over the children of a node, which is called children above. For the specific case of Python ASTs, we can use the ast module's convenience function, iter_child_nodes:

def depth_ast(root):
    return 1 + max((depth_ast(child)
                       for child in ast.iter_child_nodes(root)),
                   default = 0)

# Or:

def depth_ast(root):
    return 1 + max(map(depth_ast, ast.iter_child_nodes(root)),
                   default = 0)