Question

我正在研究Python的ast库来遍历AST并获取所有分配。如何使用ast库遍历树以仅获取python中的赋值语句？

Answer 1

您可以使用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。然后，您不需要return。

Python AST库获取所有赋值语句

1 个答案: