不要再手動批量替換了,使用 python AST 模組批量替換

語言: CN / TW / HK

前言

在我們日常協作開發時,在團隊內沒有良好的規範或者 code review 機制時,經常會出現使用語義不明的變數,比如 a、b、c 等,使得程式碼可讀性非常差,如果要將變數變更為具有語義的變數,批量替換容易替換錯,而挨個手動替換也容易遺漏或者出錯,因此,需要尋找快捷準確的方式處理這種情況。本文我們只針對 python 語言,首先我們先來了解一下 python 語言的編譯過程。

我們整天和程式碼打交道,都知道高階語言分為解釋型語言和編譯型語言。解釋型語言通過相應的直譯器,將原始碼翻譯成目的碼(機器語言,也就是二進位制形式),邊解釋邊執行,因為邊解釋邊執行的特點,因此執行效率比較低,且執行時不能脫離直譯器;編譯型語言通過編譯器,將原始碼編譯成目的碼(機器語言),因此編譯型語言可以脫離語言環境獨立執行,使用方便且效率高,但每次更改程式碼都需要重新編譯。本文我們詳細介紹 Python(本文我們基於 CPython 直譯器),python 和 Java 類似,直譯器實際上分為兩部分:編譯器和虛擬機器,先將程式碼編譯成位元組碼,然後再由虛擬機器執行。

Python執行過程

在編譯時,需要先經過語法分析,當代碼出現語法錯誤時,就會在這個階段丟擲,接下來我們先了解一下語法分析的基礎知識,編譯的過程分為六個步驟,如下所示:

編譯的六個階段

python 中可以使用 py_compile 模組將原始碼編譯成 PyCodeObject,PyCodeObject 進一步持久化到檔案 pyc 中,直譯器執行 pyc 檔案,在 python 虛擬機器中執行。

我們先來了解一下抽象語法樹是什麼?簡單來說,抽象語法樹經過語法分析(文法定義、文法分析以及消除左遞迴)後,選擇某一條產生式進行展開後的結果。如果還不是很理解,後續我們會針對詞法分析、語法分析、語義分析分別進行講解。python 中可以使用 AST 模組,將程式碼轉換為抽象語法樹,接下來我們進入 python AST 模組的詳解。

AST 基礎知識

ast 模組官方連結:https://docs.python.org/3/library/ast.html#ast-helpers

python 官方對 ast 模組的解釋如下:

The ast module helps Python applications to process trees of the Python abstract syntax grammar. 
The abstract syntax itself might change with each Python release; this module helps to find out 
programmatically what the current grammar looks like.
An abstract syntax tree can be generated by passing ast.PyCF_ONLY_AST as a flag to the compile() 
built-in function, or using the parse() helper provided in this module. The result will be a 
tree of objects whose classes all inherit from ast.AST. An abstract syntax tree can be 
compiled into a Python code object using the built-in compile() function.

更詳細的關於節點的型別描述,可以參考官方文件,就不再贅述。

AST 實戰

建立 AST 並優雅的輸出

我們使用 astpretty 模組,將 ast 物件更優雅的輸出,我們將檔案 web_util.py 中的所有程式碼轉換為 ast 物件,並優雅的輸出,web_util.py 內容如下:

# -*- encoding: utf-8 -*- 

class opUtil:
    @staticmethod
    def add_two_num(a, b):
        return a + b

    @staticmethod
    def mul_two_num(a, b):
        print(a, b)
        return a * b

ast 模組可以將檔案的 read 直接當做輸入,解析並輸出 ast 物件的程式碼如下:

# -*- encoding: utf-8 -*-
import ast
import astpretty

filename = "web_util.py"
f = open(filename)

ast_obj = ast.parse(f.read(), mode="exec")
astpretty.pprint(ast_obj)

輸出結果為:

