Python AST库：如何获取函数调用内部节点的值？

原来的代码仅查找ast.Assign节点，但当前场景中entry_points并不是赋值语句的目标变量，而是作为setup函数调用的关键字参数存在，并且嵌套在if __name__ == "__main__"的代码块（对应ast.If节点的body属性）里，所以原逻辑找不到对应节点，自然失效。

需要调整AST遍历逻辑，逐层定位到目标节点：

1. 找到if __name__ == "__main__"对应的ast.If节点
2. 在该节点的代码块中找到setup()函数调用的ast.Call节点
3. 在Call节点的关键字参数列表里，找到参数名为entry_points的项，提取其值

示例代码 #

import ast

code = '''extras_require={"dev": dev_reqs}
if __name__ == "__main__":
    setup(
        entry_points={
        "console_scripts": ["main=smamesdemo.run.main:main"]
        }
    )
'''

tree = ast.parse(code)

# 遍历AST节点
for node in ast.walk(tree):
    # 定位到if __name__ == "__main__"的节点
    if isinstance(node, ast.If):
        # 校验if条件是否匹配
        if (isinstance(node.test, ast.Compare)
            and isinstance(node.test.left, ast.Name)
            and node.test.left.id == '__name__'
            and len(node.test.ops) == 1
            and isinstance(node.test.ops[0], ast.Eq)
            and len(node.test.comparators) == 1
            and isinstance(node.test.comparators[0], ast.Constant)
            and node.test.comparators[0].value == "__main__"):
            
            # 在if代码块中查找setup函数调用
            for body_node in node.body:
                if isinstance(body_node, ast.Expr) and isinstance(body_node.value, ast.Call):
                    call_node = body_node.value
                    # 确认是setup函数
                    if isinstance(call_node.func, ast.Name) and call_node.func.id == 'setup':
                        # 遍历关键字参数找entry_points
                        for kw in call_node.keywords:
                            if kw.arg == 'entry_points':
                                # 提取并执行获取值
                                expr = ast.Expression(body=kw.value)
                                ast.fix_missing_locations(expr)
                                entry_points = eval(compile(expr, filename='', mode='eval'))
                                print(entry_points)

补充说明 #

• 通过逐层校验节点类型和属性，能精准定位到目标参数；
• 如果代码存在多个setup调用或更深嵌套，可根据实际需求调整遍历逻辑；
• 使用eval时需注意代码安全性，仅在可信代码场景下使用。

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