如何使用Python AST访问并处理FunctionDef内的Assert语句？

我正在进行AST代码转换操作，无法访问FunctionDef中的Assert语句，请问如何访问Function节点内部的Assert节点？

以下是我正在修改的示例代码：

assert True == True
def test_code():
    result = 5
    assert result == 5

我的NodeTransformer实现如下：

class TransformMe(ast.NodeTransformer):
    def visit_Assert(self, node):
        return None  # Some processing here

    def visit_FunctionDef(self, node):
        for line_node in node.body:
            if isinstance(line_node, ast.Assert):
                line_node = ast.NodeTransformer.generic_visit(self, line_node)
            for test_node in ast.walk(line_node):
                # More processing here
                pass
        return node

处理代码如下：

module_node = ast.parse(code)
new_module_node = TransformMe().visit(module_node)
print(ast.unparse(new_module_node))

当前输出为：

def test_code():
    result = 5
    assert result == 5

而我期望的输出是：

def test_code():
    result = 5

请问如何实现？如何访问并处理内部的Assert节点？

解决方案 #

你的核心问题是手动遍历FunctionDef的body却没有更新原节点列表，同时忽略了NodeTransformer内置的递归处理机制。

修正后的代码 #

不需要手动遍历函数体，直接依赖NodeTransformer的generic_visit自动递归处理所有子节点即可，最简单的实现如下：

import ast

class TransformMe(ast.NodeTransformer):
    def visit_Assert(self, node):
        return None  # 返回None会将该Assert节点从AST中移除

如果需要保留visit_FunctionDef做额外处理，也只需调用generic_visit：

import ast

class TransformMe(ast.NodeTransformer):
    def visit_Assert(self, node):
        return None

    def visit_FunctionDef(self, node):
        # 自动递归处理函数内部的所有子节点（包括Assert）
        self.generic_visit(node)
        # 这里可以添加函数级别的额外处理逻辑
        return node

关键说明 #

1. NodeTransformer的递归机制：generic_visit会自动遍历当前节点的所有属性（比如FunctionDef的body），对每个子节点调用对应的visit_*方法，所以函数内部的Assert节点会被自动传递到visit_Assert处理。
2. 节点移除逻辑：当visit_Assert返回None时，NodeTransformer会自动将该节点从其父节点的子节点列表中删除，无需手动修改node.body。
3. 原代码的问题：你之前手动遍历node.body时，只是给line_node赋值，没有修改node.body本身的列表内容，所以处理后的节点没有被替换，Assert语句依然保留。

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