rewriting code with ast; python

Viewed 3138

I am learning AST and it seems like a powerful thing but I'm confused where the code went and why it disappeared. Say I want to rewrite

example = """def fake(x):\n
    y = ['useless list']\n
    return x
"""

as

example = """def fake(x):\n
    return x
"""

I can't see any way to rewrite in this manner. I can't even find a way to get the text of the line:

In [1]: example = """def fake(x):\n
   ...:     y = ['useless list']\n
   ...:     return x
   ...: """

In [3]: import ast

In [4]: p = ast.parse(example)

In [5]: p
Out[5]: <_ast.Module at 0x7f22f7274a10>

In [6]: p.body
Out[6]: [<_ast.FunctionDef at 0x7f22f7274a50>]

In [7]: p.body
Out[7]: [<_ast.FunctionDef at 0x7f22f7274a50>]

In [8]: f = p.body[0]

In [9]: f
Out[9]: <_ast.FunctionDef at 0x7f22f7274a50>

In [10]: f.body
Out[10]: [<_ast.Assign at 0x7f22f7274b10>, <_ast.Return at 0x7f22f7274c10>]

In [11]: f.name
Out[11]: 'fake'

In [12]: newf = f.body[1:]

In [13]: newf
Out[13]: [<_ast.Return at 0x7f22f7274c10>]

In [14]: z = newf[0]

In [15]: z.value
Out[15]: <_ast.Name at 0x7f22f7274c50>

In [16]: z.value.id
Out[16]: 'x'

Even more surprising is how it gives you the lineno of the start, but not the end. So you know where a function begins, but not where it ends, which is of no use

How can I grab the code and rewrite this func without the list y? Thank you

3 Answers

you can use ast itself for unparsing expression tree:

parsed = ast.parse(text)

source = ast.unparse(parsed)
Related