跳至主要内容

產生一個產生器

為你自己學 Python

產生器（Generator）在 Python 裡是一個很有趣的東西，它可以讓你一次產生一個值，而不是一口氣產生所有的值。這樣的特性讓產生器在處理大量資料或是無限的資料集合時非常有用。在 Python 裡要做出產生器有好幾種方法，其中一種就是透過 yield 關鍵字來定義一個產生器函數，例如：

def three_numbers():
    yield 520
    yield 1450
    yield 9527

nums = three_numbers()

如果檢視 nums，你發現它就是一個產生器物件：

>>> type(nums)
<class 'generator'>

之後就可以用內建函數 next() 來拿下一個值，直到沒有值可以拿為止。這個章節我們來看看 Python 的產生器是怎麼實作的。

產生器類別​

先看一下這段程式碼的 Bytecode，先看上半段：

0           0 RESUME                   0

1           2 LOAD_CONST               0 (<code object three_numbers>)
            4 MAKE_FUNCTION            0
            6 STORE_NAME               0 (three_numbers)

7           8 PUSH_NULL
           10 LOAD_NAME                0 (three_numbers)
           12 CALL                     0
           20 STORE_NAME               1 (nums)
           22 RETURN_CONST             1 (None)

這幾個指令我們之前都看過，就跟建立一般的函數差不多。不過下半段就不是那麼回事了：

1           0 RETURN_GENERATOR
            2 POP_TOP
            4 RESUME                   0

2           6 LOAD_CONST               1 (520)
            8 YIELD_VALUE              1
           10 RESUME                   1
           12 POP_TOP

3          14 LOAD_CONST               2 (1450)
           16 YIELD_VALUE              1
           18 RESUME                   1
           20 POP_TOP

4          22 LOAD_CONST               3 (9527)
           24 YIELD_VALUE              1
           26 RESUME                   1
           28 POP_TOP
           30 RETURN_CONST             0 (None)
      >>   32 CALL_INTRINSIC_1         3 (INTRINSIC_STOPITERATION_ERROR)
           34 RERAISE                  1
ExceptionTable:
  4 to 30 -> 32 [0] lasti

哇！有好幾個沒看過的指令，而且好像還繞來繞去的，避免大家沒耐心不想往下看，先說結論：

  • RETURN_GENERATOR 這個指令滿好猜的，這會產生一個產生器物件。
  • YIELD_VALUE 把產生器目前的值傳回給呼叫者，並且暫停函數的執行，例如上面範例裡的 yield 1 就是把 1 傳給呼叫這個函數的傢伙，然後本身這個函數暫停執行。
  • RESUME 讓這個產生器恢復執行，並從 YIELD_VALUE 指令暫停的位置繼續執行，這通常發生在呼叫內建函數 next() 或是產生器本身的 .send() 方法的時候。

我們先從最開頭的 RETURN_GENERATOR 指令開始看看：

檔案：Python/bytecodes.c
inst(RETURN_GENERATOR, (--)) {
    // ... 略 ...
    PyFunctionObject *func = (PyFunctionObject *)frame->f_funcobj;
    PyGenObject *gen = (PyGenObject *)_Py_MakeCoro(func);
    // ... 略 ...
}

這個指令會從堆疊上取得函數物件，然後透過 _Py_MakeCoro() 這個函數來建立一個產生器物件，裡的 Coro 是 Coroutine 的縮寫。在 _Py_MakeCoro() 函數裡會呼叫 make_gen() 函數來建立一個產生器物件：

檔案：Objects/genobject.c
static PyObject *
make_gen(PyTypeObject *type, PyFunctionObject *func)
{
    PyCodeObject *code = (PyCodeObject *)func->func_code;
    int slots = _PyFrame_NumSlotsForCodeObject(code);
    PyGenObject *gen = PyObject_GC_NewVar(PyGenObject, type, slots);
    if (gen == NULL) {
        return NULL;
    }
    gen->gi_frame_state = FRAME_CLEARED;
    gen->gi_weakreflist = NULL;
    gen->gi_exc_state.exc_value = NULL;
    gen->gi_exc_state.previous_item = NULL;
    assert(func->func_name != NULL);
    gen->gi_name = Py_NewRef(func->func_name);
    assert(func->func_qualname != NULL);
    gen->gi_qualname = Py_NewRef(func->func_qualname);
    _PyObject_GC_TRACK(gen);
    return (PyObject *)gen;
}

