Python yield與實現

Python yield與實現



yield的功能相似於return，可是不一樣之處在於它返回的是生成器。

生成器

生成器是經過一個或多個yield表達式構成的函數，每個生成器都是一個迭代器（可是迭代器不必定是生成器）。

若是一個函數包含yield關鍵字，這個函數就會變爲一個生成器。

生成器並不會一次返回全部結果，而是每次遇到yield關鍵字後返回相應結果，並保留函數當前的運行狀態，等待下一次的調用。

因爲生成器也是一個迭代器，那麼它就應該支持next方法來獲取下一個值。

基本操做

# 經過`yield`來建立生成器
def func():
   for i in xrange(10);
        yield i

# 經過列表來建立生成器
[i for i in xrange(10)]

# 調用以下
>>> f = func()
>>> f # 此時生成器尚未運行
<generator object func at 0x7fe01a853820>
>>> f.next() # 當i=0時，遇到yield關鍵字，直接返回
0
>>> f.next() # 繼續上一次執行的位置，進入下一層循環
1
...
>>> f.next()
9
>>> f.next() # 當執行完最後一次循環後，結束yield語句，生成StopIteration異常
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>>

除了next函數，生成器還支持send函數。該函數能夠向生成器傳遞參數。

>>> def func():
...     n = 0
...     while 1:
...         n = yield n #能夠經過send函數向n賦值
... 
>>> f = func()
>>> f.next() # 默認狀況下n爲0
0
>>> f.send(1) #n賦值1
1
>>> f.send(2)
2
>>>

應用

最經典的例子，生成無限序列。

常規的解決方法是，生成一個知足要求的很大的列表，這個列表須要保存在內存中，很明顯內存限制了這個問題。

def get_primes(start):
    for element in magical_infinite_range(start):
        if is_prime(element):
            return element

若是使用生成器就不須要返回整個列表，每次都只是返回一個數據，避免了內存的限制問題。

def get_primes(number):
    while True:
        if is_prime(number):
            yield number
        number += 1

生成器源碼分析

生成器的源碼在Objects/genobject.c。

調用棧

在解釋生成器以前，須要講解一下Python虛擬機的調用原理。

Python虛擬機有一個棧幀的調用棧，其中棧幀的是PyFrameObject，位於Include/frameobject.h。

typedef struct _frame {
    PyObject_VAR_HEAD
    struct _frame *f_back;  /* previous frame, or NULL */
    PyCodeObject *f_code;   /* code segment */
    PyObject *f_builtins;   /* builtin symbol table (PyDictObject) */
    PyObject *f_globals;    /* global symbol table (PyDictObject) */
    PyObject *f_locals;     /* local symbol table (any mapping) */
    PyObject **f_valuestack;    /* points after the last local */
    /* Next free slot in f_valuestack.  Frame creation sets to f_valuestack.
       Frame evaluation usually NULLs it, but a frame that yields sets it
       to the current stack top. */
    PyObject **f_stacktop;
    PyObject *f_trace;      /* Trace function */

    /* If an exception is raised in this frame, the next three are used to
     * record the exception info (if any) originally in the thread state.  See
     * comments before set_exc_info() -- it's not obvious.
     * Invariant:  if _type is NULL, then so are _value and _traceback.
     * Desired invariant:  all three are NULL, or all three are non-NULL.  That
     * one isn't currently true, but "should be".
     */
    PyObject *f_exc_type, *f_exc_value, *f_exc_traceback;

    PyThreadState *f_tstate;
    int f_lasti;        /* Last instruction if called */
    /* Call PyFrame_GetLineNumber() instead of reading this field
       directly.  As of 2.3 f_lineno is only valid when tracing is
       active (i.e. when f_trace is set).  At other times we use
       PyCode_Addr2Line to calculate the line from the current
       bytecode index. */
    int f_lineno;       /* Current line number */
    int f_iblock;       /* index in f_blockstack */
    PyTryBlock f_blockstack[CO_MAXBLOCKS]; /* for try and loop blocks */
    PyObject *f_localsplus[1];  /* locals+stack, dynamically sized */
} PyFrameObject;

棧幀保存了給出代碼的的信息和上下文，其中包含最後執行的指令，全局和局部命名空間，異常狀態等信息。f_valueblock保存了數據，b_blockstack保存了異常和循環控制方法。

