ast模塊與astunparse模塊


 AST簡介 

Abstract Syntax Trees即抽象語法樹。Ast是python源碼到字節碼的一種中間產物,借助ast模塊可以從語法樹的角度分析源碼結構。此外,我們不僅可以修改和執行語法樹,還可以將Source生成的語法樹unparse成python源碼。因此ast給python源碼檢查、語法分析、修改代碼以及代碼調試等留下了足夠的發揮空間。可以通過將ast.PyCF_ONLY_AST作為標志傳遞給compile()內置函數,或者使用此模塊中提供的parse()幫助器生成抽象語法樹。結果將是一個對象樹,其類都繼承自ast.AST。可以使用內置的compile()函數將抽象語法樹編譯成Python代碼對象。

官網:https://docs.python.org/3.6/library/ast.html

Python官方提供的CPython解釋器對python源碼的處理過程如下:

  1. Parse source code into a parse tree (Parser/pgen.c)
  2. Transform parse tree into an Abstract Syntax Tree (Python/ast.c)
  3. Transform AST into a Control Flow Graph (Python/compile.c)
  4. Emit bytecode based on the Control Flow Graph (Python/compile.c)

即實際python代碼的處理過程如下:

源代碼解析 --> 語法樹 --> 抽象語法樹(AST) --> 控制流程圖 --> 字節碼

上述過程在python2.5之后被應用。python源碼首先被解析成語法樹,隨后又轉換成抽象語法樹。在抽象語法樹中我們可以看到源碼文件中的python的語法結構。

大部分時間編程可能都不需要用到抽象語法樹,但是在特定的條件和需求的情況下,AST又有其特殊的方便性。

下面是一個抽象語法的簡單實例。

Module(body=[
    Print(
          dest=None,
          values=[BinOp( left=Num(n=1),op=Add(),right=Num(n=2))],
          nl=True,
 )]) 

Compile函數

先簡單了解一下compile函數。compile函數將源代碼編譯成可以由exec()或eval()執行的代碼對象。返回code類型。

compile(source, filename, mode[, flags[, dont_inherit]]) 

  • source -- 字符串或者AST(Abstract Syntax Trees)對象。一般可將整個py文件內容file.read()傳入。
  • filename -- 代碼文件名稱,如果不是從文件讀取代碼則傳遞一些可辨認的值。
  • mode -- 指定編譯代碼的種類。可以指定為 exec, eval, single。
  • flags -- 變量作用域,局部命名空間,如果被提供,可以是任何映射對象。
  • flags和dont_inherit是用來控制編譯源碼時的標志。
func_def = \
"""
def add(x, y):
    return x + y
print(add(3, 5))
"""

cm = compile(func_def, '<string>', 'exec')
print(type(cm))
isinstance(cm, types.CodeType)

exec(func_def) #傳入的類型可以是str、bytes或code。
exec(cm) #傳入的類型可以是str、bytes或code。

上面func_def經過compile編譯得到字節碼,cm即code對象,True == isinstance(cm, types.CodeType)。

compile(source, filename, mode, ast.PyCF_ONLY_AST)  <==> ast.parse(sourcefilename='<unknown>'mode='exec')

生成ast

 除了python內置ast模塊可以生成抽象語法樹,還有很多第三方庫,如astunparse, codegen, unparse等。這些第三方庫不僅能夠以更好的方式展示出ast結構,還能夠將ast反向導出python source代碼。

安裝astunparse:pip install astunparse

astunparse官網:https://pypi.org/project/astunparse/

import ast, astunparse
func_def = \
"""
def add(x, y):
    return x + y
print(add(3, 5))
"""

r_node = ast.parse(func_def)
print(ast.dump(r_node))
print(astunparse.dump(r_node))
import ast
import astunparse
func_def = \
"""
a = 3
b = 5
def add(x, y):
    return x + y
print(add(a,b))
"""

def nodeTree(node:str):
    str2list = list(node.replace(' ', ''))
    count = 0
    for i, e in enumerate(str2list):
        if e == '(':
            count += 1
            str2list[i] = '(\n{}'.format('|   ' * count)
        elif e == ')':
            count -= 1
            str2list[i] = '\n{})'.format('|   ' * count)
        elif e == ',':
            str2list[i] = ',\n{}'.format('|   ' * count)
        elif e == '[':
            count += 1
            str2list[i] = '[\n{}'.format('|   ' * count)
        elif e == ']':
            count -= 1
            str2list[i] = '\n{}]'.format('|   ' * count)

    return ''.join(str2list)

# def nodeTree(node:str):
#     reStr = ''
#     count = 0
#     for e in node:
#         if e == '(':
#             count += 1
#             reStr += '(\n{}'.format('|   ' * count)
#         elif e == ')':
#             count -= 1
#             reStr += '\n{})'.format('|   ' * count)
#         elif e == ',':
#             reStr += ',\n{}'.format('|   ' * count)
#         elif e == '[':
#             count += 1
#             reStr += '[\n{}'.format('|   ' * count)
#         elif e == ']':
#             count -= 1
#             reStr += '\n{}]'.format('|   ' * count)
#         else:
#             reStr += e
#     return reStr

