生成器進化到協程 Part 2

• Part 1 傳送門
• David Beazley 的博客
• PPT 下載地址

在 Part 1 咱們已經介紹了生成器的定義和生成器的操做，如今讓咱們開始使用生成器。Part 2 主要描述瞭如何使用 yield 和 contextmanager 建立一個上下文管理器，並解釋了原理。

---

上下文管理器

理解上下文能夠聯想咱們作閱讀理解時要解讀文章某處的意思須要閱讀該處先後段落，正是先後文提供了理解的「背景」。而程序的運行的上下文也能夠理解爲程序在運行時的某些變量，正是這些變量構成了運行環境，讓程序能夠完成操做。

Python 中的上下文管理器提供這樣一種功能，爲你的程序運行時提供一個特定「空間」，當進入這個空間時 Python 上下文管理器 爲你作一些準備工做。這個「空間」中通常含有特殊的變量，你在這個「空間」中進行一些操做，而後離開。在你離開時 Python 上下文管理器又會幫你作一些收尾工做，保證不會污染運行環境。

下面是一些常見的代碼模式

# 讀取文件
f = open()
# do something
f.close()


# 使用鎖
lock.acquire()
# do somethin
lock.release()


# 進行數據庫操做
db.start_transaction()
# do something
db.commit()


# 對某段代碼進行計時
start = time.time()
# do something
end = time.time()

這些代碼進行的都是「先作這個（準備工做，好比獲取一個數據庫鏈接），而後作這個（好比寫入數據），最後整理工做環境（如提交改動，關閉連接，釋放資源等）。

若是使用 with 能夠這樣寫：

witn open(filename) as f:
    # do something
    pass
    
with lock():
    # do something
    pass

with 語句實際上使用了實現了 __enter__ 和 __exit__ 方法的上下文管理器類。一個典型的上下文管理器類以下：

clss ContextManager:
    def __enter__(self):
        return value
    def __exit__(self, exc_type, val, tb):
        if exec_type is None:
            return
        else:
            # 處理異常
            return True if handled else False

正如方法名明確告訴咱們的，__enter__ 方法負責進入上下的準備工做，若是有須要能夠返回一個值，這個值將會被賦值給 with ContextManager() as ret_value 中的 ret_value 。__exit__ 則負責收尾工做，這包括了異常處理。

對於這樣一段代碼

with ContextManager() as var:
    # do something

至關於

ctxmanager = ContextManager()
var = ctxmanager.__enter__()
# do somethin
ctxmanager.__exit__()

一個可用的例子：

import tempfile
import shutil

class TmpDir:
    def __enter__(self):
        self.dirname = tempfile.mkdtemp()
        return self.dirname
    
    def __exit__(self, exc, val, tb):
    shutil.rmtree(self.dirname)

這個上下文管理提供臨時文件的功能，在 with 語句結束後會自動刪除臨時文件夾。

with TempDir() as dirname:
    # 使用臨時文件夾進行一些操做
    pass

關於上面兩個特殊方法的文檔能夠在 Python 文檔的 Context Manager Types 找到。另外關於 with 關鍵字的詳細說明參考 PEP 343，不過這篇 PEP　不是很好讀，Good Luck :simple_smile:！

使用 yield 和 contextmanager

能看到這裏的都應該對上下文管理器有所瞭解，準備好把 yield 加入咱們的上下文管理器代碼中。

先看一個例子

import tempfile, shutil
from contextlib import contextmanager

@contextmanager
def tempdir():
    outdir = tempfile.mkdtemp()
    try:
        yield outdir
    finally:
        shutil.rmtree(outdir)

與使用上下文管理器類的實現方式不一樣，這裏咱們沒有顯式實現 __enter__ 和 __exit__，而是經過 contextmanager 裝飾器和 yield 實現，你能夠試試這兩種方式是等價的。

要理解上面的代碼，能夠把 yield 想象爲一把剪刀，把這個函數一分爲二，上部分至關於 __enter__，下部分至關於 __exit__。我這樣說你們應該明白了吧。

import tempfile, shutil
from contextlib import contextmanager

@contextmanager
def tempdir():
    outdir = tempfile.mkdtemp() #
    try:                        # __enter__
        yield outdir            #
--cut---╳-----------------------------------
    finally:                    #
        shutil.rmtree(outdir)   # __exit__

實現「剪刀」功能關鍵在於 contextmanager 。對於上面的代碼，咱們來一步一步地結構它：

contextmanager 裝飾器

contextmanager 其實使用了一個上下文管理器類，這個類在在初始化時須要提供一個生成器。

class GeneratorCM:
    def __init__(self, gen):
        self.gen = gen
    
    def __enter__(self):
       ...
        
    def __exit__(self, exc, val, tb):
        ...

contextmanager 的實現以下

def contextmanager(func):
    def run(*args, **kwargs):
        return GeneratorCM(func(*args, **kwargs))
    return run

因爲 contextmanger 所裝飾的函數裏有 yield 因此咱們在調用 func(*args, **kwargs) 時返回的是一個生成器。要使這個生成器前進，咱們須要調用 next 函數

讓生成器前進

def __enter__(self):
    return next(self.gen)

GeneratorCM 的 __ente__ 方法會讓生成器前進到 yield 語句處，並返回產出值。

收尾

def __exit__(self, exc, val, tb):
    try:
        if exc is None:
            next(self.gen)
        else:
            self.gen.throw(exc, val, tb)
        raise RuntimeError('Generator didn\'t stop')
    except StopIteration:
        return True
    except:
        if sys.exc_info()[1] is not val: raise

__exit__ 函數的邏輯比較複雜，若是沒有傳入異常，首先它會嘗試對生成器調用 next，正常狀況下這會拋出 StopIteration ，這個異常會被不作並返回 True ，告訴解釋器正常退出；若是傳入異常，會使用 throw 在 yield 處拋出這個異常；若是有其餘未捕捉的錯誤，就從新拋出該錯誤。

實際的代碼實現會更加複雜，還有一些異常狀況沒有處理

• 沒有相關值的異常
• 在 with 語句塊中拋出的 StopIteration
• 在上下文管理器中拋出的異常

若是你對怎麼實現感興趣，你能夠閱讀代碼或者再一次閱讀 PEP 343。

總結

Part 2 都是關於上下文管理器的內容，與協程關係不大。但經過這部分咱們能夠看到 yield 徹底不一樣的用法，也熟悉了控制流 (control-flow) ，這與 Part 3 的異步處理流程有很大關係。讓咱們 Part 3 再見。