簡單的說，這過程就是把一個函數物件轉換成產生器。不過這講法也太簡化了，這裡會先計算這個函數物件需要多少的空間，然後透過 PyObject_GC_NewVar() 來建立一個產生器物件。產生器本身會有好幾種狀態：

檔案：Include/internal/pycore_frame.h
typedef enum _framestate {
    FRAME_CREATED = -2,
    FRAME_SUSPENDED = -1,
    FRAME_EXECUTING = 0,
    FRAME_COMPLETED = 1,
    FRAME_CLEARED = 4
} PyFrameState;

剛剛生出來的產生器的狀態是 FRAME_CLEARED，表示這個產生器目前還沒有被執行過。

是說，大家是否有聽說過使用產生器可以比較省記憶體的說法，這個說法基本上是對的，但看到這裡不知道大家有沒有發現，其實產生器還是需要記憶體的，而且可能還會需要多一些的空間來存放一些狀態。產生器的主要優勢在於它的延遲計算（Lazy Evaluation）特性，它只在需要時才生成下一個值，而不是一次生成所有的值。對於比較大的集合體或是無限的集合體來說，產生器的確可以節省記憶體，它不需要一次性把所有的元素全部展開放在記憶體裡，但對於比較小的集合體來說就不一定了，產生器可能實際上比一般的串列或 Tuple 還需要額外的記憶體來存放狀態。

我們來看看產生器 PyGenObject 的結構：

typedef struct {
    _PyGenObject_HEAD(gi)
} PyGenObject;

巨集展開後會變成：

typedef struct {
    PyObject ob_base;
    PyObject *gi_weakreflist;
    PyObject *gi_name;
    PyObject *gi_qualname;

    _PyErr_StackItem gi_exc_state;
    PyObject *gi_origin_or_finalizer;

    char gi_hooks_inited;
    char gi_closed;
    char gi_running_async;

    int8_t gi_frame_state;

    PyObject *gi_iframe[1];
} PyGenObject;

這裡 gi_name 跟 gi_qualname 是產生器的名字，gi_frame_state 是表示目前的執行狀態，而 gi_exc_state 是用來放例外處理的東西，這個待會再看看實際的範例會比較清楚。最後面的 gi_iframe[1] 之前我們也看過這種寫法，這是個彈性陣列成員的寫法。

yield 讓一下！​

在範例面的 Bytecode 指令裡面有一個 YIELD_VALUE 1 指令，從字面猜大概跟 yield 關鍵字有點關係，來看看這在做什麼事：

檔案：Python/bytecodes.c
inst(YIELD_VALUE, (retval -- unused)) {
    assert(frame != &entry_frame);
    PyGenObject *gen = _PyFrame_GetGenerator(frame);
    gen->gi_frame_state = FRAME_SUSPENDED;
    _PyFrame_SetStackPointer(frame, stack_pointer - 1);
    tstate->exc_info = gen->gi_exc_state.previous_item;
    gen->gi_exc_state.previous_item = NULL;
    _Py_LeaveRecursiveCallPy(tstate);
    _PyInterpreterFrame *gen_frame = frame;
    frame = cframe.current_frame = frame->previous;
    gen_frame->previous = NULL;
    _PyFrame_StackPush(frame, retval);
    goto resume_frame;
}

看起來是從當前的 Frame 取得產生器物件，先把它的狀態設定成暫停狀態的 FRAME_SUSPENDED，因為 yield 關鍵字會回傳一個值，所以這裡需要調整一下堆疊的指標。在中間還做了一些例外處理的事，這待會再看，最後的 _PyFrame_StackPush(frame, retval) 會把回值設定給前一個 Frame，也就是呼叫這個函數的 Frame 裡。這就是前面說的 YIELD_VALUE 指令會把值傳回給呼叫這個函數的傢伙，然後暫停這個函數的執行。

回頭看一下這兩行，這是在處理例外狀況：

tstate->exc_info = gen->gi_exc_state.previous_item;
gen->gi_exc_state.previous_item = NULL;

這兩行在處理例外狀態的轉移，確保萬一發生例外的時候能夠正確被處理。什麼意思？我舉個例子：

def simple_generaotr():
    try:
        yield 100
        raise ValueError("嘿嘿嘿")
    except:
        yield "Hey"

gen = simple_generaotr()
print(next(gen))  # 印出 100
print(next(gen))  # 這會印出什麼？

第一次執行 next() 函數會印出 100 這很正常，照理說這個產生器已經沒有更多值可以再丟出來了，但是再執行一次 next() 函數卻會印出 "Hey"，這是錯誤處理的 except 區塊 yield 出來的。這裡發生什麼事？

