Python ast库 - 如何检索特定节点的值

Question

我有以下脚本:

extras_require={"dev": dev_reqs}
x = 3
entry_points={
    "console_scripts": ["main=smamesdemo.run.main:main"]
}

我想要将以下部分作为 Python 数据进行检索，使用字典和列表而非节点。

entry_points={
    "console_scripts": ["main=smamesdemo.run.main:main"]
}

我尝试了以下方法，但实际上无法获取这些值:

import ast

class CustomNodeTransformer(ast.NodeTransformer):
    def visit_Assign(self, node):
        print(node.value.__dict__)
        return node

with open("./setup.py") as f:
    code = f.read()

    node = ast.parse(code)
    CustomNodeTransformer().visit(node)

{'keys': [<ast.Constant object at 0x0000029E992962C0>], 'values': [<ast.Name object at 0x0000029E99296260>], 'lineno': 1, 'col_offset': 15, 'end_lineno': 1, 'end_col_offset': 32}

预期结果：

entry_points={"console_scripts": ["main=smamesdemo.run.main:main"]}

有人能帮我实现这个吗？

- Data Mastery

Answer 1

由于您不对节点进行修改，因此应使用ast.NodeVisitor而不是ast.NodeTransformer。但是，为了检索单个节点，最简单的方法是使用ast.walk以平面方式遍历AST，并找到目标id为'entry_points'的ast.Assign节点，然后将其值包装在ast.Expression节点中，以便使用compile进行编译并使用eval评估为实际的Python字典。记得在将其插入另一个树时使用ast.fix_missing_locations来重新对齐节点的行号和偏移量：

import ast

code = '''extras_require={"dev": dev_reqs}
x = 3
entry_points={
    "console_scripts": ["main=smamesdemo.run.main:main"]
}'''

for node in ast.walk(ast.parse(code)):
    if isinstance(node, ast.Assign) and node.targets[0].id == 'entry_points':
        expr = ast.Expression(body=node.value)
        ast.fix_missing_locations(expr)
        entry_points = eval(compile(expr, filename='', mode='eval'))

print(entry_points)

这将输出：

{'console_scripts': ['main=smamesdemo.run.main:main']}

演示：https://replit.com/@blhsing/OrangeredCurlyApplication

- blhsing

Answer 2

从Python 3.9开始，您可以使用ast.unparse函数将Node转换为等效的代码字符串。

import ast


class CustomNodeTransformer(ast.NodeTransformer):
    def visit_Assign(self, node):
        print(ast.unparse(node))
        return node


code = """
extras_require={"dev": dev_reqs}
x = 3
entry_points={
    "console_scripts": ["main=smamesdemo.run.main:main"]
}
"""
node = ast.parse(code)
CustomNodeTransformer().visit(node)

# Output:
# extras_require = {'dev': dev_reqs}
# x = 3
# entry_points = {'console_scripts': ['main=smamesdemo.run.main:main']}

- Iain Shelvington