PythonでASTを逆変換してソースコードにする その２

whileループを抜き出して別の関数に埋め込む例

import ast
import codegen

expr="""
def main():
    i = 0
    while(i<3):
        print(i)
        i += 1
    print("done")

if __name__=="__main__":
    main()
"""
p=ast.parse(expr)
while_part = p.body[0].body.pop(1)
p1 = ast.parse("def sub():\n    pass")
p1.body[0].body.append(while_part)
p1.body[0].body.pop(0)
p.body.insert(1,p1)

print("****** original ******")
print(expr)
print("****** converted ******")
print(codegen.to_source(p))

****** original ******

def main():
    i = 0
    while(i<3):
        print(i)
        i += 1
    print("done")

if __name__=="__main__":
    main()

****** converted ******
def main():
    i = 0
    print 'done'

def sub():
    while (i < 3):
        print i
        i += 1
if (__name__ == '__main__'):
    main()