Python中生成器和yield語句的用法詳解

在開始課程之前，我要求學(xué)生們填寫一份調(diào)查表，這個調(diào)查表反映了它們對Python中一些概念的理解情況。一些話題("if/else控制流" 或者 "定義和使用函數(shù)")對于大多數(shù)學(xué)生是沒有問題的。但是有一些話題，大多數(shù)學(xué)生只有很少，或者完全沒有任何接觸，尤其是“生成器和yield關(guān)鍵字”。我猜這對大多數(shù)新手Python程序員也是如此。

有事實(shí)表明，在我花了大功夫后，有些人仍然不能理解生成器和yield關(guān)鍵字。我想讓這個問題有所改善。在這篇文章中，我將解釋yield關(guān)鍵字到底是什么，為什么它是有用的，以及如何來使用它。

注意：最近幾年，生成器的功能變得越來越強(qiáng)大，它已經(jīng)被加入到了PEP。在我的下一篇文章中，我會通過協(xié)程(coroutine)，協(xié)同式多任務(wù)處理(cooperative multitasking)，以及異步IO(asynchronous I/O)(尤其是GvR正在研究的 "tulip" 原型的實(shí)現(xiàn))來介紹yield的真正威力。但是在此之前，我們要對生成器和yield有一個扎實(shí)的理解.

協(xié)程與子例程

我們調(diào)用一個普通的Python函數(shù)時，一般是從函數(shù)的第一行代碼開始執(zhí)行，結(jié)束于return語句、異常或者函數(shù)結(jié)束（可以看作隱式的返回None）。一旦函數(shù)將控制權(quán)交還給調(diào)用者，就意味著全部結(jié)束。函數(shù)中做的所有工作以及保存在局部變量中的數(shù)據(jù)都將丟失。再次調(diào)用這個函數(shù)時，一切都將從頭創(chuàng)建。

對于在計(jì)算機(jī)編程中所討論的函數(shù)，這是很標(biāo)準(zhǔn)的流程。這樣的函數(shù)只能返回一個值，不過，有時可以創(chuàng)建能產(chǎn)生一個序列的函數(shù)還是有幫助的。要做到這一點(diǎn)，這種函數(shù)需要能夠“保存自己的工作”。

我說過，能夠“產(chǎn)生一個序列”是因?yàn)槲覀兊暮瘮?shù)并沒有像通常意義那樣返回。return隱含的意思是函數(shù)正將執(zhí)行代碼的控制權(quán)返回給函數(shù)被調(diào)用的地方。而"yield"的隱含意思是控制權(quán)的轉(zhuǎn)移是臨時和自愿的，我們的函數(shù)將來還會收回控制權(quán)。

在Python中，擁有這種能力的“函數(shù)”被稱為生成器，它非常的有用。生成器（以及yield語句）最初的引入是為了讓程序員可以更簡單的編寫用來產(chǎn)生值的序列的代碼。 以前，要實(shí)現(xiàn)類似隨機(jī)數(shù)生成器的東西，需要實(shí)現(xiàn)一個類或者一個模塊，在生成數(shù)據(jù)的同時保持對每次調(diào)用之間狀態(tài)的跟蹤。引入生成器之后，這變得非常簡單。

為了更好的理解生成器所解決的問題，讓我們來看一個例子。在了解這個例子的過程中，請始終記住我們需要解決的問題：生成值的序列。

注意：在Python之外，最簡單的生成器應(yīng)該是被稱為協(xié)程（coroutines）的東西。在本文中，我將使用這個術(shù)語。請記住，在Python的概念中，這里提到的協(xié)程就是生成器。Python正式的術(shù)語是生成器；協(xié)程只是便于討論，在語言層面并沒有正式定義。

例子：有趣的素?cái)?shù)

假設(shè)你的老板讓你寫一個函數(shù)，輸入?yún)?shù)是一個int的list，返回一個可以迭代的包含素?cái)?shù)1 的結(jié)果。