舉一個例子來講明，

def foo():
    x = 1
    def bar(y):
        z = y + 2  # <--- (3) ... and the interpreter is here.
        return z
    return bar(x)  # <--- (2) ... which is returning a call to bar ...
foo()              # <--- (1) We're in the middle of a call to foo ...

那麼，相應的調用棧以下，一個py文件，一個類，一個函數都是一個代碼塊，對應者一個Frame，保存着上下文環境以及字節碼指令。

c   ---------------------------
a  | bar Frame                 | -> block stack: []
l  |     (newest)              | -> data stack: [1, 2]
l   ---------------------------
   | foo Frame                 | -> block stack: []
s  |                           | -> data stack: [<Function foo.<locals>.bar at 0x10d389680>, 1]
t   ---------------------------
a  | main (module) Frame       | -> block stack: []
c  |       (oldest)            | -> data stack: [<Function foo at 0x10d3540e0>]
k   ---------------------------

每個棧幀都擁有本身的數據棧和block棧，獨立的數據棧和block棧使得解釋器能夠中斷和恢復棧幀（生成器正式利用這點）。

Python代碼首先被編譯爲字節碼，再由Python虛擬機來執行。通常來講，一條Python語句對應着多條字節碼（因爲每條字節碼對應着一條C語句，而不是一個機器指令，因此不能按照字節碼的數量來判斷代碼性能）。

調用dis模塊能夠分析字節碼，

from dis import dis

dis(foo)

  5           0 LOAD_CONST               1 (1) # 加載常量1
              3 STORE_FAST               0 (x) # x賦值爲1

  6           6 LOAD_CONST               2 (<code object bar at 0x7f3cdee3a030, file "t1.py", line 6>) # 加載常量2
              9 MAKE_FUNCTION            0 # 建立函數
             12 STORE_FAST               1 (bar) 

  9          15 LOAD_FAST                1 (bar) 
             18 LOAD_FAST                0 (x)
             21 CALL_FUNCTION            1  # 調用函數
             24 RETURN_VALUE

其中，

第一行爲代碼行號；
第二行爲偏移地址；
第三行爲字節碼指令；
第四行爲指令參數；
第五行爲參數解釋。

生成器源碼分析

由了上面對於調用棧的理解，就能夠很容易的明白生成器的具體實現。

生成器的源碼位於object/genobject.c。

生成器的建立

PyObject *
PyGen_New(PyFrameObject *f)
{
    PyGenObject *gen = PyObject_GC_New(PyGenObject, &PyGen_Type); # 建立生成器對象
    if (gen == NULL) {
        Py_DECREF(f);
        return NULL;
    }
    gen->gi_frame = f; # 賦予代碼塊
    Py_INCREF(f->f_code); # 引用計數+1
    gen->gi_code = (PyObject *)(f->f_code);
    gen->gi_running = 0; # 0表示爲執行，也就是生成器的初始狀態
    gen->gi_weakreflist = NULL;
    _PyObject_GC_TRACK(gen); # GC跟蹤
    return (PyObject *)gen;
}

send與next

next與send函數，以下

static PyObject *
gen_iternext(PyGenObject *gen)
{
    return gen_send_ex(gen, NULL, 0);
}


static PyObject *
gen_send(PyGenObject *gen, PyObject *arg)
{
    return gen_send_ex(gen, arg, 0);
}

從上面的代碼中能夠看到，send和next都是調用的同一函數gen_send_ex，區別在因而否帶有參數。

