从Python AST获取与给定名称的特定变量对应的所有节点 [英] Getting all the nodes from Python AST that correspond to a particular variable with a given name

问题描述

请考虑以下代码：

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 。现在，如何获取所有类为 Name 且其 id s为 x 并引用第1、10、12和13行中使用的 x ？

The value of x after execution of the code above is 10. Now, how can I get all the nodes with class Name whose ids are x and refer to the x that's being used in lines 1, 10, 12 and 13?

换句话说， f 内部的 x 与的其余部分不同x s。

In other words, the x inside of f is different from the rest of the xs. Is it possible to get their AST nodes, having only the script and script's AST while not executing it?

推荐答案

在AST上行走时，是否可以获取仅具有脚本和脚本AST的AST节点？树，跟踪上下文；从全局上下文开始，然后遇到 FunctionDef 或 ClassDef 或 Lambda 节点，将该上下文记录为堆栈（退出相关节点时再次弹出堆栈）。

When walking the AST tree, track the context; start with a global context, then as you encounter FunctionDef or ClassDef or Lambda nodes, record that context as a stack (pop the stack again when exiting the relevant node).

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

You can then simply only look at Name nodes in the global context. You can track global identifiers too (I'd use a set per stack level).

使用 NodeVisitor 子类：

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))

Demo（在 t ）：

>>> 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获取与给定名称的特定变量对应的所有节点的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！