記住，迭代器（Iterable） 只是對象每次返回特定成員的一種能力。

你肯定認(rèn)為"這很簡單"，然后很快寫出下面的代碼：
def get_primes(input_list):
  result_list = list()
  for element in input_list:
    if is_prime(element):
      result_list.append()
 
  return result_list
 
# 或者更好一些的...
 
def get_primes(input_list):
  return (element for element in input_list if is_prime(element))
 
# 下面是 is_prime 的一種實(shí)現(xiàn)...
 
def is_prime(number):
  if number > 1:
    if number == 2:
      return True
    if number % 2 == 0:
      return False
    for current in range(3, int(math.sqrt(number) + 1), 2):
      if number % current == 0:
        return False
    return True
  return False

上面 is_prime 的實(shí)現(xiàn)完全滿足了需求，所以我們告訴老板已經(jīng)搞定了。她反饋說我們的函數(shù)工作正常，正是她想要的。

處理無限序列

噢，真是如此嗎？過了幾天，老板過來告訴我們她遇到了一些小問題：她打算把我們的get_primes函數(shù)用于一個很大的包含數(shù)字的list。實(shí)際上，這個list非常大，僅僅是創(chuàng)建這個list就會用完系統(tǒng)的所有內(nèi)存。為此，她希望能夠在調(diào)用get_primes函數(shù)時帶上一個start參數(shù)，返回所有大于這個參數(shù)的素?cái)?shù)（也許她要解決 Project Euler problem 10）。

我們來看看這個新需求，很明顯只是簡單的修改get_primes是不可能的。 自然，我們不可能返回包含從start到無窮的所有的素?cái)?shù)的列表 (雖然有很多有用的應(yīng)用程序可以用來操作無限序列)?？瓷先ビ闷胀ê瘮?shù)處理這個問題的可能性比較渺茫。

在我們放棄之前，讓我們確定一下最核心的障礙，是什么阻止我們編寫滿足老板新需求的函數(shù)。通過思考，我們得到這樣的結(jié)論：函數(shù)只有一次返回結(jié)果的機(jī)會，因而必須一次返回所有的結(jié)果。得出這樣的結(jié)論似乎毫無意義；“函數(shù)不就是這樣工作的么”，通常我們都這么認(rèn)為的?？墒牵粚W(xué)不成，不問不知，“如果它們并非如此呢？”

想象一下，如果get_primes可以只是簡單返回下一個值，而不是一次返回全部的值，我們能做什么？我們就不再需要創(chuàng)建列表。沒有列表，就沒有內(nèi)存的問題。由于老板告訴我們的是，她只需要遍歷結(jié)果，她不會知道我們實(shí)現(xiàn)上的區(qū)別。
 

不幸的是，這樣做看上去似乎不太可能。即使是我們有神奇的函數(shù)，可以讓我們從n遍歷到無限大，我們也會在返回第一個值之后卡住：
 
def get_primes(start):
  for element in magical_infinite_range(start):
    if is_prime(element):
      return element
假設(shè)這樣去調(diào)用get_primes：
 
def solve_number_10():
  # She *is* working on Project Euler #10, I knew it!
  total = 2
  for next_prime in get_primes(3):
    if next_prime < 2000000:
      total += next_prime
    else:
      print(total)
      return

顯然，在get_primes中，一上來就會碰到輸入等于3的，并且在函數(shù)的第4行返回。與直接返回不同，我們需要的是在退出時可以為下一次請求準(zhǔn)備一個值。

不過函數(shù)做不到這一點(diǎn)。當(dāng)函數(shù)返回時，意味著全部完成。我們保證函數(shù)可以再次被調(diào)用，但是我們沒法保證說，“呃，這次從上次退出時的第4行開始執(zhí)行，而不是常規(guī)的從第一行開始”。函數(shù)只有一個單一的入口：函數(shù)的第1行代碼。
 

