How to use Python ast module's NodeTransformer class

Viewed 2005

I'm trying switch all ">=" in my program to "<". This is what I have so far:

import ast
from astor import to_source

class Visitor(ast.NodeTransformer):
    def generic_visit(self, node):
        ast.NodeVisitor.generic_visit(self, node)

    def visit_GtE(self, node):
        node2 = ast.Lt()
        ast.copy_location(node2, node)
        ast.NodeVisitor.generic_visit(self, node2)
        return node2

def mutate(filename, numMutations):
    file = open(filename)
    contents = file.read()
    parsed = ast.parse(contents)
    nodeVisitor = Visitor()
    nodeVisitor.visit(parsed)
    print(to_source(parsed))

When I call mutate, nothing get's changed. What's wrong with my visit_GtE method?

1 Answers

The NodeTransformer does not transform an AST in-place. The visit method returns a new AST. So you need to assign this to a variable in your mutate function, and output the new AST.

def mutate(filename, numMutations):
    file = open(filename)
    contents = file.read()
    parsed = ast.parse(contents)
    nodeVisitor = Visitor()
    # Capture the new AST.
    transformed = nodeVisitor.visit(parsed)
    print(to_source(transformed))
Related