Changing python AST while saving comments

I am currently working with AST in python. I take a python file, create an AST, modify it and then recompile it back to source. I am using a transformer that adds a getter to the class (I am using the visitor template with ast.NodeTransformer). Currently my code is working as expected, but not saving comments, which is my problem. Below is my code:

#visits nodes and generates getters or setters
def genGet(file,type,func):
    global things
    things['func'] = func
    things['type'] = type
    with open(file) as f:
        code = f.read()             #get the code
    tree = ast.parse(code)          #make the AST from the code
    genTransformer().visit(tree)    #lets generate getters or setters depending on type argument given in our transformer so the genTransformer function
    source = meta.asttools.dump_python_source(tree) #recompile the modified ast to source code
    newfile = "{}{}".format(file[:-3],"_mod.py")
    print "attempting to write source code new file: {}".format(newfile) #tell everyone we will write our new source code to a file
    outputfile = open(newfile,'w+')
    outputfile.write(source)        #write our new source code to a file
    outputfile.close()


class genTransformer(ast.NodeTransformer):
    ...

      

I've done some research on lib2to3, which seems to be able to store comments, but nothing has been found so far, but it helps with my problem. For example, I found the code below but don't understand it. It looks like it saves comments but doesn't allow my changes. I am getting a missing attribute error when I run it.

import lib2to3
from lib2to3.pgen2 import driver
from lib2to3 import pygram, pytree
import ast

def main():
    filename = "%s" % ("exfunctions.py")
    with open(filename) as f:
        code = f.read()
    drv = driver.Driver(pygram.python_grammar, pytree.convert)
    tree = drv.parse_string(code, True)
    # ast transfomer breaks if it is placed here
    print str(tree)
    return

      

I'm having trouble finding a package or strategy to preserve comments when converting an AST. So far, my research hasn't worked for me. What can I use that will allow me to modify the AST but also keep the comments?

+3


source to share





All Articles