为何Python AST无法检测函数内部的if语句？

问题场景 #

统计学背景的AST新手遇到如下情况：

• 当if语句位于函数内部时，使用ast.iter_child_nodes遍历顶层节点无法检测到该if语句；
• 当if语句不在函数内时，同样的遍历方式可正常检测到。

示例代码 #

第一段代码（函数内包含if语句） #

import ast
from pprint import pprint

tree = ast.parse("""
def add(a, b):
  return a + b

def subr(a,b):
    if 2>3:
        print("true")
    else:
        print("false")
    return 0
""")

for node in ast.iter_child_nodes(tree):
    print(isinstance(node, ast.If))

运行结果：两次输出False，遍历顶层节点无法找到函数内的if。

第二段代码（顶层if语句） #

import ast
from pprint import pprint

tree = ast.parse("""
if 2>3:
    print("true")
else:
    print("false")
""")

for node in ast.iter_child_nodes(tree):
    print(isinstance(node, ast.If))

运行结果：输出True，if是顶层节点的直接子节点，可被遍历到。

问题原因 #

Python的AST是层级结构：

• ast.parse返回的根节点是Module类型，它的直接子节点是顶层定义（比如FunctionDef函数定义节点）。
• 函数内部的if语句属于FunctionDef节点的body属性下的子节点，并非Module节点的直接子节点。
• ast.iter_child_nodes(tree)仅遍历Module的直接子节点（也就是两个FunctionDef），不会递归进入函数内部，所以检测不到函数内的if语句。

解决方法 #

如果需要检测所有层级的if语句，推荐使用ast.walk()，它会递归遍历AST的所有节点：

import ast

tree = ast.parse("""
def add(a, b):
  return a + b

def subr(a,b):
    if 2>3:
        print("true")
    else:
        print("false")
    return 0
""")

# 递归遍历所有节点，检测If类型
for node in ast.walk(tree):
    if isinstance(node, ast.If):
        print("找到If节点：")
        print(ast.dump(node, indent=2))

运行后会输出函数内部的if节点信息。

也可以手动遍历层级：先遍历顶层的FunctionDef，再遍历其body属性下的子节点：

import ast

tree = ast.parse("""
def add(a, b):
  return a + b

def subr(a,b):
    if 2>3:
        print("true")
    else:
        print("false")
    return 0
""")

for node in ast.iter_child_nodes(tree):
    if isinstance(node, ast.FunctionDef):
        # 遍历函数体内部的节点
        for inner_node in ast.iter_child_nodes(node):
            if isinstance(inner_node, ast.If):
                print("找到函数内的If节点：")
                print(ast.dump(inner_node, indent=2))

内容的提问来源于stack exchange，提问作者desert_ranger