python?ast模塊詳析與用法

更新時間：2023年07月25日 10:28:43 作者：愛幻想-hjyp

這篇文章主要給大家介紹了關(guān)于python?ast模塊詳析與用法的相關(guān)資料, Python的ast(Abstract Syntax Trees,抽象語法樹)模塊是一個內(nèi)置模塊,用于解析Python代碼并生成語法樹,需要的朋友可以參考下

基本概念

在 python 中，我們可以通過自帶的 ast 模塊來對解析遍歷語法樹，通過ast.parse()可以將字符串代碼解析為抽象語法樹，然后通過ast.dump()可以打印這棵語法樹。

除了ast模塊外，還有 astor 模塊，其中的 astor.to_sourse() 函數(shù)可以將語法樹Node轉(zhuǎn)換為代碼， astor.dump_tree() 可以很好地格式化整棵樹。

除了這些基礎(chǔ)操作外，我們還可以遍歷和修改整棵語法樹。

比如，對于a = 10 來說，我們可以先解析成抽象語法樹，然后打印所有的結(jié)點，如下所示。根據(jù)輸出，我們可以看到根節(jié)點是Module類型的，然后其body是Assign類型的。對于Assign類型的結(jié)點，可以繼續(xù)劃分為Name結(jié)點（表示變量名）和Constant結(jié)點（表示變量內(nèi)容）。

node = ast.parse('a = 10')
print(astor.dump_tree(node))
# Module(body=[Assign(targets=[Name(id='a')], value=Constant(value=10, kind=None), type_comment=None)], type_ignores=[])

節(jié)點類型

上面的簡單示例向我們展示了幾種基本結(jié)點類型（Assign、Name、Constant），接下來我們將會展示其他幾種常見的結(jié)點類型和示例，完整的節(jié)點類型可以查閱節(jié)點類型。大體上，我們可以把結(jié)點類型分為葉子結(jié)點類型和非葉子結(jié)點類型，比如Assign就是非葉子結(jié)點類型，Name和Constant是葉子結(jié)點類型，因為他們不會有子結(jié)點了。

ast.Assign

Assign 類型用來表示賦值語句，比如a = 10、 b = a 這樣的賦值語句都是Assign結(jié)點類型，他并不是一個葉子結(jié)點，因為它的下面一般還有 Name 結(jié)點。

ast.Name

Name類型用來表示一個變量的名稱，是一個葉子結(jié)點。比如對于b = a 這樣的賦值語句，子結(jié)點就是兩個Name。

node = ast.parse('a = b')
print(astor.dump_tree(node.body[0]))
# Assign(targets=[Name(id='a')], value=Name(id='b'), type_comment=None)

ast.Constant

表示一個不可變內(nèi)容，它可以是Number 、string，只要其內(nèi)容是不可變的，都是ast.Constant類型的結(jié)點，它是一個葉子結(jié)點。

node = ast.parse('a = 100')
print(astor.dump_tree(node.body[0]))
# Assign(targets=[Name(id='a')], value=Constant(value=100, kind=None), type_comment=None)
node = ast.parse('a = "paddle"')
print(astor.dump_tree(node.body[0]))
# Assign(targets=[Name(id='a')], value=Constant(value='paddle', kind=None), type_comment=None)

ast.Call

表示函數(shù)的調(diào)用，比如paddle.to_tensor()。非葉子節(jié)點類型，一般包含三個屬性：func、args、 keywords。

func：代表調(diào)用函數(shù)的名稱，一般是一個ast.Name或ast.Constant類型的結(jié)點，如果是連續(xù)調(diào)用，會是一個ast.Call結(jié)點。
args：代表函數(shù)傳入的位置參數(shù)和可變參數(shù)。
keywords：代表函數(shù)傳入的關(guān)鍵字參數(shù)。

node = ast.parse('paddle.to_tensor(1, a = 10)')
print(astor.dump_tree(node.body[0]))

# Expr(
    value=Call(func=Attribute(value=Name(id='paddle'), attr='to_tensor'),
        args=[Constant(value=1, kind=None)],
        keywords=[keyword(arg='a', value=Constant(value=10, kind=None))]))

