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 區塊的內容而有些微的差異。

如果區塊內沒有 yield 表達式，稱為「普通函式（Normal Function）」
如果區塊內有 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 的行為其實是：

建立一個暫時目錄
解壓縮
呼叫 list_files 生成器函式，建立一個可以列舉暫時目錄下所有檔案的生成器物件
執行 finally 子句，刪除暫時目錄
回傳第 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

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

參考資料

PEP 380 -- Syntax for Delegating to a Subgenerator