技術文章 » Python 中生成器的原理

Python 中生成器的原理

2022-07-10 06:00:35

生成器的使用

在 Python 中，如果一個函數定義的內部使用了 yield 關鍵字，那麼在執行函數的時候返回的是一個生成器，而不是常規函數的返回值。

我們先來看一個常規函數的定義，下面的函數 f() 通過 return 語句返回 1，那麼 print 列印的就是數位 1。

def f():
    return 1
print(f())

如果我們將上面的 return 改成 yield，也就是下面這樣

def f():
    yield 1
    yield 2
g = f()
print(g)
print(next(g))
print(next(g))
print(next(g))

最終的輸出如下，呼叫函數 f() 得到的是一個生成器（generator）物件 g，通過 Python 內建的 next() 函數可以驅動生成器往下執行，每呼叫一次 next() 函數，生成器就會執行到下一個 yield 語句處，並將 yield 語句中的表示式返回，當沒有更多 yield 語句時繼續執行 next() 函數會觸發 StopIteration 異常。

<generator object f at 0x10c963c50>
1
2
Traceback (most recent call last):
  File "<string>", line 8, in <module>
StopIteration

當然更優雅的使用生成器的方式是使用 for 迴圈，如下所示，會依次列印 1、2，並且不會丟擲 StopIteration 異常，因為本質上生成器也是一種迭代器，所以可以用 for 迴圈遍歷。另外，生成器也可以用生成器表示式如 g = (i for i "hello world") 來建立，這不是本文重點，就不詳細介紹了。

def f():
    yield 1
    yield 2
for i in f():
    print(i)

生成器的原理

要理解 Python 中生成器的原理其實就是要搞清楚下面兩個問題

  • 呼叫包含 yield 語句的函數為什麼同普通函數不一樣，返回的是一個生成器物件，而不是普通的返回值
  • next() 函數驅動生成器執行的時候為什麼可以在函數體中返回 yield 後面的表示式後暫停，下次呼叫 next() 的時候可以從暫停處繼續執行

這兩個問題都跟 Python 程式執行機制有關。Python 程式碼首先會經過 Python 編譯器編譯成位元組碼，然後由 Python 直譯器解釋執行，機制上跟其他直譯語言一樣。Python 編譯器和直譯器配合，就能完成上面兩個問題中的功能，這在編譯型語言中很難做到。像 C、Golang 會編譯成機器語言，函數呼叫通過 CALL 指令來完成，被呼叫的函數中遇到 RET 指令就會返回，釋放掉被呼叫函數的棧幀，無法在中途返回，下次繼續執行。

雖然作業系統線上程切換的時候也會中斷正在執行的函數，再次切換回來的時候繼續執行，但是被中斷的函數在切換的時候並沒有返回值產生，這點與 Python 生成器是不同的，不要混淆了。

下面我們具體來看一下 Python 是如何解決上面兩個問題的（基於 CPython 3.10.4）。

生成器的建立

Python 編譯器在編譯 Python 程式碼的時候分為詞法分析、語法分析、語意分析和位元組碼生成這幾個階段，在進行語意分析的時候有一項重要的工作是構建符號表，主要用於確定各個變數的作用域，順帶做了一件跟生成器相關的事，也就是在分析過程中如果遇到了 yield 語句就將當前程式碼塊的符號表標記為是生成器。
相關原始碼如下

static int
symtable_visit_expr(struct symtable *st, expr_ty e)
{
    if (++st->recursion_depth > st->recursion_limit) {
        PyErr_SetString(PyExc_RecursionError, "maximum recursion depth exceeded during compilation");
        VISIT_QUIT(st, 0);
    }
    switch (e->kind) {
    ...
    case Yield_kind:
        if (!symtable_raise_if_annotation_block(st, "yield expression", e)) {
            VISIT_QUIT(st, 0);
        }
        if (e->v.Yield.value)
            VISIT(st, expr, e->v.Yield.value);
        st->st_cur->ste_generator = 1; // 如果遇到了 yield 語句，就將 ste_generator 標誌位置 1
        if (st->st_cur->ste_comprehension) {
            return symtable_raise_if_comprehension_block(st, e);
        }
        break; 
    ...
    }
    ...
}

最後在生成位元組碼的時候，會根據符號表的屬性計算位元組碼物件的標誌位，如果 ste_generator 為 1，就將位元組碼物件的標誌位加上 CO_GENERATOR，相關原始碼如下