走進(jìn)生成器

這類問題極其常見以至于Python專門加入了一個結(jié)構(gòu)來解決它：生成器。一個生成器會“生成”值。創(chuàng)建一個生成器幾乎和生成器函數(shù)的原理一樣簡單。

一個生成器函數(shù)的定義很像一個普通的函數(shù)，除了當(dāng)它要生成一個值的時候，使用yield關(guān)鍵字而不是return。如果一個def的主體包含yield，這個函數(shù)會自動變成一個生成器（即使它包含一個return）。除了以上內(nèi)容，創(chuàng)建一個生成器沒有什么多余步驟了。

生成器函數(shù)返回生成器的迭代器。這可能是你最后一次見到“生成器的迭代器”這個術(shù)語了， 因?yàn)樗鼈兺ǔ＞捅环Q作“生成器”。要注意的是生成器就是一類特殊的迭代器。作為一個迭代器，生成器必須要定義一些方法(method)，其中一個就是__next__()。如同迭代器一樣，我們可以使用next()函數(shù)來獲取下一個值。

為了從生成器獲取下一個值，我們使用next()函數(shù)，就像對付迭代器一樣。

(next()會操心如何調(diào)用生成器的__next__()方法)。既然生成器是一個迭代器，它可以被用在for循環(huán)中。

每當(dāng)生成器被調(diào)用的時候，它會返回一個值給調(diào)用者。在生成器內(nèi)部使用yield來完成這個動作(例如yield 7)。為了記住yield到底干了什么，最簡單的方法是把它當(dāng)作專門給生成器函數(shù)用的特殊的return(加上點(diǎn)小魔法)。**

yield就是專門給生成器用的return(加上點(diǎn)小魔法)。
下面是一個簡單的生成器函數(shù)：    
>>> def simple_generator_function():
>>>  yield 1
>>>  yield 2
>>>  yield 3
這里有兩個簡單的方法來使用它：
 
>>> for value in simple_generator_function():
>>>   print(value)
1
2
3
>>> our_generator = simple_generator_function()
>>> next(our_generator)
1
>>> next(our_generator)
2
>>> next(our_generator)
3

魔法?

那么神奇的部分在哪里?我很高興你問了這個問題!當(dāng)一個生成器函數(shù)調(diào)用yield，生成器函數(shù)的“狀態(tài)”會被凍結(jié)，所有的變量的值會被保留下來，下一行要執(zhí)行的代碼的位置也會被記錄，直到再次調(diào)用next()。一旦next()再次被調(diào)用，生成器函數(shù)會從它上次離開的地方開始。如果永遠(yuǎn)不調(diào)用next()，yield保存的狀態(tài)就被無視了。

我們來重寫get_primes()函數(shù)，這次我們把它寫作一個生成器。注意我們不再需要magical_infinite_range函數(shù)了。使用一個簡單的while循環(huán)，我們創(chuàng)造了自己的無窮串列。
 
def get_primes(number):
    while True:
        if is_prime(number):
            yield number
        number += 1

如果生成器函數(shù)調(diào)用了return，或者執(zhí)行到函數(shù)的末尾，會出現(xiàn)一個StopIteration異常。 這會通知next()的調(diào)用者這個生成器沒有下一個值了(這就是普通迭代器的行為)。這也是這個while循環(huán)在我們的get_primes()函數(shù)出現(xiàn)的原因。如果沒有這個while，當(dāng)我們第二次調(diào)用next()的時候，生成器函數(shù)會執(zhí)行到函數(shù)末尾，觸發(fā)StopIteration異常。一旦生成器的值用完了，再調(diào)用next()就會出現(xiàn)錯誤，所以你只能將每個生成器的使用一次。下面的代碼是錯誤的：    
>>> our_generator = simple_generator_function()
>>> for value in our_generator:
>>>   print(value)
 
