Python3: 淺談 Python 3.3 的 Yield From 表達式

我最近需要使用 Python 2.7 解開 zip 壓縮檔並分析檔案內容。所以我寫了下面的程式。然而這個程式是有問題的。它沒有辦法列出 zip 壓縮檔裡面的檔案。為什麼呢?

#!/usr/bin/env python

import os
import shutil
import sys
import tempfile
import zipfile

def unzip(zip_file_path, out_dir):
    with zipfile.ZipFile(zip_file_path) as zip_file:
        zip_file.extractall(out_dir)

def list_files(path):
    for base, _, filenames in os.walk(path):
        for filename in filenames:
            yield os.path.join(base, filename)

def list_files_in_zip_file(zip_file_path):
    tmp_dir = tempfile.mkdtemp()
    try:
        unzip(zip_file_path, tmp_dir)
        return list_files(tmp_dir)
    finally:
        shutil.rmtree(tmp_dir)

def main():
    for path in list_files_in_zip_file(sys.argv[1]):
        print(path)

if __name__ == '__main__':
    main()

生成器函式

在 Python 程式語言中,def 述句的用途是定義函式。但是 def 定義的函式會隨著 def 區塊的內容而有些微的差異。

  1. 如果區塊內沒有 yield 表達式,稱為「普通函式(Normal Function)」
  2. 如果區塊內 yield 表達式,稱為「生成器函式(Generator Function)」

普通函式被「呼叫表達式(Calls Expression)」呼叫時,呼叫表達式會被解析(Evaluate)為 return 述句回傳的回傳值。例如:

>>> def example1():
...     return 'hello'
...

>>> example1()
'hello'

生成器函式被「呼叫表達式」呼叫時,生成器函式會回傳一個「生成器物件(Generator Object)」。此時,生成器函式內的代碼還不會被執行。例如:

>>> def example2():
...     for i in range(3):
...         print 'debug:', i
...         yield i
...

>>> example2()
<generator object example2 at 0x7fdacd720820>

接著,我們可以使用 next 函式向「生成器物件」請求一個數值。例如:

>>> g = example2()
>>> next(g)
debug: 0
0
>>> next(g)
debug: 1
1
>>> next(g)
debug: 2
2
>>> next(g)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

我們可以觀察到 next 函式會開始執行 example2 函式的代碼。當 yield 表達式被解折時,yield 表達式的參數會成為 next 函式的回傳值,而 example2 函式的執行狀態會被記錄在「生成器物件」裡面。當我們再次使用 next 函式向同一個「生成器物件」請求一個數值時,example2 函式就會從上次暫停的地方繼續執行。若 example2 函式已經執行完畢,則 next 會拋出一個 StopIteration 例外。

附帶一題,上面的 next() 呼叫也可以用 for 述句改寫:

>>> for i in example2():
...     print 'got:', i
debug: 0
got: 0
debug: 1
got: 1
debug: 2
got: 2

我們可以注意到在上面的例子中,執行 example2() 所需「執行狀態」是儲存於「生成器物件」裡面。與此對照,普通函式回傳後就只剩回傳值。普通函式執行過程中的「執行狀態」在回傳前就會被釋放與清除。

問題與解法

回到本文開頭的程式,哪裡寫錯了呢?

首先我們必需要知道哪些 def 述句會定義生成器函式、哪些 def 述句會定義普通函式。如前所述,只有 list_files 是生成器函式,其他的都是普通函式(包含 list_files_in_zip_file)。這意謂著 list_files_in_zip_file 的行為其實是:

  1. 建立一個暫時目錄
  2. 解壓縮
  3. 呼叫 list_files 生成器函式,建立一個可以列舉暫時目錄下所有檔案的生成器物件
  4. 執行 finally 子句,刪除暫時目錄
  5. 回傳第 3 步建立的生成器物件

最後,main() 得到「生成器物件」後才開始走訪已經被刪掉的暫時目錄。因此這個程式印不出任何資訊。

正確的作法是要把 list_files_in_zip_file() 也變成生成器函式以確保「暫時目錄的生命週期」長於「生成器物件的生命週期」:

def list_files_in_zip_file(zip_file_path):
    tmp_dir = tempfile.mkdtemp()
    try:
        unzip(zip_file_path, tmp_dir)
        for x in list_files(tmp_dir):  # Modified
            yield x                    # Modified
    finally:
        shutil.rmtree(tmp_dir)

類似的錯誤

其實一些看似無害的重構也會產生類似的錯誤。例如:下面的「生成器函式」本來會從一個檔案中抽取若干資訊,並在例外被拋出時印出檔名:

def extract_info(path):
    try:
        with open(path) as fp:
            for line in fp:             # To be refactored
                yield parse_line(line)  # To be refactored
    except ParseError:
        print >> sys.stderr, 'Failed to parse', path
        raise

接著,把中間的 for line in fp 重構為獨立的函式:

def extract_info_from_file(fp):
    for line in fp:
        yield parse_line(line)

此時,如果把 for line in fp 那二行代換成 return 述句,則 except 子句就永遠抓不到例外。因為下面重構過後的 extract_info() 的行為是直接把 extract_info_from_file() 回傳的生成器物件直接回傳給 extract_info() 的呼叫者,這個過程不會有 ParseError 例外。ParseError 例外只可能會在 extract_info() 的呼叫者迭代生成器物件時被拋出。

def extract_info(path):
    try:
        with open(path) as fp:
            return extract_info_from_file(fp)  # XXX BUG
    except ParseError:
        print >> sys.stderr, 'Failed to parse', path
        raise

正確的作法應該將 extract_info() 改寫為:

def extract_info(path):
    try:
        with open(path) as fp:
            for x in extract_info_from_file(fp):  # Rewritten
                yield x                           # Rewritten
    except ParseError:
        print >> sys.stderr, 'Failed to parse', path
        raise

Yield From 表達式

標題中的「yield from 表達式」是 PEP 380 在 Python 3.3 引入的新機制。「yield from 表達式」讓我們把生成器物件的工作「委任」給另一個生成器物件。簡單的說,它讓我們可以把「另一個生成器生成的值」當成「自己要生成的值」。例如,文章開頭的程式在 Python 3.3 可以改寫成:

def list_files_in_zip_file(zip_file_path):
    tmp_dir = tempfile.mkdtemp()
    try:
        unzip(zip_file_path, tmp_dir)
        yield from list_files(tmp_dir)  # Python 3.3
    finally:
        shutil.rmtree(tmp_dir)

Yield From 表達式與 Return 述句

除此之外,PEP 380 也稍微放寬了 yield 和 return 的規定。在此之前,同一個 def 述句不能同時包含 yield 述句和 return 述句。在 Python 3.3 之後,「yield from 表達式」會為被解析為「生成器函式中 return 述句回傳的回傳值」。舉例來說:

>>> def example3():
...     for i in range(3):
...         yield i
...     return 'end'
...

>>> def example4():
...     x = yield from example3()
...     print('example4: x:', x)
...

>>> for i in example4():
...     print('i:', i)
...

i: 0
i: 1
i: 2
example4: x: end

這個回傳值可以有不同的用途。例如,被委任的生成器可以回傳其生成的數值個數、是否有發生錯誤、是否要重試等等。