Use AST module to mutate and delete assignment/function calls

问题

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")

回答1:

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.

来源：https://stackoverflow.com/questions/48771051/use-ast-module-to-mutate-and-delete-assignment-function-calls

标签

python

abstract-syntax-tree

mutation