Path: csiph.com!usenet.pasdenom.info!gegeweb.org!de-l.enfer-du-nord.net!feeder1.enfer-du-nord.net!feeds.phibee-telecom.net!newsfeed.xs4all.nl!newsfeed1.news.xs4all.nl!xs4all!newsgate.cistron.nl!newsgate.news.xs4all.nl!post.news.xs4all.nl!not-for-mail
To: python-list@python.org
From: Peter Otten <__peter__@web.de>
Subject: Re: Modify code with AST
Date: Wed, 12 Jun 2013 11:00:44 +0200
Organization: None
References: <84561fb1-8784-47a6-b0ca-fc3449c1adf2@googlegroups.com>
Mime-Version: 1.0
Content-Type: text/plain; charset="ISO-8859-1"
Content-Transfer-Encoding: 7Bit
User-Agent: KNode/4.7.3
Precedence: list
Newsgroups: comp.lang.python
Message-ID: <mailman.3055.1371027635.3114.python-list@python.org>
Lines: 79
NNTP-Posting-Host: 2001:888:2000:d::a6
Xref: csiph.com comp.lang.python:47759

Ronny Mandal wrote:

> Hello,
> 
> I am trying to write a script which will parse a code segment (with
> ast.parse()), locate the correct function/method node (by name) in the
> resulting tree and replace this function (node) with another function
> (node), e.g.:
> 
> MyMod1.py:
> 
> class FooBar():
>   def Foo(self): #I want to replace this and only this
>     return 1
> 
>   def Bar(self):
>     return 2
> 
> Here is the parser-class:
> 
> class FindAndTransform(NodeTransformer):
>   """Visit the function and check name"""
>   def visit_FunctionDef(self, node):
>     if node.name == 'Foo': #Only replace if name is "Foo"
>       #Create a new function and assign it to node
>       node = parse('''
> def add(n, m):
>   return n + m
> ''')
>       return node
> 
> When I run the parser on MyMod1.py and generate code (with codegen), the
> output is:
> 
> class FooBar():
>   def add(n, m):
>     return n + m
> 
> i.e. both methods are replaced. It seems like "node" in the parser
> contains all method bodies of class FooBar, not only Foo. When ran through
> a debugger, it iterates both methods. What I really wanted to do, was to
> replace only one method (Foo) and leave the other untouched.
> 
> I hope this was understandable conveyed.

I think the main problem is that you have to return the unchanged node (you 
return None which might be an indentation accident). I also had to take the 
add() FunctionDef out of the enclosing Module. So (I don't have codegen or 
is it part of the stdlib?):

import ast

class FindAndTransform(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        if node.name == 'Foo':
            node = ast.parse('''
def add(n, m):
  return n + m
''').body[0]
        return node

if __name__ == "__main__":
    orig = """
class FooBar():
  def Foo(self): #I want to replace this and only this
    return 1

  def Bar(self):
    return 2
"""

    p = ast.parse(orig)
    q = FindAndTransform().visit(p)
    qq = compile(q, "<nofile>", "exec")
    exec(qq)
    assert {n for n in dir(FooBar) if not n.startswith("_")} == {"Bar", 
"add"}