第一次的 next()，產生器開始動起來，遇到 yield 100，暫停原本 simple_generator 函數並回傳 100。此時產生器的狀態變為暫停（FRAME_SUSPENDED），而且還沒走到 raise 那一行。第二次執行 next() 函數的時候，產生器繼續動起來，程式從之前暫停的位置繼續執行。執行到 raise ValueError("嘿嘿嘿") 引發例外。產生器抓到這個例外後進到 except 區塊，在這個 except 區塊中產生器 yield 出 "Hey"，程式再次暫停。

產生器會把例外狀態保存在 previous_item 裡，然後繼續執行下去，以便在產生器恢復時能夠正確的處理例外。

下面一位！​

要取得產生器的下一個值，可使用內建函數 next()，但這個函數是怎麼實作的呢？

檔案：Python/bltinmodule.c
static PyObject *
builtin_next_impl(PyObject *module, PyObject *iterator,
                  PyObject *default_value)
{
    PyObject *res;

    // ... 錯誤處理 ...

    res = (*Py_TYPE(iterator)->tp_iternext)(iterator);
    if (res != NULL) {
        return res;
    } else if (default_value != NULL) {
        if (PyErr_Occurred()) {
            if(!PyErr_ExceptionMatches(PyExc_StopIteration))
                return NULL;
            PyErr_Clear();
        }
        return Py_NewRef(default_value);
    } else if (PyErr_Occurred()) {
        return NULL;
    } else {
        PyErr_SetNone(PyExc_StopIteration);
        return NULL;
    }
}

看起來是找 tp_iternext 成員，追看看 PyGen_Type 的定義：

檔案：Objects/genobject.c
static PyObject *
gen_iternext(PyGenObject *gen)
{
    PyObject *result;
    assert(PyGen_CheckExact(gen) || PyCoro_CheckExact(gen));
    if (gen_send_ex2(gen, NULL, &result, 0, 0) == PYGEN_RETURN) {
        if (result != Py_None) {
            _PyGen_SetStopIterationValue(result);
        }
        Py_CLEAR(result);
    }
    return result;
}

重點應該就是這裡的 gen_send_ex2() 函數了，在往下看這個函數的實作之前，先看看這個判斷：

if (result != Py_None) {
    _PyGen_SetStopIterationValue(result);
}

這好像有點不太直覺，為什麼當 result 不是 Py_None 的時候就呼叫看起來要丟出 StopIteration 的函數呢？先從外面的 if 判斷來看：

if (gen_send_ex2(gen, NULL, &result, 0, 0) == PYGEN_RETURN) {
    // ... 略 ...
}

gen_send_ex2 函數會執行傳進去的產生器的 Frame，如果 Frame 執行完畢會回傳 PYGEN_RETURN，並且將結果設定在 result 變數。一般產生器應該不需要主動寫 return，而是讓它自己拿到不能拿而產生 StopIteration 例外就會自動結束。但是如果產生器函數中裡有明確的寫了 return 的話，像這樣：

def simple_gen():
    yield 100
    return 9999

在執行 gen_send_ex2() 函數的過程中，result 會被設定成 9999。我們進到 REPL 試用一下：

>>> s = simple_gen()
>>> next(s)
100

第一次印出 100 沒問題，不過這時候產生器裡的值都被拿完了，如果再執行一次 next() 函數應該會丟出 StopIteration 例外，但仔細看結果：

>>> next(s)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration: 9999

在 StopIteration 例外裡的值就是 9999。這就是 _PyGen_SetStopIterationValue(result) 這個函數幹的好事。

所以 gen_send_ex2() 函數主要在做什麼事？

  • 根據目前不同的狀態，例如 FRAME_CREATED、FRAME_EXECUTING 或是 FRAME_COMPLETED 做相對應的處理。
  • 開始執行的時候，會把狀態設定成 FRAME_EXECUTING 並透過 _PyEval_EvalFrame() 函數來執行產生器的 Frame。

不得不說，這個產生器用起來可能很簡單，但是背後的實作卻是有點複雜，看的有點暈.. orz

工商服務

想學 Python 嗎？我教你啊 :)

想要成為軟體工程師嗎？這不是條輕鬆的路，除了興趣之外，還需要足夠的決心、設定目標並持續學習，我們的ASTROCamp 軟體工程師培訓營提供專業的前後端課程培訓，幫助你在最短時間內建立正確且扎實的軟體開發技能，有興趣而且不怕吃苦的話不妨來試試看！