static PyObject *
gen_send_ex(PyGenObject *gen, PyObject *arg, int exc)
{
    PyThreadState *tstate = PyThreadState_GET();
    PyFrameObject *f = gen->gi_frame;
    PyObject *result;

    if (gen->gi_running) { # 判斷生成器是否已經運行
        PyErr_SetString(PyExc_ValueError,
                        "generator already executing");
        return NULL;
    }
    if (f==NULL || f->f_stacktop == NULL) { # 若是代碼塊爲空或調用棧爲空，則拋出StopIteration異常
        /* Only set exception if called from send() */
        if (arg && !exc)
            PyErr_SetNone(PyExc_StopIteration);
        return NULL;
    }

    if (f->f_lasti == -1) { # f_lasti=1 表明首次執行
        if (arg && arg != Py_None) { # 首次執行不容許帶有參數
            PyErr_SetString(PyExc_TypeError,
                            "can't send non-None value to a "
                            "just-started generator");
            return NULL;
        }
    } else {
        /* Push arg onto the frame's value stack */
        result = arg ? arg : Py_None;
        Py_INCREF(result); # 該參數引用計數+1
        *(f->f_stacktop++) = result; # 參數壓棧
    }

    /* Generators always return to their most recent caller, not
     * necessarily their creator. */
    f->f_tstate = tstate;
    Py_XINCREF(tstate->frame);
    assert(f->f_back == NULL);
    f->f_back = tstate->frame;

    gen->gi_running = 1; # 修改生成器執行狀態
    result = PyEval_EvalFrameEx(f, exc); # 執行字節碼
    gen->gi_running = 0; # 恢復爲未執行狀態

    /* Don't keep the reference to f_back any longer than necessary.  It
     * may keep a chain of frames alive or it could create a reference
     * cycle. */
    assert(f->f_back == tstate->frame);
    Py_CLEAR(f->f_back);
    /* Clear the borrowed reference to the thread state */
    f->f_tstate = NULL;

    /* If the generator just returned (as opposed to yielding), signal
     * that the generator is exhausted. */
    if (result == Py_None && f->f_stacktop == NULL) {
        Py_DECREF(result);
        result = NULL;
        /* Set exception if not called by gen_iternext() */
        if (arg)
            PyErr_SetNone(PyExc_StopIteration);
    }

    if (!result || f->f_stacktop == NULL) {
        /* generator can't be rerun, so release the frame */
        Py_DECREF(f);
        gen->gi_frame = NULL;
    }

    return result;
}

字節碼的執行

PyEval_EvalFrameEx函數的功能爲執行字節碼並返回結果。

# 主要流程以下，
for (;;) {
   switch(opcode) { # opcode爲操做碼，對應着各類操做
        case NOP:
            goto  fast_next_opcode;
        ...
        ...
        case YIELD_VALUE: # 若是操做碼是yield
            retval = POP(); 
            f->f_stacktop = stack_pointer;
            why = WHY_YIELD;
            goto fast_yield; # 利用goto跳出循環
    }
}

fast_yield:
    ... 
return vetval; # 返回結果

舉一個例子，f_back上一個Frame，f_lasti上一次執行的指令的偏移量，

import sys
from dis import dis


def func():
    f = sys._getframe(0)
    print f.f_lasti
    print f.f_back
    yield 1

    print f.f_lasti
    print f.f_back
    yield 2


a = func()
dis(func)
a.next()
a.next()

結果以下，其中第三行的英文爲操做碼，對應着上面的opcode，每次switch都是在不一樣的opcode之間進行選擇。

6           0 LOAD_GLOBAL              0 (sys)
              3 LOAD_ATTR                1 (_getframe)
              6 LOAD_CONST               1 (0)
              9 CALL_FUNCTION            1
             12 STORE_FAST               0 (f)

  7          15 LOAD_FAST                0 (f)
             18 LOAD_ATTR                2 (f_lasti) 
             21 PRINT_ITEM          
             22 PRINT_NEWLINE       

  8          23 LOAD_FAST                0 (f)
             26 LOAD_ATTR                3 (f_back)
             29 PRINT_ITEM          
             30 PRINT_NEWLINE       

  9          31 LOAD_CONST               2 (1)
             34 YIELD_VALUE     # 此時操做碼爲YIELD_VALUE，直接跳轉上述goto語句，此時f_lasti爲當前指令，f_back爲當前frame
             35 POP_TOP             

 11          36 LOAD_FAST                0 (f)
             39 LOAD_ATTR                2 (f_lasti)
             42 PRINT_ITEM          
             43 PRINT_NEWLINE       

 12          44 LOAD_FAST                0 (f)
             47 LOAD_ATTR                3 (f_back)
             50 PRINT_ITEM          
             51 PRINT_NEWLINE       

 13          52 LOAD_CONST               3 (2)
             55 YIELD_VALUE         
             56 POP_TOP             
             57 LOAD_CONST               0 (None)
             60 RETURN_VALUE        
18
<frame object at 0x7fa75fcebc20> #和下面的frame相同，屬於同一個frame，也就是說在同一個函數（命名空間）內，frame是同一個。
39
<frame object at 0x7fa75fcebc20>