>>> # 我們的生成器沒有下一個值了...
>>> print(next(our_generator))
Traceback (most recent call last):
 File "<ipython-input-13-7e48a609051a>", line 1, in <module>
  next(our_generator)
StopIteration
 
>>> # 然而，我們總可以再創(chuàng)建一個生成器
>>> # 只需再次調(diào)用生成器函數(shù)即可
 
>>> new_generator = simple_generator_function()
>>> print(next(new_generator)) # 工作正常
1

因此，這個while循環(huán)是用來確保生成器函數(shù)永遠(yuǎn)也不會執(zhí)行到函數(shù)末尾的。只要調(diào)用next()這個生成器就會生成一個值。這是一個處理無窮序列的常見方法（這類生成器也是很常見的）。

執(zhí)行流程

讓我們回到調(diào)用get_primes的地方：solve_number_10。    
def solve_number_10():
  # She *is* working on Project Euler #10, I knew it!
  total = 2
  for next_prime in get_primes(3):
    if next_prime < 2000000:
      total += next_prime
    else:
      print(total)
      return

我們來看一下solve_number_10的for循環(huán)中對get_primes的調(diào)用，觀察一下前幾個元素是如何創(chuàng)建的有助于我們的理解。當(dāng)for循環(huán)從get_primes請求第一個值時，我們進(jìn)入get_primes，這時與進(jìn)入普通函數(shù)沒有區(qū)別。

        進(jìn)入第三行的while循環(huán)
        停在if條件判斷（3是素?cái)?shù)）
        通過yield將3和執(zhí)行控制權(quán)返回給solve_number_10

接下來，回到insolve_number_10：

        for循環(huán)得到返回值3
        for循環(huán)將其賦給next_prime
        total加上next_prime
        for循環(huán)從get_primes請求下一個值

這次，進(jìn)入get_primes時并沒有從開頭執(zhí)行，我們從第5行繼續(xù)執(zhí)行，也就是上次離開的地方。    
def get_primes(number):
  while True:
    if is_prime(number):
      yield number
    number += 1 # <<<<<<<<<<

最關(guān)鍵的是，number還保持我們上次調(diào)用yield時的值（例如3）。記住，yield會將值傳給next()的調(diào)用方，同時還會保存生成器函數(shù)的“狀態(tài)”。接下來，number加到4，回到while循環(huán)的開始處，然后繼續(xù)增加直到得到下一個素?cái)?shù)（5）。我們再一次把number的值通過yield返回給solve_number_10的for循環(huán)。這個周期會一直執(zhí)行，直到for循環(huán)結(jié)束（得到的素?cái)?shù)大于2,000,000）。

更給力點(diǎn)

在PEP 342中加入了將值傳給生成器的支持。PEP 342加入了新的特性，能讓生成器在單一語句中實(shí)現(xiàn)，生成一個值（像從前一樣），接受一個值，或同時生成一個值并接受一個值。

我們用前面那個關(guān)于素?cái)?shù)的函數(shù)來展示如何將一個值傳給生成器。這一次，我們不再簡單地生成比某個數(shù)大的素?cái)?shù)，而是找出比某個數(shù)的等比級數(shù)大的最小素?cái)?shù)（例如10， 我們要生成比10，100，1000，10000 ... 大的最小素?cái)?shù)）。我們從get_primes開始：
    
def print_successive_primes(iterations, base=10):
  # 像普通函數(shù)一樣，生成器函數(shù)可以接受一個參數(shù)
   
  prime_generator = get_primes(base)
  # 這里以后要加上點(diǎn)什么
  for power in range(iterations):
    # 這里以后要加上點(diǎn)什么
 
def get_primes(number):
  while True:
    if is_prime(number):
    # 這里怎么寫?
 get_primes的后幾行需要著重解釋。yield關(guān)鍵字返回number的值，而像 other = yield foo 這樣的語句的意思是，"返回foo的值，這個值返回給調(diào)用者的同時，將other的值也設(shè)置為那個值"。你可以通過send方法來將一個值”發(fā)送“給生成器。
 
