在 Python 中获取抽象语法树的深度

Get depth of an abstract syntax tree in Python

我正在尝试使用 Python 3.7 中的抽象语法树,我使用 Python 库中的库 ast。我想知道用 ast.parse 创建的树的深度。我还没有在库中找到任何内置函数。

我已经用这个玩具示例试过了:

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))

这段代码的深度是 8,而它应该是 6 (Module -> Expr -> Call -> Attribute -> Name -> 'pd'),我不明白为什么。

这里有两个相关的概念:

  • 树中一个节点的深度,也就是那个节点到树根的距离,
  • 树中一个节点的高度(更准确地说,是以该节点为根的子树的高度),这是该节点到任何后代的最大距离。

树的深度是任意节点的最大深度;一棵树的高度就是根的高度。应该清楚,这些是相同的值。

当然可以递归地遍历树,通过递归传递当前深度。如果出于某种原因需要用深度注释每个 ast 节点,这甚至可能很有用。但是如果你只想计算树的高度,有一个更自然的递归解决方案:

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

这依赖于迭代节点子节点的函数的存在,上面称为 children。对于 Python AST 的特定情况,我们可以使用 ast 模块的便利函数,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)