Python内核源码解析：初步了解语法树AST优化的过程_python 源码解析

代码优化分为前端优化和后端优化。前端优化是指在源代码层面进行优化，例如语法树优化、中间代码优化等。后端优化是指在目标代码层面进行优化，例如机器码优化等。

Python AST 语法树优化的实现在文件 Python/ast_opt.c 中，主要是通过调用 PyAST_Optimize 函数来实现。PyAST_Optimize 函数的实现如下：

int
_PyAST_Optimize(mod_ty mod, PyArena *arena, _PyASTOptimizeState *state)
{
    PyThreadState *tstate;
    int recursion_limit = Py_GetRecursionLimit();
    int starting_recursion_depth;

    /* Setup recursion depth check counters */
    tstate = _PyThreadState_GET();
    if (!tstate) {
        return 0;
    }
    /* Be careful here to prevent overflow. */
    starting_recursion_depth = (tstate->recursion_depth < INT_MAX / COMPILER_STACK_FRAME_SCALE) ?
        tstate->recursion_depth * COMPILER_STACK_FRAME_SCALE : tstate->recursion_depth;
    state->recursion_depth = starting_recursion_depth;
    state->recursion_limit = (recursion_limit < INT_MAX / COMPILER_STACK_FRAME_SCALE) ?
        recursion_limit * COMPILER_STACK_FRAME_SCALE : recursion_limit;

    int ret = astfold_mod(mod, arena, state);
    assert(ret || PyErr_Occurred());

    /* Check that the recursion depth counting balanced correctly */
    if (ret && state->recursion_depth != starting_recursion_depth) {
        PyErr_Format(PyExc_SystemError,
            "AST optimizer recursion depth mismatch (before=%d, after=%d)",
            starting_recursion_depth, state->recursion_depth);
        return 0;
    }

    return ret;
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

astfold_mod 函数是 AST 优化的核心函数，对body，stmt，expr等节点进行优化。其实现如下：

static int
astfold_mod(mod_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
    switch (node_->kind) {
    case Module_kind:
        CALL(astfold_body, asdl_seq, node_->v.Module.body);
        break;
    case Interactive_kind:
        CALL_SEQ(astfold_stmt, stmt, node_->v.Interactive.body);
        break;
    case Expression_kind:
        CALL(astfold_expr, expr_ty, node_->v.Expression.body);
        break;
    // The following top level nodes don't participate in constant folding
    case FunctionType_kind:
        break;
    // No default case, so the compiler will emit a warning if new top level
    // compilation nodes are added without being handled here
    }
    return 1;
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

如果语法树节点的类型是 Module_kind，那么调用 astfold_body 函数。如果语法树的类型是 Interactive_kind，那么调用 astfold_stmt 函数。如果语法树的类型是 Expression_kind，那么调用 astfold_expr 函数。

语法树的优化是通过将语法树的节点进行折叠（折叠无效，累赘，重复，可编译时计算的节点），从而减少语法树的节点，最终达到减少执行指令的目的。

astfold_body body代码块折叠优化

static int
astfold_body(asdl_stmt_seq *stmts, PyArena *ctx_, _PyASTOptimizeState *state)
{
    int docstring = _PyAST_GetDocString(stmts) != NULL;
    CALL_SEQ(astfold_stmt, stmt, stmts);
    if (!docstring && _PyAST_GetDocString(stmts) != NULL) {
        stmt_ty st = (stmt_ty)asdl_seq_GET(stmts, 0);
        asdl_expr_seq *values = _Py_asdl_expr_seq_new(1, ctx_);
        if (!values) {
            return 0;
        }
        asdl_seq_SET(values, 0, st->v.Expr.value);
        expr_ty expr = _PyAST_JoinedStr(values, st->lineno, st->col_offset,
                                        st->end_lineno, st->end_col_offset,
                                        ctx_);
        if (!expr) {
            return 0;
        }
        st->v.Expr.value = expr;
    }
    return 1;
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

astfold_body 函数的作用是对语法树的 body 进行优化，如果 body 中包含 docstring，那么调用 _PyAST_GetDocString 函数获取 docstring；然后调用 astfold_stmt 函数对 body 进行优化；最后，如果 body 中包含 docstring，那么调用 _PyAST_JoinedStr 函数对 docstring 进行优化。

body里面的每一个节点都会调用 astfold_stmt 函数进行优化。这是一个递归的过程，直到所有的节点都被优化。

astfold_stmt 语句折叠优化

static int
astfold_stmt(stmt_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
    if (++state->recursion_depth > state->recursion_limit) {
        PyErr_SetString(PyExc_RecursionError,
                        "maximum recursion depth exceeded during compilation");
        return 0;
    }
    switch (node_->kind) {
    case FunctionDef_kind:
        CALL(astfold_arguments, arguments_ty, node_->v.FunctionDef.args);
        CALL(astfold_body, asdl_seq, node_->v.FunctionDef.body);
        CALL_SEQ(astfold_expr, expr, node_->v.FunctionDef.decorator_list);
        if (!(state->ff_features & CO_FUTURE_ANNOTATIONS)) {
            CALL_OPT(astfold_expr, expr_ty, node_->v.FunctionDef.returns);
        }
        break;
        // ...
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

astfold_stmt 函数的作用是对语法树的 stmt 进行优化，如果 stmt 的类型是 FunctionDef_kind，那么调用 astfold_arguments 函数对 args 进行优化，调用 astfold_body 函数对 body 进行优化，调用 astfold_expr 函数对 decorator_list 进行优化，如果没有 CO_FUTURE_ANNOTATIONS，那么调用 astfold_expr 函数对 returns 进行优化。

stmt 的优化是通过节点类型的数据结构进行逐步优化，这也是一个递归的过程，直到所有的节点都被优化。例如 if 包含 test, body, orelse 三个节点，test是一个表达式，而body和orelse是一个语句序列（stmt），所以需要分别调用 astfold_expr 和 astfold_stmt 函数进行优化。

case If_kind:
    CALL(astfold_expr, expr_ty, node_->v.If.test);
    CALL_SEQ(astfold_stmt, stmt, node_->v.If.body);
    CALL_SEQ(astfold_stmt, stmt, node_->v.If.orelse);
    break;

1
2
3
4
5
6

astfold_expr 表达式折叠优化

static int
astfold_expr(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
{
    if (++state->recursion_depth > state->recursion_limit) {
        PyErr_SetString(PyExc_RecursionError,
                        "maximum recursion depth exceeded during compilation");
        return 0;
    }
    switch (node_->kind) {
    case BoolOp_kind:   // bool
        CALL_SEQ(astfold_expr, expr, node_->v.BoolOp.values);
        break;
    case BinOp_kind: // bin?
        CALL(astfold_expr, expr_ty, node_->v.BinOp.left);
        CALL(astfold_expr, expr_ty, node_->v.BinOp.right);
        CALL(fold_binop, expr_ty, node_);
        break;
    // ...
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

astfold_expr 函数的作用是对语法树的 expr（表达式）进行优化，如果 expr 的类型是 BinOp_kind，那么调用 astfold_expr 函数对 left 和 right 进行优化，然后调用 fold_binop 函数对 expr 进行优化。astfold_expr 是一个递归的操作（我们可以不用关心这里的递归），最终调用 fold_binop 函数对 expr 进行优化。

简单了解一下 fold_binop 的实现

static int
fold_binop(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
{
    expr_ty lhs, rhs;
    lhs = node->v.BinOp.left;
    rhs = node->v.BinOp.right;
    if (lhs->kind != Constant_kind || rhs->kind != Constant_kind) {
        return 1;
    }

    PyObject *lv = lhs->v.Constant.value;
    PyObject *rv = rhs->v.Constant.value;
    PyObject *newval = NULL;

    switch (node->v.BinOp.op) {
    case Add:
        newval = PyNumber_Add(lv, rv);
        break;
        // ... 其他操作符
    }
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

fold_binop 是对表达式的优化，如果表达式的左右节点都是常量，那么就可以对表达式进行折叠优化。如 1 + 2，可以直接折叠成 3。这个操作是通过调用 PyNumber_Add 函数来实现数值的相加，那么最终的结果就是一个常量。

接下来我们验证这个优化，可以通过 dis 模块来查看 Python 代码的字节码。

import dis
code = """
a = 1 + 2
"""
dis.dis(code)

1
2
3
4
5
6

可以看到 1 + 2 被优化成了 3。调用 a 变成了 LOAD_CONST，将3加载到栈中，然后调用 STORE_NAME 将 3 存储到 a 中。输出字节码如下：

  2           0 LOAD_CONST               0 (3)
              2 STORE_NAME               0 (a)
              4 LOAD_CONST               1 (None)
              6 RETURN_VALUE

1
2
3
4
5

以上就是“Python内核源码解析：初步了解语法树AST优化的过程”的全部内容，希望对你有所帮助。
​
​关于Python技术储备

学好 Python 不论是就业还是做副业赚钱都不错，但要学会 Python 还是要有一个学习规划。最后大家分享一份全套的 Python 学习资料，给那些想学习 Python 的小伙伴们一点帮助！

一、Python所有方向的学习路线

Python所有方向的技术点做的整理，形成各个领域的知识点汇总，它的用处就在于，你可以按照上面的知识点去找对应的学习资源，保证自己学得较为全面。

二、Python必备开发工具

三、Python视频合集

观看零基础学习视频，看视频学习是最快捷也是最有效果的方式，跟着视频中老师的思路，从基础到深入，还是很容易入门的。

四、实战案例

光学理论是没用的，要学会跟着一起敲，要动手实操，才能将自己的所学运用到实际当中去，这时候可以搞点实战案例来学习。

五、Python练习题

检查学习结果。

六、面试资料

我们学习Python必然是为了找到高薪的工作，下面这些面试题是来自阿里、腾讯、字节等一线互联网大厂最新的面试资料，并且有阿里大佬给出了权威的解答，刷完这一套面试资料相信大家都能找到满意的工作。

最后祝大家天天进步！！

上面这份完整版的Python全套学习资料已经上传至CSDN官方，朋友如果需要可以直接微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】。