r_node = ast.parse(func_def)
cm = ast.dump(r_node)
print(nodeTree(cm))
自定義展示出ast結構的函數

通過ast的parse方法得到ast tree的根節點r_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返回抽象語法樹。

節點類型分析

import ast
root_node = ast.parse("print('hello world')")
print(ast.dump(root_node))

輸出:

Module(body=[Expr(value=Call(func=Name(id='print', ctx=Load()), args=[Str(s='hello world')], keywords=[]))])

語法樹中的每個節點都對應ast下的一種類型,根節點是ast.Moudle類型,在分析的時候可以通過isinstance函數方便的進行節點類型的判斷。

import ast
root_node = ast.parse("print('hello world')")
print(ast.dump(root_node))
print(isinstance(root_node,ast.Module))
print(isinstance(root_node,ast.Expr))
print(isinstance(root_node.body[0],ast.Expr))

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去分析和修改代碼結構。

語法樹的遍歷分析

1. visitor的定義

可以通過ast模塊的提供的visitor來對語法樹進行遍歷。
ast.NodeVisitor是一個專門用來遍歷語法樹的工具,我們可以通過繼承這個類來完成對語法樹的遍歷以及遍歷過程中的處理。

import ast
import astunparse
func_def = \
"""
a = 3
b = 5
def add(x, y):
    return x + y
print(add(a,b))
"""
class CodeVisitor(ast.NodeVisitor):
    def generic_visit(self, node):
        print(type(node).__name__,end=', ')
        ast.NodeVisitor.generic_visit(self, node)

    def visit_FunctionDef(self, node):
        print(type(node).__name__,end=', ')
        ast.NodeVisitor.generic_visit(self, node)

    def visit_Assign(self, node):
        print(type(node).__name__,end=', ')
        ast.NodeVisitor.generic_visit(self, node)
r_node = ast.parse(func_def)
visitor = CodeVisitor()
visitor.visit(r_node)
View Code
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方式遍歷

 

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 = 10
b = "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_object

ast作用在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這些必要的屬性,必須手動修改添加指定,新添加的節點代碼的行位置以及偏移位置。

修復節點位置

我們可以通過相應的方法,對默認沒有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信息拷貝給新節點。

AST應用

AST模塊實際編程中很少用到,但是作為一種源代碼輔助檢查手段是非常有意義的;語法檢查,調試錯誤,特殊字段檢測等。

上面通過為函數添加調用日志的信息是一種調試python源代碼的一種方式,不過實際中我們是通過parse整個python文件的方式遍歷修改源碼

相關代碼

# -- encoding:utf-8 --
"""
Greate by ibf on 2019
"""
import ast
import astunparse
func_def = \
"""
a = 3
b = 5
def add(x, y):
    return x + y
print(add(a,b))
"""

def nodeTree(node:str):
    str2list = list(node.replace(' ', ''))
    count = 0
    for i, e in enumerate(str2list):
        if e == '(':
            count += 1
            str2list[i] = '(\n{}'.format('|   ' * count)
        elif e == ')':
            count -= 1
            str2list[i] = '\n{})'.format('|   ' * count)
        elif e == ',':
            str2list[i] = ',\n{}'.format('|   ' * count)
        elif e == '[':
            count += 1
            str2list[i] = '[\n{}'.format('|   ' * count)
        elif e == ']':
            count -= 1
            str2list[i] = '\n{}]'.format('|   ' * count)

    return ''.join(str2list)

# def nodeTree(node:str):
#     reStr = ''
#     count = 0
#     for e in node:
#         if e == '(':
#             count += 1
#             reStr += '(\n{}'.format('|   ' * count)
#         elif e == ')':
#             count -= 1
#             reStr += '\n{})'.format('|   ' * count)
#         elif e == ',':
#             reStr += ',\n{}'.format('|   ' * count)
#         elif e == '[':
#             count += 1
#             reStr += '[\n{}'.format('|   ' * count)
#         elif e == ']':
#             count -= 1
#             reStr += '\n{}]'.format('|   ' * count)
#         else:
#             reStr += e
#     return reStr

'''遍歷節點'''
class CodeVisitor(ast.NodeVisitor):
    def generic_visit(self, node):
        print(type(node).__name__,end=', ')
        ast.NodeVisitor.generic_visit(self, node)

    def visit_FunctionDef(self, node):
        print(type(node).__name__,end=', ')
        ast.NodeVisitor.generic_visit(self, node)

    def visit_Assign(self, node):
        print(type(node).__name__,end=', ')
        ast.NodeVisitor.generic_visit(self, node)

'''修改節點'''
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 == 3:
            node.n = 100
        return node
    def visit_BinOp(self,node):
            node.op = ast.Sub()
            return node

r_node = ast.parse(func_def)
visitor = ReWriteName()
visitor.visit(r_node)
print(astunparse.unparse(r_node))#將ast反向導出python source代碼。
# ast.fix_missing_locations(r_node)
View Code

 

 

參考鏈接:https://www.cnblogs.com/yssjun/p/10069199.html  http://www.dalkescientific.com/writings/diary/archive/2010/02/22/instrumenting_the_ast.html  https://pycoders-weekly-chinese.readthedocs.io/en/latest/issue3/static-modification-of-python-with-python-the-ast-module.html#cpython

 


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM