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

问题描述

考虑下面的代码:

1 |x = 202 |3 |定义 f():4 |x = 05 |对于范围内的 x(10):6 |x += 107 |返回 x8 |F()9 |10|对于范围内的 x(10):11|经过12|x += 113|打印(x)

执行上述代码后x的值为10.现在，我如何获得 Name 类的所有节点，其 ids 是 x 并引用 x在第 1、10、12 和 13 行中使用了吗?

换句话说，f 中的 x 与 x 的其余部分不同.是否有可能获得他们的 AST 节点，只有脚本和脚本的 AST 而没有执行它?

解决方案

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

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

使用 NodeVisitor 子类:

导入 ast类 GlobalUseCollector(ast.NodeVisitor):def __init__(self, name):self.name = 姓名# 跟踪上下文名称和标记为全局"的名称集self.context = [('global', ())]def visit_FunctionDef(self, node):self.context.append(('function', set()))self.generic_visit(节点)self.context.pop()# 以同样的方式对待协程visit_AsyncFunctionDef = visit_FunctionDefdef visit_ClassDef(self, node):self.context.append(('class', ()))self.generic_visit(节点)self.context.pop()def visit_Lambda(self, node):# lambdas 只是函数，虽然没有语句，所以没有赋值self.context.append(('function', ()))self.generic_visit(节点)self.context.pop()def visit_Global(self, node):断言 self.context[-1][0] == 'function'self.context[-1][1].update(node.names)defvisit_Name(self, node):ctx, g = self.context[-1]if node.id == self.name and (ctx == 'global' or node.id in g):print('{} 在行 {}'.format(node.id, node.lineno) 处使用

Demo(在 t 中给出示例代码的 AST 树):

<预><代码>>>>GlobalUseCollector('x').visit(t)x 在第 1 行使用x 在第 10 行使用x 在第 12 行使用x 在第 13 行使用

并在函数中使用global x:

<预><代码>>>>u = ast.parse('''... x = 20...... def g():... 全球 x... x = 0... 对于范围内的 x(10):... x += 10...返回 x...... G()... 对于范围内的 x(10):... 经过... x += 1... 打印(x)... ''')>>>GlobalUseCollector('x').visit(u)x 在第 1 行使用x 在第 5 行使用x 在第 6 行使用x 在第 7 行使用x 在第 8 行使用x 在第 11 行使用x 在第 13 行使用x 在第 14 行使用

Consider the code below:

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)

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?

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?

解决方案

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

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

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

Demo (given the AST tree for your sample code in 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

And using global x in a function:

>>> 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屋！