ast manipulation

Tsize tsize69 at gmail.com
Tue Nov 17 09:55:03 EST 2009


Hello,

I am hoping for a little help.  I have been playing with the python
ast module and have run into
an issue that I need a little push on.  I would like to be able to
change a specific element in a
specific node in an ast then compile the resulting ast.  Consider the
simplified example below
with its output. In this example I would like a way to change a
specific addition operation.  With the NodeTransformer I see how to
change every addition operator but not how to change a specific one.
I would like this to work on both the 2.6 and 3.1 branches.  Ideally I
would like to read a file, count the instances of an operation of
interest and then use an index to make the
changes.

I am probably missing something simple but I am lost right now.

import ast

class SwitchMinusPlus(ast.NodeTransformer):

    def visit_BinOp(self, node):
        node = self.generic_visit(node)
        if isinstance(node.op, ast.Add):
            node.op = ast.Sub()
        return node

myfile = open('trivial.py').read()
print myfile
tree = compile(myfile, '<string>', 'exec', ast.PyCF_ONLY_AST)
print ast.dump(tree, annotate_fields=False, include_attributes=False)
node = SwitchMinusPlus().visit(ast.parse(myfile))
print ast.dump(node, annotate_fields=False, include_attributes=False)

Which gives the following output:  Note that this code changes the
addition operator to an
subtraction operator at the AST level for every instance.

a = 8
b = 6
c = b + a
d =  c + a
Module([Assign([Name('a', Store())], Num(8)), Assign([Name('b', Store
())], Num(6)),
Assign([Name('c', Store())], BinOp(Name('b', Load()), Add(), Name('a',
Load()))),
Assign([Name('d', Store())], BinOp(Name('c', Load()), Add(), Name('a',
Load())))])

Module([Assign([Name('a', Store())], Num(8)), Assign([Name('b', Store
())], Num(6)),
Assign([Name('c', Store())], BinOp(Name('b', Load()), Sub(), Name('a',
Load()))),
Assign([Name('d', Store())], BinOp(Name('c', Load()), Sub(), Name('a',
Load())))])

Thanks in advance,
Thomas



More information about the Python-list mailing list