談一談Python的上下文管理器 – 思誠之道

經常在Python程式碼中看到with語句，仔細分析下，會發現這個with語句功能好強，可以自動關閉資源。這個在Python中叫上下文管理器Context Manager。那我們要怎麼用它，什麼時候用它呢。這裡我們就來聊一聊。

上下文管理器的作用

很多情況，當我們使用完一個資源後，我們需要手動的關閉掉它，比如操作檔案，建立資料庫連線等。但是，在使用資源的過程中，如果遇到異常，很可能錯誤被直接丟擲，導致來不及關閉資源。所以在大部分程式語言裡，我們使用”try-finally”語句來確保資源會關閉。比如下面的Python寫檔案程式碼：

這樣做固然沒有問題，但是當”try-finally”中間的邏輯複雜，而且還帶有各種巢狀的話，程式碼就很不容易維護。Python的with語句，可以說功能同上面的”try-finally”幾乎一樣，但程式碼看上去簡潔的多，我們來實現同樣的功能：

with語句後面跟著open()方法，如果它有返回值的話，可以使用as語句將其賦值給f。在with語句塊退出時，”f.close()”方法會自動被呼叫，即使”f.write()”出現異常，也能確保close()方法被呼叫。

自定義類來使用上下文管理器

上例中”open()”方法是Python自帶的，那我們怎麼定義自己的型別來使用with語句呢。其實只要你的類定義了”__enter__()”和”__exit__()”方法，就可以使用Python的上下文管理器了。”__enter__()”方法會在with語句進入時被呼叫，其返回值會賦給as關鍵字後的變數；而”__exit__()”方法會在with語句塊退出後自動被呼叫。

我們來實現個跟上節一樣的檔案寫入功能：

異常處理

肯定有朋友注意到上面的”__exit__()”帶了三個引數，是的，他們是用來異常處理的。大部分情況下，我們希望with語句中遇到的異常最後被丟擲，但也有時候，我們想處理這些異常。”__exit__()”方法中的三個引數exc_type, exc_val, exc_tb分別代表異常型別，異常值，和異常的Traceback。當你處理完異常後，你可以讓”__exit__()”方法返回True，此時該異常就會不會再被丟擲。比如我們將上例中的”__exit__()”方法改一下：

現在，如果遇到SyntaxError的話，異常會被正常丟擲，而其他異常的話都會被忽略。

contextlib模組

Python中還有一個contextlib模組提供一些簡便的上下文管理器功能。

closing()方法

如果說with語句塊在退出時會自動呼叫”__exit__()”方法的話，那用了”contextlib.closing()”的with語句塊則在退出時會自動呼叫”close()”方法。看一下示例：

程式執行後，會列印出

說明Resource類建立的物件被賦給了as關鍵字後面的變數r，而with語句塊退出時，自動呼叫了”r.close()”方法。

contextmanager裝飾器

“@contextlib.contextmanager”是一個裝飾器，由它修飾的方法會有兩部分構成，中間由yield關鍵字分開。由此方法建立的上下文管理器，在程式碼塊執行前會先執行yield上面的語句；在程式碼塊執行後會再執行yield下面的語句。看個例子比較容易明白：

這個”timeit()”方法實現了一個計時器，它會計算由他生成的with語句塊執行時間。可以看出，yield上面的語句就如同之間介紹過的”__enter__()”方法，而yield下面的語句就如同”__exit__()”方法。而yield部分就是with語句塊中的程式碼。如果yield後面帶引數的話，我們就可以用as關鍵字賦值給後面的變數，比如上例：

需要注意的是，”@contextlib.contextmanager”不像之前介紹的”__exit__()”方法，遇到異常也會執行。也就是with語句塊丟擲異常的話，yield後面的程式碼將不會被執行。所以，必要時你需要對yield語句使用”try-finally”。