深入理解 python 虛擬機器：生成器停止背後的魔法

一無是處的研究僧發表於2023-10-04
Python虛擬機

深入理解 python 虛擬機器：生成器停止背後的魔法

在本篇文章當中主要給大家介紹 Python 當中生成器的實現原理，尤其是生成器是如何能夠被停止執行，而且還能夠被恢復的，這是一個非常讓人疑惑的地方。因為這與我們通常使用的函式的直覺是相違背的，函式之後執行完成之後才會返回，而生成表面是函式的形式，但是這違背了我們正常的程式設計直覺。

深入理解生成器與函式的區別

為了從根本上建立對生成器的認識，我們首先就需要深入理解一下生成器和函式的區別。其實在從虛擬機器的層面來看，他們兩個都是物件，只不過一個是生成器物件，一個是函式物件。在 Python 當中，如果你在函式里面使用了 yield 語句，那麼你的這個函式在被呼叫的時候就不會被執行，而是會返回一個生成器物件。

>>> def bar():
...     print("before yield")
...     res = yield 1
...     print(f"{res = }")
...     print("after yield")
...     return "Return Value"
...
>>> generator = bar()
>>> generator
<generator object bar at 0x105267510>
>>> bar
<function bar at 0x10562fc40>
>>>

在 Python 當中有的物件是可以直接呼叫的，比如你自己的類如果實現了__call__方法的話，這個類生成的物件就是一個可呼叫物件，在 Python 當中一個最常見的可呼叫物件就是函式了，生成器和函式的區別之一就是，生成器不能夠直接被呼叫，而函式可以。

>>> generator()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'generator' object is not callable
>>>

在上面的程式碼當中我們要明確 bar 是一個函式，但是這個函式和正常的函式有一點區別，這個函式在被呼叫的時候不會直接執行程式碼，而是會返回一個生成器物件，因為在這個函式體當中使用了 yield 語句，我們稱這種函式為生成器函式 (generator function)，在 Python 當中你可以透過檢視一個函式的 co_flags 欄位檢視一個函式的屬性，如果這個欄位和 0x0020 進行 & 操作之後的結果大於 0，那麼就說明這個函式是一個生成器函式。

>>> (bar.__code__.co_flags & 0x0020) > 0
True
>>> bar.__code__.co_flags & 0x0020
32

從上面的程式碼當中我們可以看到 bar 就是一個生成器函式，除了上面的方法 Python 的標準庫也提供了方法去輔助我們進行判斷。

>>> import inspect
>>> inspect.isgeneratorfunction(bar)
True

上面的特性在 Python 程式進行編譯的時候，編譯器可以做到這一點，當發現一個函式當中存在類似 yield 的語句的時候就在函式的 co_flags 欄位當中和 0x0020 進行或操作，然後將這個值儲存在 co_flags 當中。

總之生成器和函式之間的關係為：生成器物件是透過呼叫生成器函式得到的，呼叫生成器函式的返回物件是生成器。

虛實交錯的時空魔法

首先我們需要了解的是，如果我們想讓一個生成器物件執行下去的話，我們可以使用 next 或者 send 函式，進行實現：

>>> next(generator)
before yield
1
>>> next(generator)
res = None
after yield
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration: Return Value

在 CPython 實現的虛擬機器當中，如果我們想要正確的使用 send 函式首先需要讓生成器物件執行到第一個 yield 語句，我們可以使用 next(generator) 或者 generator.send(None)。比如在上面的第一條語句當中執行 next(generator)，執行到語句 res = yield 1，但是這條語句還沒有執行完，需要我們呼叫 send 函式之後才能夠完成賦值操作，send 函式的引數會被賦值給變數 res 。當整個函式體執行完成之後虛擬機器就會丟擲 StopIteration 異常，並且將返回值儲存到 StopIteration 異常物件當中：

>>> generator = bar()
>>> next(generator)
before yield
1
>>> try:
...     generator.send("None")
... except StopIteration as e:
...     print(f"{e.value = }")
...
res = 'None'
after yield
e.value = 'Return Value'
>>>

上面的程式碼當中可以看到，我們正確的執行力我們在上面談到的生成器的使用方法，並且將生成器執行完成之後的返回值儲存到異常的 value 當中。

生成器內部實現原理

從上面的關於生成器的使用方式來看，生成器可以在函式執行到一半的時候停止，然後繼續恢復執行，為了實現這一點我們就需要有一種手段去儲存函式執行的狀態。但是我們需要儲存函式執行的那些狀態呢？最重要的兩點就是程式碼現在執行到什麼位置了，因為我們之後要繼續從下一條指令開始恢復執行，同時我們需要儲存虛擬機器的棧空間，就是在執行位元組碼的時候使用到的 valuestack，注意這不是棧幀，同時還有執行函式的區域性變數表，這裡主要是儲存一些區域性變數的。而這些東西都儲存在虛擬機器的棧幀當中了，這一點我們在前面的文章當中已經詳細介紹過了。

因此根據這些分析我們應該知道了，生成器裡面最重要的就是一個虛擬機器的棧幀資料結構了。一個生成器物件當中一定需要有一個虛擬機器的棧幀，在 CPython 的實現當中，生成器物件的資料結構如下：

