Python 中生成器的原理 - zikcheng

生成器的使用

在 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 版本中就引进了生成器。