Python ast Libary - how to retreive value of a specific node-CodePudding

I have the following script:

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

I want to retrieve the following part as python data, dict and list not Nodes.

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

I have tried the following, but I am not actually able to get the values with this:

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}

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

Can anyone help me to archieve this?

CodePudding user response：

You should use ast.NodeVisitor rather than ast.NodeTransformer since you are not making modifications to nodes. But for the purpose of retrieving a single node, it is simpler to use ast.walk instead to traverse the AST in a flat manner and find the ast.Assign node whose target id is 'entry_points', and then wrap its value in an ast.Expression node for compilation with compile and evaluation with eval into an actual Python dict. Remember to use ast.fix_missing_locations to realign the line numbers and offsets of the node when you plant it into another tree:

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)

This outputs:

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

Demo: https://replit.com/@blhsing/OrangeredCurlyApplication

CodePudding user response：

From Python 3.9 you can use the ast.unparse function to get the equivalent code as a string from a 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']}