使用 AST 模块来改变和删除赋值/函数调用

Question

例如，如果我想将大于更改为小于或等于，我已成功执行：

def visit_Gt(self, node):
    new_node = ast.GtE()
    return ast.copy_location(new_node, node)

我如何访问/检测赋值操作（=）和函数调用()并简单地删除它们？我正在阅读 AST 文档，但找不到访问赋值或函数调用类然后不返回任何内容的方法。

我正在寻找的赋值操作的一个例子：

print("Start")
x = 5
print("End")

变成：

print("Start")

print("End")

我正在寻找删除函数调用的示例：

 print("Start")
 my_function_call(Args)
 print("End")

成为

print("Start")

print("End")

Answer 1

您可以使用ast.NodeTransformer()子类来改变现有的 AST 树：

import ast

class RemoveAssignments(ast.NodeTransformer):
    def visit_Assign(self, node):
        return None

    def visit_AugAssign(self, node):
        return None

new_tree = RemoveAssignments().visit(old_tree)

上面的类返回None以从输入树中完全删除节点。和节点包含整个赋值语句、生成结果Assign的AugAssign表达式以及目标列表（将结果分配到的 1 个或多个名称）。

这意味着上面的内容将会变成

print('Start!')
foo = 'bar'
foo += 'eggs'
print('Done!')

进入

print('Start!')


print('Done!')

如果您需要做出更细粒度的决策，请直接查看分配的子节点，或者通过将子节点传递给来self.visit()让转换器进一步调用visit_*它们的钩子（如果它们存在）：

class RemoveFunctionCallAssignments(NodeTransformer):
    """Remove assignments of the form "target = name()", so a single name being called

    The target list size plays no role.

    """
    def visit_Assign(self, node):
        if isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name):
            return None
        return node

在这里，只有None当赋值的值侧（右侧的表达式）是Call应用于直接节点的节点时，我们才返回Name。返回传入的原始节点对象意味着它不会被替换。

要替换顶级函数调用（即那些没有赋值或进一步表达式的函数调用），请查看Expr节点；这些是表达式语句，而不仅仅是作为其他构造一部分的表达式。如果您有一个Expr带有 的节点Call，则可以将其删除：

def visit_Expr(self, node):
    # stand-alone call to a single name is to be removed
    if isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name):
        return None
    return node

另请参阅优秀的Green Tree Snakes文档，其中介绍了 AST 树的工作以及更多示例。