def get_primes(number):
  while True:
    if is_prime(number):
      number = yield number
    number += 1

通過這種方式，我們可以在每次執(zhí)行yield的時候?yàn)閚umber設(shè)置不同的值。現(xiàn)在我們可以補(bǔ)齊print_successive_primes中缺少的那部分代碼：    
def print_successive_primes(iterations, base=10):
  prime_generator = get_primes(base)
  prime_generator.send(None)
  for power in range(iterations):
    print(prime_generator.send(base ** power))

這里有兩點(diǎn)需要注意：首先，我們打印的是generator.send的結(jié)果，這是沒問題的，因?yàn)閟end在發(fā)送數(shù)據(jù)給生成器的同時還返回生成器通過yield生成的值（就如同生成器中yield語句做的那樣）。

第二點(diǎn)，看一下prime_generator.send(None)這一行，當(dāng)你用send來“啟動”一個生成器時（就是從生成器函數(shù)的第一行代碼執(zhí)行到第一個yield語句的位置），你必須發(fā)送None。這不難理解，根據(jù)剛才的描述，生成器還沒有走到第一個yield語句，如果我們發(fā)生一個真實(shí)的值，這時是沒有人去“接收”它的。一旦生成器啟動了，我們就可以像上面那樣發(fā)送數(shù)據(jù)了。

綜述

在本系列文章的后半部分，我們將討論一些yield的高級用法及其效果。yield已經(jīng)成為Python最強(qiáng)大的關(guān)鍵字之一?，F(xiàn)在我們已經(jīng)對yield是如何工作的有了充分的理解，我們已經(jīng)有了必要的知識，可以去了解yield的一些更“費(fèi)解”的應(yīng)用場景。

不管你信不信，我們其實(shí)只是揭開了yield強(qiáng)大能力的一角。例如，send確實(shí)如前面說的那樣工作，但是在像我們的例子這樣，只是生成簡單的序列的場景下，send幾乎從來不會被用到。下面我貼一段代碼，展示send通常的使用方式。對于這段代碼如何工作以及為何可以這樣工作，在此我并不打算多說，它將作為第二部分很不錯的熱身。    
import random
 
def get_data():
  """返回0到9之間的3個隨機(jī)數(shù)"""
  return random.sample(range(10), 3)
 
def consume():
  """顯示每次傳入的整數(shù)列表的動態(tài)平均值"""
  running_sum = 0
  data_items_seen = 0
 
  while True:
    data = yield
    data_items_seen += len(data)
    running_sum += sum(data)
    print('The running average is {}'.format(running_sum / float(data_items_seen)))
 
def produce(consumer):
  """產(chǎn)生序列集合，傳遞給消費(fèi)函數(shù)（consumer）"""
  while True:
    data = get_data()
    print('Produced {}'.format(data))
    consumer.send(data)
    yield
 
if __name__ == '__main__':
  consumer = consume()
  consumer.send(None)
  producer = produce(consumer)
 
  for _ in range(10):
    print('Producing...')
    next(producer)
請謹(jǐn)記……
我希望您可以從本文的討論中獲得一些關(guān)鍵的思想：
        generator是用來產(chǎn)生一系列值的
        yield則像是generator函數(shù)的返回結(jié)果
        yield唯一所做的另一件事就是保存一個generator函數(shù)的狀態(tài)
        generator就是一個特殊類型的迭代器（iterator）
        和迭代器相似，我們可以通過使用next()來從generator中獲取下一個值
        通過隱式地調(diào)用next()來忽略一些值

我希望這篇文章是有益的。如果您還從來沒有聽說過generator，我希望現(xiàn)在您可以理解它是什么以及它為什么是有用的，并且理解如何使用它。如果您已經(jīng)在某種程度上比較熟悉generator，我希望這篇文章現(xiàn)在可以讓您掃清對generator的一些困惑。