1. 程式人生 > 實用技巧 >python AST 抽象語法樹

python AST 抽象語法樹

Abstract Sytax Tree
參考:
https://docs.python.org/3/library/ast.html#ast.NodeTransformer
https://www.cnblogs.com/yssjun/p/10069199.html
Abstract Syntax Trees即抽象語法樹。Ast是python原始碼到位元組碼的一種中間產物,藉助ast模組可以從語法樹的角度分析原始碼結構。此外,我們不僅可以修改和執行語法樹,還可以將Source生成的語法樹unparse成python原始碼。因此ast給python原始碼檢查、語法分析、修改程式碼以及程式碼除錯等留下了足夠的發揮空間。

1. AST簡介

Python官方提供的CPython直譯器對python原始碼的處理過程如下:

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

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

  • 原始碼解析 --> 語法樹 --> 抽象語法樹(AST) --> 控制流程圖 --> 位元組碼

上述過程在python2.5之後被應用。python原始碼首先被解析成語法樹,隨後又轉換成抽象語法樹。在抽象語法樹中我們可以看到原始碼檔案中的python的語法結構。
大部分時間程式設計可能都不需要用到抽象語法樹,但是在特定的條件和需求的情況下,AST又有其特殊的方便性。

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

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

其中 三引號可以根據書寫的方式智慧換行,輸出如下:


def add(x, y):
    return x+y
    
print(add(3,5))

2. 建立AST

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

這是python自帶的函式

  • source -- 字串或者AST(Abstract Syntax Trees)物件。一般可將整個py檔案內容file.read()傳入。
  • filename -- 程式碼檔名稱,如果不是從檔案讀取程式碼則傳遞一些可辨認的值。
  • mode -- 指定編譯程式碼的種類。可以指定為 exec, eval, single。
  • flags -- 變數作用域,區域性名稱空間,如果被提供,可以是任何對映物件。
  • flags和dont_inherit是用來控制編譯原始碼時的標誌。
>>> cm = compile(func_def, filename='<string>', mode='exec')
>>> exec(cm)
8
>>> type(cm)
code

上面func_def經過compile編譯得到位元組碼,cm即code物件,True == isinstance(cm, types.CodeType)。

2.2 生成AST

>>> cm1 = ast.parse(func_def,filename='<unknown>', mode='exec')
>>> type(cm1)
_ast.Module
>>> ast.dump(cm1)
(
body=[
    FunctionDef(name='add', 
                args=arguments(
                                args=[arg(arg='x', annotation=None), arg(arg='y', annotation=None)], 
                                vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]
                              ), 
                body=[Return(
                                value=BinOp(left=Name(id='x', ctx=Load()), op=Add(), right=Name(id='y', ctx=Load()))
                            )
                     ], 
                decorator_list=[], 
                returns=None), 
    Expr(value=Call(
                    func=Name(id='print', ctx=Load()), 
                    args=[Call(func=Name(id='add', ctx=Load()), args=[Num(n=3), Num(n=5)], keywords=[])], 
                    keywords=[])
                    )
    ]
)

可以看到,這裡對原始碼進行了解析

  • 首先是原始碼字串的主體 body,可以看到,一個是FunctionDef,也就是我們定義的add函式,另外一個是下面使用的print函式
  • 對於第一個主體 FunctionDef,可以看到裡面的 name是 ‘add’,也就是函式的名字是 add, 再一個就是args,引數,可以看到一個是 'x',annotation=None,另外一個引數是 y, annntation=None; 然後裡面又有一個body,裡面可以看到是 return返回值,其中BinOp表示雙目操作符,操作符的左值為x,操作符 opAdd(),也就是將我們原始碼中的 +轉換成了 Add()函式,最後就是右值 y
  • 最後就是 print函式,可以看到,valuesCall 呼叫了一個函式,其中 函式名func為 add,引數有兩個,一個是3,一個是5

3. 遍歷語法樹

