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 */
<span class="hljs-comment">/* 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".
 */</span>
PyObject *f_exc_type, *f_exc_value, *f_exc_traceback;

PyThreadState *f_tstate;
<span class="hljs-type">int</span> f_lasti;		<span class="hljs-comment">/* Last instruction if called */</span>
<span class="hljs-comment">/* 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. */</span>
<span class="hljs-type">int</span> f_lineno;		<span class="hljs-comment">/* Current line number */</span>
<span class="hljs-type">int</span> f_iblock;		<span class="hljs-comment">/* index in f_blockstack */</span>
PyTryBlock f_blockstack[CO_MAXBLOCKS]; <span class="hljs-comment">/* for try and loop blocks */</span>
PyObject *f_localsplus[<span class="hljs-number">1</span>];	<span class="hljs-comment">/* locals+stack, dynamically sized */</span>

} 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

nextsend函数,如下

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);
}

从上面的代码中可以看到,sendnext都是调用的同一函数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;
<span class="hljs-title function_ invoke__">if</span> (gen<span class="hljs-punctuation">-&gt;</span>gi_running) { # 判断生成器是否已经运行
    <span class="hljs-title function_ invoke__">PyErr_SetString</span>(PyExc_ValueError,
                    <span class="hljs-string">"generator already executing"</span>);
    <span class="hljs-keyword">return</span> NULL;
}
<span class="hljs-title function_ invoke__">if</span> (f==NULL || f<span class="hljs-punctuation">-&gt;</span>f_stacktop == NULL) { # 如果代码块为空或调用栈为空,则抛出StopIteration异常
    <span class="hljs-comment">/* Only set exception if called from send() */</span>
    <span class="hljs-title function_ invoke__">if</span> (arg &amp;&amp; !exc)
        <span class="hljs-title function_ invoke__">PyErr_SetNone</span>(PyExc_StopIteration);
    <span class="hljs-keyword">return</span> NULL;
}

<span class="hljs-title function_ invoke__">if</span> (f<span class="hljs-punctuation">-&gt;</span>f_lasti == -<span class="hljs-number">1</span>) { # f_lasti=<span class="hljs-number">1</span> 代表首次执行
    <span class="hljs-title function_ invoke__">if</span> (arg &amp;&amp; arg != Py_None) { # 首次执行不允许带有参数
        <span class="hljs-title function_ invoke__">PyErr_SetString</span>(PyExc_TypeError,
                        <span class="hljs-string">"can't send non-None value to a "</span>
                        <span class="hljs-string">"just-started generator"</span>);
        <span class="hljs-keyword">return</span> NULL;
    }
} <span class="hljs-keyword">else</span> {
    <span class="hljs-comment">/* Push arg onto the frame's value stack */</span>
    result = arg ? arg : Py_None;
    <span class="hljs-title function_ invoke__">Py_INCREF</span>(result); # 该参数引用计数+<span class="hljs-number">1</span>
    *(f<span class="hljs-punctuation">-&gt;</span>f_stacktop++) = result; # 参数压栈
}

<span class="hljs-comment">/* Generators always return to their most recent caller, not
 * necessarily their creator. */</span>
f<span class="hljs-punctuation">-&gt;</span>f_tstate = tstate;
<span class="hljs-title function_ invoke__">Py_XINCREF</span>(tstate<span class="hljs-punctuation">-&gt;</span>frame);
<span class="hljs-title function_ invoke__">assert</span>(f<span class="hljs-punctuation">-&gt;</span>f_back == NULL);
f<span class="hljs-punctuation">-&gt;</span>f_back = tstate<span class="hljs-punctuation">-&gt;</span>frame;

gen<span class="hljs-punctuation">-&gt;</span>gi_running = <span class="hljs-number">1</span>; # 修改生成器执行状态
result = <span class="hljs-title function_ invoke__">PyEval_EvalFrameEx</span>(f, exc); # 执行字节码
gen<span class="hljs-punctuation">-&gt;</span>gi_running = <span class="hljs-number">0</span>; # 恢复为未执行状态

<span class="hljs-comment">/* 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. */</span>
<span class="hljs-title function_ invoke__">assert</span>(f<span class="hljs-punctuation">-&gt;</span>f_back == tstate<span class="hljs-punctuation">-&gt;</span>frame);
<span class="hljs-title function_ invoke__">Py_CLEAR</span>(f<span class="hljs-punctuation">-&gt;</span>f_back);
<span class="hljs-comment">/* Clear the borrowed reference to the thread state */</span>
f<span class="hljs-punctuation">-&gt;</span>f_tstate = NULL;

<span class="hljs-comment">/* If the generator just returned (as opposed to yielding), signal
 * that the generator is exhausted. */</span>
<span class="hljs-title function_ invoke__">if</span> (result == Py_None &amp;&amp; f<span class="hljs-punctuation">-&gt;</span>f_stacktop == NULL) {
    <span class="hljs-title function_ invoke__">Py_DECREF</span>(result);
    result = NULL;
    <span class="hljs-comment">/* Set exception if not called by gen_iternext() */</span>
    <span class="hljs-title function_ invoke__">if</span> (arg)
        <span class="hljs-title function_ invoke__">PyErr_SetNone</span>(PyExc_StopIteration);
}

<span class="hljs-title function_ invoke__">if</span> (!result || f<span class="hljs-punctuation">-&gt;</span>f_stacktop == NULL) {
    <span class="hljs-comment">/* generator can't be rerun, so release the frame */</span>
    <span class="hljs-title function_ invoke__">Py_DECREF</span>(f);
    gen<span class="hljs-punctuation">-&gt;</span>gi_frame = NULL;
}

<span class="hljs-keyword">return</span> 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

<span class="hljs-built_in">print</span> f.f_lasti
<span class="hljs-built_in">print</span> f.f_back
<span class="hljs-keyword">yield</span> <span class="hljs-number">2</span>

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>