對于上面的例子，我們通過可視化可以看到，頂層是一個ast.Expr類型的結(jié)點，表示一個表達(dá)式。下面是ast.Call 結(jié)點，Call 結(jié)點包含一個ast.Attribute結(jié)點，表示調(diào)用者和調(diào)用的方法名，paddle是調(diào)用者，to_tensor是方法名；一個ast.Constant類型的args，表示函數(shù)的位置參數(shù)；一個ast.keyword，表示函數(shù)的關(guān)鍵字參數(shù)。

下面我們看一個比較復(fù)雜的示例，多個函數(shù)的連續(xù)調(diào)用。根據(jù)輸出結(jié)果可以看到，最后的調(diào)用reshape在最外層，然后一直向內(nèi)遞歸，子結(jié)點還是ast.Call類型的結(jié)點。

node = ast.parse('a.to_tensor(1, a = 10).reshape(1)')
print(astor.dump_tree(node.body[0]))

Expr(
    value=Call(
        func=Attribute(
            value=Call(func=Attribute(value=Name(id='a'), attr='to_tensor'),      
                args=[Constant(value=1, kind=None)],
                keywords=[keyword(arg='a', value=Constant(value=10, kind=None))]),
            attr='reshape'),
        args=[Constant(value=1, kind=None)],
        keywords=[]))

ast.Attribute

上面的例子中出現(xiàn)了ast.Attribute結(jié)點，Attribute結(jié)點可以理解為屬性，是一個非葉子結(jié)點。它包含兩個字段，value字段和attr字段。對于a.shape來說value指明調(diào)用者，即a；attr指明調(diào)用的方法名，即shape。

node = ast.parse('a.shape')
print(astor.dump_tree(node.body[0]))

Expr(value=Attribute(value=Name(id='a'), attr='shape'))

結(jié)點的遍歷

在ast模塊中，可以借助繼承ast.NodeVisitor類來完成結(jié)點的遍歷，該類具有兩種訪問結(jié)點的方法，一種是針對所有結(jié)點類型通用的訪問方法generic_visit()，另一種是針對某個類型結(jié)點的訪問方法 visit_xxx，其中xxx代表具體的結(jié)點類型。generic_visit()函數(shù)是遍歷每個結(jié)點的入口函數(shù)，隨后會調(diào)用visitor()函數(shù)，獲取該結(jié)點的類型，然后判斷是否有遍歷該類型結(jié)點的函數(shù)，如果有則調(diào)用 visit_xxx類型的方法，如果沒有則調(diào)用通用generic_visit()方法。

ast源碼

class NodeVisitor(object):
    def visit(self, node):
        """Visit a node."""
        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)
    def generic_visit(self, node):
    	# 可以看到 generic_visit函數(shù)會調(diào)用visit函數(shù)，然后尋找并調(diào)用特定類型的visit函數(shù)。 
        """Called if no explicit visitor function exists for a node."""
        for field, value in iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, AST):
                        self.visit(item)
            elif isinstance(value, AST):
                self.visit(value)
    def visit_Constant(self, node):
        value = node.value
        type_name = _const_node_type_names.get(type(value))
        if type_name is None:
            for cls, name in _const_node_type_names.items():
                if isinstance(value, cls):
                    type_name = name
                    break
        if type_name is not None:
            method = 'visit_' + type_name
            try:
                visitor = getattr(self, method)
            except AttributeError:
                pass
            else:
                import warnings
                warnings.warn(f"{method} is deprecated; add visit_Constant",
                              PendingDeprecationWarning, 2)
                return visitor(node)
        return self.generic_visit(node)

示例

下面是一個例子，我們定義了一個繼承ast.NodeVisitor的類，并且重寫了visit_attribute方法，這樣在遍歷到ast.Attribute結(jié)點時，會輸出當(dāng)前調(diào)用的屬性名或方法名，對于其他類型的結(jié)點則會輸出結(jié)點類型。