typedef struct
{
    /* The gi_ prefix is intended to remind of generator-iterator. */
    PyObject ob_base;
    struct _frame *gi_frame;
    char gi_running;
    PyObject *gi_code;
    PyObject *gi_weakreflist;
    PyObject *gi_name;
    PyObject *gi_qualname;
    _PyErr_StackItem gi_exc_state;
} PyGenObject;
  • gi_frame: 這個欄位就是表示生成器所擁有的棧幀。
  • gi_running: 表示生成器是否在執行。
  • gi_code: 表示對應生成器函式的程式碼（位元組碼）。
  • gi_weakreflist: 用於儲存這個棧幀物件儲存的弱引用物件。
  • gi_name 和 gi_qualname 都是表示生成器的名字，後者更加詳細。
  • gi_exc_state: 用於儲存執行生成器程式碼之前的程式狀態，因為之前的程式碼可能已經產生一些異常了，這個主要用於儲存之前的程式狀態，待生成器返回之後就進行恢復。
class A:
	def hello(self):
		yield 1


if __name__ == '__main__':
	g = A().hello()
	print(g.__name__)
	print(g.__qualname__)

上面的程式輸出結果為:

hello
A.hello

生成器對應的位元組碼行為

我們透過下面的例子來分析一下，生成器 yield 對應的位元組碼：

>>> import dis
>>> def hello():
...     yield 1
...     yield 2
...
>>> dis.dis(hello)
  2           0 LOAD_CONST               1 (1)
              2 YIELD_VALUE
              4 POP_TOP

  3           6 LOAD_CONST               2 (2)
              8 YIELD_VALUE
             10 POP_TOP
             12 LOAD_CONST               0 (None)
             14 RETURN_VALUE

在上面的程式當中只有和生成器相關的位元組碼為 YIELD_VALUE，在載入完常量 1 之後就會執行 YIELD_VALUE 指令，虛擬機器在執行完 yield 指令之後，就會直接返回，此時虛擬機器的狀態——valuestack 和當前指令執行的位置（在上面的這個例子當中就是 4）都會被儲存到虛擬機器棧幀當中，當下一次執行生成器的程式碼的時候就會直接從 POP_TOP 指令直接執行。

我們再來看一下另外一個比較重要的指令 YIELD_FROM:

>>> def generator_b(gen):
...     yield from gen
...
>>> dis.dis(generator_b)
  2           0 LOAD_FAST                0 (gen)
              2 GET_YIELD_FROM_ITER
              4 LOAD_CONST               0 (None)
              6 YIELD_FROM
              8 POP_TOP
             10 LOAD_CONST               0 (None)
             12 RETURN_VALUE

我們現在用一個簡單的例子重新回顧一下程式的行為：

def generator_a():
	yield 1
	yield 2


def generator_b(gen):
	yield from gen


if __name__ == '__main__':
	gen = generator_b(generator_a())
	print(gen.send(None))
	print(gen.send(None))
	try:
		gen.send(None)
	except StopIteration:
		print("generator exit")

上面的程式輸出結果如下所示：

1
2
generator exit

從上面程式的輸出結果我們可以看到 generator_a 的兩個值都會被返回，這些魔法隱藏在位元組碼 YIELD_FROM 當中。YIELD_FROM 位元組碼會呼叫棧頂上的生成器物件的 send 方法，並且將引數生成器物件 gen 的返回結果返回，比如 1 和 2 這兩個值會被返回到 generator_b ，然後 generator_b 會將這個結果繼續傳播出來。

  • 在這個位元組碼執行最後會進行判斷虛擬機器當中是否出現了 StopIteration 異常，如果出現了則說 yield from 的生成器已經執行完了，則 generator_b 繼續往下執行。
  • 如果沒有 StopIteration 異常，則說明 yield from 的生成器沒有執行完成，這個時候虛擬機器會將當前棧幀的位元組碼執行位置往前移動，這麼做的目的是讓下一次生成器執行的時候繼續執行 YIELD_FROM 位元組碼，這就是 YIELD_FROM 能夠將另一個生成器物件執行完整的秘密。

總結

在本篇文章當中主要分析的生成器內部實現原理和相關的兩個重要的位元組碼，分析了生成器能夠停下來還能夠恢復執行的原因。本文最重要的兩點就是區分函式和生成器和 YIELD 、YIELD_FROM 兩個位元組碼，生成器是生成器函式返回的物件，YIELD 會直接進行函式返回，虛擬機器不會繼續往下執行，YIELD_FROM 除了會進行函式返回還會將位元組碼的執行位置往前移動，以保證 YIELD_FROM 下一次還能夠被執行。

本篇文章是深入理解 python 虛擬機器系列文章之一，文章地址：https://github.com/Chang-LeHung/dive-into-cpython

更多精彩內容合集可訪問專案：https://github.com/Chang-LeHung/CSCore

關注公眾號：一無是處的研究僧，瞭解更多計算機（Java、Python、計算機系統基礎、演算法與資料結構）知識。

相關文章