python提供了兩種方式來遍歷整個語法樹

  • 節點的訪問就只需要重寫 visit_nodename函式,在裡面定義引數即可
  • 這裡節點的visit 會預設根據 ast中的 nodename 去訪問 visit_nodename 函式,同時如果當前節點存- 在children,比如 FunctionDef 中存在 BinOp 節點,若想 visit BinOp這個節點,就需要在 FunctionDef中增加一句 self.generic_visit()來達到遞迴訪問;如果不加,就只能訪問當前節點

generic_visit(node)
This visitor calls visit() on all children of the node. Note that child nodes of nodes that have a custom visitor method won’t be visited unless the visitor calls generic_visit() or visits them itself.

3.1 ast.NodeVisitor

比如 我們將 func_def 的 add 函式中的加法運算改為減法

class CodeVisitor(ast.NodeVisitor):
    def visit_BinOp(self, node):# 這個函式的訪問是由於 Visit_FunctionDef的先訪問再generic_visit才訪問的
        print('Bin')            # 如果Visit_FunctionDef中沒有generic_visit的話,則這個函式是不會訪問的
        if isinstance(node.op, ast.Add):
            node.op = ast.Sub()
            
        self.generic_visit(node)
    
    def visit_FunctionDef(self, node):
        print('Function Name: %s'% node.name)
        self.generic_visit(node) # FunctionDef中還包含有 BinOp,因此會進去visit BinOP
        
    def visit_Call(self, node):
        print("call") 
        self.generic_visit(node) # 因為AST的Call中還包含有一個Call,因此會重複再訪問一次 
    
        
r_node = ast.parse(func_def)
visitor = CodeVisitor()
visitor.visit(r_node) # 這裡的visit函式會根據 node 的語法樹去遍歷裡面的函式,

輸出:

Function Name: add
Bin
call
call

3.2 ast.NodeTransformer

ANodeVisitorsubclass that walks the abstract syntax tree and allows modification of nodes

使用NodeVisitor主要是通過修改語法樹上節點的方式改變AST結構,NodeTransformer主要是替換ast中的節點。

class CodeTransformer(ast.NodeTransformer):
    def visit_BinOp(self, node):
        if isinstance(node.op, ast.Add):
            node.op = ast.Sub()
        self.generic_visit(node)
        return node

    def visit_FunctionDef(self, node):
        self.generic_visit(node) # 這裡表示先去訪問裡面的children node        
        if node.name == 'add':
            node.name = 'sub'
        args_num = len(node.args.args)
        
        args_num = len(node.args.args)
        args = tuple([arg.arg for arg in node.args.args])
        print(str(args))
        func_log_stmt = ''.join(["print('calling func: %s', " % node.name, "'args:'", ", %s" * args_num % args ,')'])
        node.body.insert(0, ast.parse(func_log_stmt))
        
        #func_log_stmt = ''.join(["print 'calling func: %s', " % node.name, "'args:'", ", %s" * args_num % args])
        #node.body.insert(0, ast.parse(func_log_stmt))

        return node

    def visit_Name(self, node):
        replace = {'add': 'sub', 'x': 'a', 'y': 'b'}
        re_id = replace.get(node.id, None)
        node.id = re_id or node.id
        self.generic_visit(node)
        return node
    
    def visit_arg(self, node):
        self.generic_visit(node)
        replace = {'x':'a', 'y':'b'}
        node.arg = replace[node.arg]
        return node
        
        
r_node = ast.parse(func_def)
transformer = CodeTransformer()
r_node = transformer.visit(r_node)
#print(astunparse.dump(r_node))
source = astunparse.unparse(r_node) # astunparse 一般python不自帶,需要conda 或者 pip安裝
print(source)

輸出:

('a', 'b')


def sub(a, b):
    print('calling func: sub', 'args:', a, b)
    return (a - b)
print(sub(3, 5))

可以看加入了一個print語句,然後將變數名字由 x, y 改為了 a, b


Keep in mind that if the node you’re operating on has child nodes you must either transform the child nodes yourself or call the generic_visit() method for the node first.


Don’t use theNodeVisitorif you want to apply changes to nodes during traversal. For this a special visitor exists (NodeTransformer) that allows modifications.