static int compute_code_flags(struct compiler *c)
{
    PySTEntryObject *ste = c->u->u_ste;
    int flags = 0;
    if (ste->ste_type == FunctionBlock) {
        flags |= CO_NEWLOCALS | CO_OPTIMIZED;
        if (ste->ste_nested)
            flags |= CO_NESTED;
        if (ste->ste_generator && !ste->ste_coroutine)
            flags |= CO_GENERATOR; // 如果符號表中 ste_generator 標誌位為 1，就將 code 物件的 flags 加上 CO_GENERATOR 
        if (!ste->ste_generator && ste->ste_coroutine)
            flags |= CO_COROUTINE;
        if (ste->ste_generator && ste->ste_coroutine)
            flags |= CO_ASYNC_GENERATOR;
        if (ste->ste_varargs)
            flags |= CO_VARARGS;
        if (ste->ste_varkeywords)
            flags |= CO_VARKEYWORDS;
    }
    ...
    return flags;
}

最終 g = f() 會生成下面的位元組碼

0 LOAD_NAME                0 (f)
2 CALL_FUNCTION            0
4 STORE_NAME               1 (g)

Python 直譯器會執行 CALL_FUNCTION 指令，將函數 f() 的呼叫返回值賦值給 g。CALL_FUNCTION 指令在執行的時候會先檢查對應的位元組碼物件的 co_flags 標誌，如果包含 CO_GENERATOR 標誌就返回一個生成器物件。相關原始碼簡化後如下

PyObject *
_PyEval_Vector(PyThreadState *tstate, PyFrameConstructor *con, PyObject *locals, PyObject* const* args, size_t argcount, PyObject *kwnames)
{
    PyFrameObject *f = _PyEval_MakeFrameVector(tstate, con, locals, args, argcount, kwnames);
    if (f == NULL) {
        return NULL;
    }
    // 如果 code 物件有 CO_GENERATOR 標誌位，就直接返回一個生成器物件
    if (((PyCodeObject *)con->fc_code)->co_flags & CO_GENERATOR) { 
        return PyGen_NewWithQualName(f, con->fc_name, con->fc_qualname);
    }
    ...
}

可以看到編譯器和直譯器的配合，讓生成器得以建立。

生成器的執行

Python 直譯器用軟體的方式模擬了 CPU 執行指令的流程，每個程式碼塊（模組、類、函數）在執行的時候，直譯器首先為其建立一個棧幀，主要用於儲存程式碼塊執行時所需要的各種變數的值，同時指向呼叫方的棧幀，使得當前程式碼塊執行結束後能夠順利返回到呼叫方繼續執行。與物理棧幀不同的是，Python 直譯器中的棧幀是在程序的堆區建立的，如此一來棧幀就完全是直譯器控制的，即使直譯器自己的物理棧幀結束了，只要不主動釋放，程式碼塊的棧幀依然會存在。

執行位元組碼的主邏輯在 _PyEval_EvalFrameDefault 函數中，其中有個 for 迴圈依次取出程式碼塊中的各條指令並執行，next(g) 在執行的時候經過層層的呼叫最終也會走到這個迴圈裡，其中跟生成器相關的原始碼簡化後如下

PyObject* _Py_HOT_FUNCTION _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
{
    ...
    for (;;) {
        opcode = _Py_OPCODE(*next_instr);
        switch (opcode) {
        case TARGET(YIELD_VALUE): {
            retval = POP(); // 將 yiled 後面的表示式的值賦給返回值 retval

            if (co->co_flags & CO_ASYNC_GENERATOR) {
                PyObject *w = _PyAsyncGenValueWrapperNew(retval);
                Py_DECREF(retval);
                if (w == NULL) {
                    retval = NULL;
                    goto error;
                }
                retval = w;
            }
            f->f_state = FRAME_SUSPENDED; // 設定當前棧幀為暫停狀態
            f->f_stackdepth = (int)(stack_pointer - f->f_valuestack);
            goto exiting; // 結束本次函數呼叫，返回上級函數
        }
        }
    }
    ...
}

可以看出 Python 直譯器在執行 yield 語句時會將 yield 後面的值作為返回值直接返回，同時設定當前棧幀為暫停狀態。由於這裡的棧幀是儲存在程序的堆區的，所以當這次對生成器的呼叫結束之後，其棧幀依然存在，各個變數的值依然儲存著，下次呼叫的時候可以繼續當前的狀態往下執行。

總結

本文介紹了 Python 中生成器的使用方法，然後介紹了 Python 程式碼的執行機制，並結合原始碼對生成器的工作原理做了介紹。Python 直譯器能實現生成器，主要是因為其是用軟體來模擬硬體的行為，既然是軟體，在實現的時候就可以新增很多功能，對直譯器的一頓魔改，在 Python 2.2 版本中就引進了生成器。