生成器進化到協程 Part 2

PrivateRookie發表於2018-07-07

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


上下文管理器

理解上下文可以聯想我們做閱讀理解時要解讀文章某處的意思需要閱讀該處前後段落,正是前後文提供了理解的“背景”。而程式的執行的上下文也可以理解為程式在執行時的某些變數,正是這些變數構成了執行環境,讓程式可以完成操作。

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 ,告訴直譯器正常退出;如果傳入異常,會使用 throwyield 處丟擲這個異常;如果有其他未捕捉的錯誤,就重新丟擲該錯誤。

實際的程式碼實現會更加複雜,還有一些異常情況沒有處理

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

如果你對怎麼實現感興趣,你可以閱讀程式碼或者再一次閱讀 PEP 343

總結

Part 2 都是關於上下文管理器的內容,與協程關係不大。但透過這部分我們可以看到 yield 完全不同的用法,也熟悉了控制流 (control-flow) ,這與 Part 3 的非同步處理流程有很大關係。讓我們 Part 3 再見。

本作品採用《CC 協議》,轉載必須註明作者和本文連結
多少事,從來急。天地轉,光陰迫。

相關文章