python - Python AST库获取所有赋值语句

我正在查看Python的ast库，以遍历AST并获取所有分配。如何遍历树以使用ast库只获取python中的赋值语句？

最佳答案

您可以使用ast.NodeTransformer和visit_Assign和visit_AugAssign方法。下面是一个快速演示：

import ast
import inspect

def f():
    a = 1
    # Something else
    print("test")
    a += 1

class CustomNodeTransformer(ast.NodeTransformer):
    def visit_Assign(self, node):
        print("\nvisit_Assign")
        print(node.__dict__)
        print(node.targets[0].__dict__)
        print(node.value.__dict__)
        return node
    def visit_AugAssign(self, node):
        print("\nvisit_AugAssign")
        print(node.__dict__)
        print(node.target.__dict__)
        print(node.value.__dict__)
        return node

nodes = ast.parse(inspect.getsource(f))
CustomNodeTransformer().visit(nodes)

打印内容：

visit_Assign
{'targets': [<_ast.Name object at 0x7fdca78a10b8>], 'value': <_ast.Num object at 0x7fdca78a1080>, 'lineno': 2, 'col_offset': 4}
{'id': 'a', 'ctx': <_ast.Store object at 0x7fdca78b0198>, 'lineno': 2, 'col_offset': 4}
{'n': 1, 'lineno': 2, 'col_offset': 8}

visit_AugAssign
{'target': <_ast.Name object at 0x7fdca7585048>, 'op': <_ast.Add object at 0x7fdca78b0f60>, 'value': <_ast.Num object at 0x7fdca7585080>, 'lineno': 5, 'col_offset': 4}
{'id': 'a', 'ctx': <_ast.Store object at 0x7fdca78b0198>, 'lineno': 5, 'col_offset': 4}
{'n': 1, 'lineno': 5, 'col_offset': 9}

关于这个例子：
AST树的节点是从f函数的源代码中提取的。你可以用任何你喜欢的树。
visit_Assign有一个列表，用于targets之类的赋值。这个例子只使用a = b = 1。
targets[0]用于更改节点。因此，如果上面的方法中没有NodeTransformer，则类似于将节点从树中移除的return。
如果您的任何代码都不应该更改任何节点，那么无论如何也不需要使用return None，您可以只使用ast.NodeTransformer。那你就不需要ast.NodeVisitor。