CODEWITHSUNDEEP
pythonabstract-syntax-tree
Stuart Axon
Stuart Axon

Reputation: 1874

How can I write an ast.NodeTransformer to convert a=1 to a function call a=variables("a")?

The python docs and greentreesnakes show it's possible to use ast.NodeTransformer to change python code via ASTs:

I want to change assignments like a=1 to a=variables(1).

I'm getting a bit confused over what ctx is and when it's needed, the output of astor doesn't seem to have it, yet it seems to be needed.

Upvotes: 1

Views: 1584

Answers (1)

Prodipta Ghosh
Prodipta Ghosh

Reputation: 529

you do not necessarily need the NodeTransformer to do this. If you already have the AST and know the node (an Assign node in this case) you can modify in-place as below.

import ast

src = 'a = 1'
code = 'variables(1)'

tree = ast.parse(src)
assign_node = tree.body[0]
assign_node.value = ast.parse(code).body[0].value

The updated AST can be used further for whatever is your purpose. For example, you can generate the source code back (using say the astor module as astor.to_source(tree) in the above example). NodeTransformer is handy if you need to find the node given some condition. It has a similar visit_* functions like NodeVisitor that allows you to do this in place. Like below:

from ast import parse, NodeTransformer

class ChangeAssignment(NodeTransformer):
    
    def visit_Assign(self, node):
        if assign_node.targets[0].id == 'a':
            node.value = parse('variables(1)').body[0].value
            
        return node

src = 'a = 1'            

t = ChangeAssignment()
tree = parse(src)
tree = t.visit(tree)

The ctx for variable determines if it is for loading, storing or deleting a variable. It is required if you create the node directly (above I use parse to create nodes implicitly). For modifying the AST or generating the source back (like astor.to_source) it does not matter what the ctx you use. If you, however, directly pass the AST to Python compile function, it will complain if the correct ctx is not used.

Upvotes: 2

Related Questions