For example if I wanted to change greater than to less than or equal to I have successfully executed:
def visit_Gt(self, node):
new_node = ast.GtE()
return ast.copy_location(new_node, node)
How would I visit/detect an assignment operation (=
) and a function call ()
and simply delete them? I'm reading through the AST documentation and I can't find a way to visit the assignment or function call classes and then return nothing.
An example of what I'm seeking for assignment operations:
print("Start")
x = 5
print("End")
Becomes:
print("Start")
print("End")
And an example of what I'm seeking for deleting function calls:
print("Start")
my_function_call(Args)
print("End")
Becomes
print("Start")
print("End")
You can use a ast.NodeTransformer()
subclass to mutate an existing AST tree:
import ast
class RemoveAssignments(ast.NodeTransformer):
def visit_Assign(self, node):
return None
def visit_AugAssign(self, node):
return None
new_tree = RemoveAssignments().visit(old_tree)
The above class removes None
to completely remove the node from the input tree. The Assign
and AugAssign
nodes contain the whole assignment statement, so the expression producing the result, and the target list (1 or more names to assign the result to).
This means that the above will turn
print('Start!')
foo = 'bar'
foo += 'eggs'
print('Done!')
into
print('Start!')
print('Done!')
If you need to make more fine-grained decisions, look at the child nodes of the assignment, either directly, or by passing the child nodes to self.visit()
to have the transformer further call visit_*
hooks for them if they exist:
class RemoveFunctionCallAssignments(NodeTransformer):
"""Remove assignments of the form "target = name()", so a single name being called
The target list size plays no role.
"""
def visit_Assign(self, node):
if isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name):
return None
return node
Here, we only return None
if the value side of the assignment (the expression on the right-hand side) is a Call
node that is applied to a straight-forward Name
node. Returning the original node object passed in means that it'll not be replaced.
To replace top-level function calls (so those without an assignment or further expressions), look at Expr
nodes; these are expression statements, not just expressions that are part of some other construct. If you have a Expr
node with a Call
, you can remove it:
def visit_Expr(self, node):
# stand-alone call to a single name is to be removed
if isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name):
return None
return node
Also see the excellent Green Tree Snakes documentation, which covers working on the AST tree with further examples.