class CustomVisitor(ast.NodeVisitor):
    def visit_Attribute(self, node):
        print('----' + node.attr)
        ast.NodeVisitor.generic_visit(self, node)
    def generic_visit(self, node):
        print(node.__class__.__name__)
        ast.NodeVisitor.generic_visit(self, node)
code = textwrap.dedent(
    '''
    import paddle
    x = paddle.to_tensor([1, 2, 3])
    axis = 0
    y = paddle.max(x, axis=axis)
    '''
)
node = ast.parse(code)
visitor = CustomVisitor()
visitor.generic_visit(node)

需要注意的是，當(dāng)我們重寫visit_xxx函數(shù)后，一定要記得再次調(diào)用ast.NodeVisitor.generic_visit(self, node)，這樣才會繼續(xù)遍歷整棵語法樹。

結(jié)點的修改

對于結(jié)點的修改可以借助ast.NodeTransformer 類來完成，ast.NodeTransformer繼承自ast.NodeVisitor類，重寫了generic_visit方法，該方法可以傳入一個結(jié)點，并且返回修改后的結(jié)點，從而完成語法樹的修改。

示例

在該示例中，我們定義了CustomVisitor類來修改ast.Call 結(jié)點。具體來說，當(dāng)遍歷到Call類型的結(jié)點后，流程如下：

首先會調(diào)用get_full_attr方法獲取整個api名稱，如果是普通方法調(diào)用，則會返回完整的調(diào)用名稱，比如torch.tensor()會返回torch.tensor；如果是連續(xù)的方法調(diào)用，比如x.exp().floor()，則會返回ClassMethod.floor。
然后調(diào)用 ast.NodeVisitor.generic_visit(self, node) ，進(jìn)行深度優(yōu)先的修改，這樣就可以一層層遞歸，先修改內(nèi)層，再修改外層。
如果是普通的方法調(diào)用，則修改結(jié)點后返回；
如果是連續(xù)的方法調(diào)用，需要先通過astor.to_source(node)獲取前綴方法，即調(diào)用者，保留前綴方法名稱的同時，修改目前的方法名后返回。具體是通過'{}.{}()'實現(xiàn)的。

def get_full_attr(node):
        # torch.nn.fucntional.relu
        if isinstance(node, ast.Attribute):
            return get_full_attr(node.value) + '.' + node.attr
        # x.abs() -> 'x'
        elif isinstance(node, ast.Name):
            return node.id
        # for example ast.Call
        else:
            return 'ClassMethod'
            
class CustomVisitor(ast.NodeTransformer):
    
    def visit_Call(self, node):
        # 獲取api的全稱
        full_func = get_full_attr(node.func)

        # post order
        ast.NodeVisitor.generic_visit(self, node)
        
        # 如果是普通方法調(diào)用，直接改寫整個結(jié)點即可
        if full_func == 'torch.tensor':
            # 將 torch.tensor() 改寫為 paddle.to_tensor()
            code = 'paddle.to_tensor()'
            new_node = ast.parse(code).body[0]
            return new_node.value
        
        # 如果是類方法調(diào)用，需要取前面改寫后的方法作為 func.value 
        if full_func == 'ClassMethod.floor':
            # 獲取前綴方法作為 func.value
            new_func = astor.to_source(node).strip('\n')
            new_func = new_func[0: new_func.rfind('.')]
            # 將 floor() 改寫為 floor2()
            code = '{}.{}()'.format(new_func, 'floor2')
            new_node = ast.parse(code).body[0]
            return new_node.value

        # 其余結(jié)點不修改
        return node

code = textwrap.dedent(
    '''
    import torch
    x = torch.tensor([1, 2, 3])
    x = x.exp().floor()
    '''
)
node = ast.parse(code)
visitor = CustomVisitor()
node = visitor.generic_visit(node)
result_code = astor.to_source(node)
print(result_code)

參考鏈接

https://blog.csdn.net/ThinkTimes/article/details/110831176?ydreferer=aHR0cHM6Ly9jbi5iaW5nLmNvbS8%3D
https://greentreesnakes.readthedocs.io/en/latest/
https://github.com/PaddlePaddle/PaConvert