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

Data Mastery 0
Viewed 47

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 achieve this?

abstract-syntax-tree python
2 Answers

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

blhsing0

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']}
Iain Shelvington0
Related