通過ast的parse方法得到ast tree的根節點root_node, 我看可以通過根節點來遍歷語法樹,從而對python代碼進行分析和修改。ast.parse(可以直接查看ast模塊的源代碼)方法實際上是調用內置函數compile進行編譯,如下所示:
def parse(source, filename='<unknown>', mode='exec'): """ Parse the source into an AST node. Equivalent to compile(source, filename, mode, PyCF_ONLY_AST). """ return compile(source, filename, mode, PyCF_ONLY_AST)傳遞給compile特殊的flag = PyCF_ONLY_AST, 來通過compile返回抽象語法樹。(關于compile的第四個參數flag有時間可以深入研究下)
2. 節點類型分析
語法樹中的每個節點都對應ast下的一種類型,根節點是ast.Moudle類型,在分析的時候可以通過isinstance函數方便的進行節點類型的判斷。ast中存在的節點的所有類型可以參考:ast節點類型比如 a = 10這樣一條語句對應ast.Assign節點類型,而Assign節點類型分別有兩個子節點, 分別為ast.Name類型的a和ast.Num類型的10等。我們可以通過ast.dump(node)函數來將node格式化,并進行打印,以查看節點內容,以“a = 10”這行代碼為例。
Module(body=[Assign(targets=[Name(id='a', ctx=Store())], value=Num(n=10))])(1) root節點
Module(body=[Assign(targets=[Name(id='a', ctx=Store())], value=Num(n=10))])root節點是Module類型,由于只有一行代碼,所有root節點只有Assign這樣一個子節點。
(2) 子節點
Assign(targets=[Name(id='a', ctx=Store())], value=Num(n=10))上述的Assign節點有三個子節點,分別是Name, Store和Num.Name(id='a', ctx=Store())Num(n=10)而Name有一個子節點,Store.Store()(Store表示Name中操作時賦值, 類型的有Load,del, 具體參考節點類型的文檔)一個簡單的“a = 10”的這樣一行代碼,我們就可以通過上述的這種ast tree去分析和修改代碼結構。二. 語法樹的遍歷分析
可以通過ast模塊的提供的visitor來對語法樹進行遍歷。ast.NodeVisitor是一個專門用來遍歷語法樹的工具,我們可以通過繼承這個類來完成對語法樹的遍歷以及遍歷過程中的處理。
1. visitor的定義
class CodeVisitor(ast.NodeVisitor): def generic_visit(self, node): print type(node).__name__ ast.NodeVisitor.generic_visit(self, node) def visit_FunctionDef(self, node): print type(node).__name__ ast.NodeVisitor.generic_visit(self, node) def visit_Assign(self, node): print type(node).__name__ ast.NodeVisitor.generic_visit(self, node)如上述代碼,定義類CodeVisitor,繼承自NodeVisitor,這里面主要有兩種類型的函數,一種的generic_visit,一種是"visit_" + "Node類型"。visitor首先從根節點root進行遍歷,在遍歷的過程中,假設節點類型為Assign,如果存在visit_Assign類型的函數,則調用visit_Assgin函數,如果不存在則調用generic_visit函數。總的來說就是每個節點類型都有專用的類型處理函數,如果不存在,則調用通用的的處理函數generic_visit.關于visitor進行語法樹的遍歷,stackoverflow上有一篇文章講的比較詳細:Simple example of how to use ast.NodeVisitor注意:在每個函數處理中,根據需求需要加上ast.NodeVisitor.generic_visit(self, node)這段代碼,否則visitor不會繼續訪問當前節點的子節點。e.g. 如果定義如下的函數:def visit_Moudle(self, node): print type(node).__name__那么,首先訪問根節點root,root為Moudle類型,會調用visit_Moudle函數,由于visit_Moudle函數中沒有調用NodeVisitor.generic_visit(self, node),所以此次遍歷只遍歷了根節點root,并沒有遍歷其他節點。
2. walk方式遍歷
也可以通過ast.walk對ast tree進行遍歷,如下:
for node in ast.walk(tree): if isinstance(node, ast.FunctionDef): print(node.name)三. 節點的修改
ast模塊同樣提供了一個NodeTransfomer節點來支持對node的修改,NodeTransfomer繼承自NodeVisitor,并重寫了generic_visit函數。對于NodeTransfomer的generic_visit以及visit_ + 節點類型的函數,都需要返回一個node,可以返回原始node,一個新的替代的node,或者是返回Node代表remove掉這個節點。假設我們有如下的代碼:
"""ast test code"""a = 10b = "test"print a我們定義一個NodeTransform的visitor如下:
class ReWriteName(ast.NodeTransformer): def generic_visit(self, node): has_lineno = getattr(node, "lineno", "None") col_offset = getattr(node, "col_offset", "None") print type(node).__name__, has_lineno, col_offset ast.NodeTransformer.generic_visit(self, node) return node def visit_Name(self, node): new_node = node if node.id == "a": new_node = ast.Name(id = "a_rep", ctx = node.ctx) return new_node def visit_Num(self, node): if node.n == 10: node.n = 100 return node在visit_Name中,將變量"a"替換成了變量"a_rep",執行到a = 10以及print a的時候,都會將a替換成a_rep,并返回一個新節點。在visit_Num中,簡單粗暴的將10替換成了100,返回修改后的原節點。我們通過如下方式運用這個NodeTransfomer visitor:
file = open("code.py", "r")source = file.read()visitor = ReWriteName()root = ast.parse(source)root = visitor.visit(root)ast.fix_missing_locations(root) code_object = compile(root, "<string>", "exec")exec code_objectast作用在python解析語法之后,編譯成pyCodeObject字節碼結構之前,通過NodeTransformer修改后,返回修改后的語法樹,我們通過內置模塊compile編譯成pyCodeObject對象,交給python虛擬機執行。執行結果:100可以看到,我們同時將a = 10和print a兩處將a名字換成了a_rep,并將10替換成了100,最后打印的結果是100,成功修改了語法樹的節點。關于節點的修改,這里有比較好的例子可以參考:https://greentreesnakes.readthedocs.org/en/latest/examples.html注意:修改語法樹節點,尤其是刪除一個語法樹節點時要慎重,因為修改或者刪除后有可能返回錯誤的語法樹,直到compile或者執行的時候才會發現問題。通過節點修改python code就可以通過上述方法進行,不過請注意,在運用visitor的代碼中有ast.fix_missing_locations(root)這樣一行代碼,這是因為我們自己創建的節點是不包含lineno以及col_offset這些必要的屬性,必須手動修改添加指定,新添加的節點代碼的行位置以及偏移位置。
四. 修復節點位置
1 屬性分析
每個節點都有一些相應的屬性,lineno以及col_offset是每個節點都必須有的屬性,分別代表行號以及在這行中的偏移。另外每個節點都有一些自己的特殊屬性,如上訴的Module含有body屬性,Assign含有targets屬性等。lineno以及col_offset這兩個屬性,如果是python中原本代碼的節點,如Assign、Name、Num等(注:Moudle和Store這樣的節點是沒有lineno以及col_offset屬性的),但是如果我們通過NodeTransFormer新增的節點,默認是不存在這些屬性的,我們可以通過三種方法來fix這些節點的lineno以及col_offset屬性。2 屬性的修復
我們可以通過相應的方法,對默認沒有lineno以及col_offset的節點進行位置的修復,以方便在代碼中獲取每個節點的位置信息,主要有三種方法進行修復。1)ast.fix_missing_locations(node)函數遞歸的將父節點的位置信息(lineno以及col_offset)賦值給沒有位置信息的子節點。2)ast.copy_location(new_node, node)將node的位置信息拷貝給new_node節點,并返回new_node節點。當我們將舊節點替換成一個新節點的時候,這種方法比較適用。3)ast.increment_lineno(node, n=1)將node節點以及其所以子節點的行號加上n。3 分析
我們通過“三. 節點的修改"中的例子來分析location信息。在例子中,我們只有在visit_Name的時候返回的新的節點,這時候節點是沒有lineno以及col_offset屬性,我們可以通過兩種方式獲取。一是如上述代碼中,利用ast.fix_missing_locations函數來修復,在"a = 10"以及"print a"中,Name節點a跟父節點的lineno相同,但是此時col_offset會有差異。二是我們將visit_Name的代碼修改如下:def visit_Name(self, node): new_node = node if node.id == "a": new_node = ast.Name(id = "a_rep", ctx = node.ctx) ast.copy_location(new_node, node) return new_node通過copy_location將舊節點的location信息拷貝給新節點。參考:Green Tree Snakes - the missing Python AST docs --非常詳盡的ast的模塊的分析文檔。Simple example of how to use ast.NodeVisitor --stackoverflow上一篇比較明了的回答Instrumenting the AST --簡單的分析與應用ast — Abstract Syntax Trees --官方文檔AST 模塊:用 Python 修改 Python 代碼 --比較詳細的介紹了修改