从 Python AST 获取对应于具有给定名称的特定变量的所有节点

Question

考虑以下代码：

1 | x = 20
2 | 
3 | def f():
4 |     x = 0
5 |     for x in range(10):
6 |         x += 10
7 |     return x
8 | f()
9 |
10| for x in range(10):
11|     pass
12| x += 1
13| print(x)

执行上述代码后x的值为10。现在，我怎样才能得到 class Name 的所有节点，其 id 是 x 并引用第 1、10 行中使用的 x , 12 和 13?

换句话说，f里面的x和其他的x是不一样的。是否有可能得到他们的 AST 节点，只有脚本和脚本的 AST 而没有执行它？

Answer 1

在遍历 AST 树时，跟踪上下文；从全局上下文开始，然后当您遇到 FunctionDef 或 ClassDef 或 Lambda 节点时，将该上下文记录为堆栈（退出相关节点时再次弹出堆栈）。

然后您可以只查看全局上下文中的 Name 个节点。您也可以跟踪 global 标识符（我会在每个堆栈级别使用一组）。

使用 NodeVisitor subclass:

import ast

class GlobalUseCollector(ast.NodeVisitor):
    def __init__(self, name):
        self.name = name
        # track context name and set of names marked as `global`
        self.context = [('global', ())]

    def visit_FunctionDef(self, node):
        self.context.append(('function', set()))
        self.generic_visit(node)
        self.context.pop()

    # treat coroutines the same way
    visit_AsyncFunctionDef = visit_FunctionDef

    def visit_ClassDef(self, node):
        self.context.append(('class', ()))
        self.generic_visit(node)
        self.context.pop()

    def visit_Lambda(self, node):
        # lambdas are just functions, albeit with no statements, so no assignments
        self.context.append(('function', ()))
        self.generic_visit(node)
        self.context.pop()

    def visit_Global(self, node):
        assert self.context[-1][0] == 'function'
        self.context[-1][1].update(node.names)

    def visit_Name(self, node):
        ctx, g = self.context[-1]
        if node.id == self.name and (ctx == 'global' or node.id in g):
            print('{} used at line {}'.format(node.id, node.lineno))

演示（在 t 中给出了示例代码的 AST 树）：

>>> GlobalUseCollector('x').visit(t)
x used at line 1
x used at line 10
x used at line 12
x used at line 13

并在函数中使用 global x：

>>> u = ast.parse('''\
... x = 20
...
... def g():
...     global x
...     x = 0
...     for x in range(10):
...         x += 10
...     return x
...
... g()
... for x in range(10):
...     pass
... x += 1
... print(x)
... ''')
>>> GlobalUseCollector('x').visit(u)
x used at line 1
x used at line 5
x used at line 6
x used at line 7
x used at line 8
x used at line 11
x used at line 13
x used at line 14

从 Python AST 获取对应于具有给定名称的特定变量的所有节点

Getting all the nodes from Python AST that correspond to a particular variable with a given name

python

abstract-syntax-tree