Module(
    body=[
        ClassDef(
            lineno=3,
            col_offset=0,
            name='opUtil',
            bases=[],
            keywords=[],
            body=[
                FunctionDef(
                    lineno=4,
                    col_offset=4,
                    name='add_two_num',
                    args=arguments(
                        args=[
                            arg(lineno=5, col_offset=20, arg='a', annotation=None),
                            arg(lineno=5, col_offset=23, arg='b', annotation=None),
                        ],
                        vararg=None,
                        kwonlyargs=[],
                        kw_defaults=[],
                        kwarg=None,
                        defaults=[],
                    ),
                    body=[
                        Return(
                            lineno=6,
                            col_offset=8,
                            value=BinOp(
                                lineno=6,
                                col_offset=15,
                                left=Name(lineno=6, col_offset=15, id='a', ctx=Load()),
                                op=Add(),
                                right=Name(lineno=6, col_offset=19, id='b', ctx=Load()),
                            ),
                        ),
                    ],
              decorator_list=[Name(lineno=4, col_offset=5, id='staticmethod', ctx=Load())],
                    returns=None,
                ),
                FunctionDef(
                    lineno=8,
                    col_offset=4,
                    name='mul_two_num',
                    args=arguments(
                        args=[
                            arg(lineno=9, col_offset=20, arg='a', annotation=None),
                            arg(lineno=9, col_offset=23, arg='b', annotation=None),
                        ],
                        vararg=None,
                        kwonlyargs=[],
                        kw_defaults=[],
                        kwarg=None,
                        defaults=[],
                    ),
                    body=[
                        Expr(
                            lineno=10,
                            col_offset=8,
                            value=Call(
                                lineno=10,
                                col_offset=8,
                                func=Name(lineno=10, col_offset=8, id='print', ctx=Load()),
                                args=[
                                    Name(lineno=10, col_offset=14, id='a', ctx=Load()),
                                    Name(lineno=10, col_offset=17, id='b', ctx=Load()),
                                ],
                                keywords=[],
                            ),
                        ),
                        Return(
                            lineno=11,
                            col_offset=8,
                            value=BinOp(
                                lineno=11,
                                col_offset=15,
                                left=Name(lineno=11, col_offset=15, id='a', ctx=Load()),
                                op=Mult(),
                                right=Name(lineno=11, col_offset=19, id='b', ctx=Load()),
                            ),
                        ),
                    ],
              decorator_list=[Name(lineno=8, col_offset=5, id='staticmethod', ctx=Load())],
                    returns=None,
                ),
            ],
            decorator_list=[],
        ),
    ],
)

遍歷 AST 並修改節點

採用 ast.NodeTransformer 的方式遍歷抽象語法樹,我們在遍歷的過程中,將函式中引數 a 命名改為更有意義的 first_num,這個場景在日常開發中很常見,經常會有變數名命名格式不規範,但手動改起來又容易遺漏或者改錯,成本還是很高的。程式碼如下:

# -*- encoding: utf-8 -*-
import ast, astunparse
import astpretty

filename = "web_util.py"
f = open(filename)

ast_obj = ast.parse(f.read(), mode="exec")
astpretty.pprint(ast_obj)


class visitor_ast(ast.NodeTransformer):
    def generic_visit(self, node):
        print("ALL", type(node).__name__)
        fields = node._fields
        if "id" in fields and node.id == "a":
            print("field id", node.id)
            node.id = "first_num"
        ast.NodeVisitor.generic_visit(self, node)

    def visit_FunctionDef(self, node):
        ast.NodeVisitor.generic_visit(self, node)
        args_num = len(node.args.args)
        args = tuple([arg.arg for arg in node.args.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))

    def visit_arg(self, node):
        fields = node._fields
        if "arg" in fields and node.arg == "a":
            print("field arg", node.arg)
            node.arg = "first_num"
        print("ARG", type(node).__name__, node.arg)
        ast.NodeVisitor.generic_visit(self, node)


v = visitor_ast()
v.visit(ast_obj)
astpretty.pprint(ast_obj)
print(astunparse.unparse(ast_obj))

可以使用 astunparse 模組,將 ast 物件還原為程式碼,在這段程式碼中做了兩件事:1、將引數 a 統一修改命名為 first_num;2、在函式中新增 print 日誌。

在 ast 輸出資料 15-18 行中,這幾行表示方法的傳入引數,在遍歷節點時,可以通過 arg 取出入參。

在 ast 輸出資料 65-80 行中,id 為 a 的,是引數 a 的引用。

因此,我們在遍歷時,可以根據節點的 id 或 arg 挑出指定的引數,然後進行替換即可,最後把抽象語法樹再轉化為程式碼。遍歷節點時,我們的程式碼 16-18 行中,將方法內對引數 a 的引用,都改為 first_num,在 30-32 行中,將函式的入參修改為 first_name,24-26 行中,在抽象語法樹中新增列印日誌的節點,第 40 行中,將抽象語法樹轉化為程式碼,引數 a 都被替換為 first_num,第 10 行新增了呼叫函式的日誌輸出,實現了我們想要的效果,最終轉化的程式碼如下:

class opUtil():

    @staticmethod
    def add_two_num(first_num, b):
        print('calling func: add_two_num', 'args:', first_num, b)
        return (first_num + b)

    @staticmethod
    def mul_two_num(first_num, b):
        print('calling func: mul_two_num', 'args:', first_num, b)
        print(first_num, b)
        return (first_num * b)

總結

其實 AST 在我們日常的業務開發中極少用到,AST 模組作為程式碼輔助檢查功能非常有意義,比如語法檢查,除錯錯誤等等,我們上面僅僅用來全域性替換變數及列印日誌,還可以縮小範圍修改某個函式或者某個類裡的。除此之外,還可以用於檢測漢字、closure 檢查等,後續我們有更多的使用案例,也會單獨介紹。