Groups | Search | Server Info | Keyboard shortcuts | Login | Register [http] [https] [nntp] [nntps]

Groups > comp.lang.python > #47746 > unrolled thread

Modify code with AST

Back to article view | Back to comp.lang.python

Contents

  Modify code with AST Ronny Mandal <ronnyma@volatile.no> - 2013-06-12 00:22 -0700
    Re: Modify code with AST Peter Otten <__peter__@web.de> - 2013-06-12 11:00 +0200
      Re: Modify code with AST Ronny Mandal <ronnyma@volatile.no> - 2013-06-12 03:36 -0700

#47746 — Modify code with AST

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.


Answers are highly appreciated.

Regards,

Ronny Mandal

[toc] | [next] | [standalone]

#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"}

[toc] | [prev] | [next] | [standalone]

#47772

> 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?):

Thank you, my problem is now solved. It was the indentation of the return statement which was the culprit. I forgot that visit() (in this context) is called once for each FunctionDef-node. I moved the return-statement one click to the left, then it worked.

codegen is not a part of the stdlib. It is written by Armin Ronacher (the same person who wrote ast) and it uses ast to generate python code from the AST.

Thanks!


Regards,

Ronny Mandal

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

[toc] | [prev] | [standalone]

Back to top | Article view | comp.lang.python

csiph-web