Python 中生成器的原理

生成器的使用

在 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 解释器中的栈帧是在进程的堆区创建的,如此一来栈帧就完全是解释器控制的,即使解释器自己的物理栈帧结束了,只要不主动释放,代码块的栈帧依然会存在。

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

发表评论

相关文章