Context Managers 是我最喜歡的 python feature 之一，在恰當的時機使用 context manager 使程式碼更加簡潔、清晰，更加安全，複用性更好，更加 pythonic。本文簡單介紹一下其使用方法以及常見使用場景。
本文地址：https://www.cnblogs.com/xybaby/p/13202496.html
with statement and context manager
Python’s with statement supports the concept of a runtime context defined by a context manager
new statement "with" to the Python language to make it possible to factor out standard uses of try/finally statements.
在 pep0343 中，通過引入 context manager protocol 來支援 With statement ， context manager 是用來管理 context（上下文）的，即保證程式要保持一種特定的狀態 -- 無論是否發生異常。可以說，context manager 簡化了對 try-finally 的使用，而且更加安全，更加便於使用。
Transforming Code into Beautiful, Idiomatic Python 中，指出了 context manager 的最顯著的優點：
- Helps separate business logic from administrative logic
- Clean, beautiful tools for factoring code and improving code reuse
最廣為人知的例子，就是通過 with statement 來讀寫檔案，程式碼如下：
with open('test.txt') as f:
contect = f.read()
handle_content(content)
上面的程式碼幾乎等價於
f = open('test.txt')
try:
contect = f.read()
handle_content(content)
finally:
f.close()
注意，上面的finally的作用就是保證file.close一定會被呼叫，也就是資源一定會釋放。不過，很多時候，都會忘了去寫這個finally，而 with statement 就徹底避免了這個問題。
從上述兩段程式碼也可以看出，with statement 更加簡潔，而且將核心的業務邏輯（從檔案中讀取、處理資料）與其他邏輯（開啟、關係檔案）相分離，可讀性更強。
實現context manager protocol
一個類只要定義了__enter__、__exit__方法就實現了context manager 協議
object.__enter__(self)
Enter the runtime context related to this object. The with statement will bind this method’s return value to the target(s) specified in the as clause of the statement, if any.
object.__exit__(self, exc_type, exc_value, traceback)
Exit the runtime context related to this object. The parameters describe the exception that caused the context to be exited. If the context was exited without an exception, all three arguments will be None.
If an exception is supplied, and the method wishes to suppress the exception (i.e., prevent it from being propagated), it should return a true value. Otherwise, the exception will be processed normally upon exit from this method.
Note that __exit__() methods should not reraise the passed-in exception; this is the caller’s responsibility.
__enter__方法在進入這個 context 的時候呼叫，返回值賦值給 with as X 中的 X
__exit__方法在退出 context 的時候呼叫，如果沒有異常，後三個引數為 None。如果返回值為 True，則Suppress Exception，所以除非特殊情況都應返回 False。另外注意， __exit__方法本身不應該丟擲異常。
例子：BlockGuard
在看c++程式碼（如mongodb原始碼）的時候，經常看見其用 RAII 實現BlockGuard, 用以保證在離開 Block 的時候執行某些動作，同時，也提供手段來取消執行。
下面用python實現一下：
class BlockGuard(object):
def __init__(self, fn, *args, **kwargs):
self._fn = fn
self._args = args
self._kwargs = kwargs
self._canceled = False
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
if not self._canceled:
self._fn(*self._args, **self._kwargs)
self._fn = None
self._args = None
self._kwargs = None
return False
def cancel(self):
self._canceled = True
def foo():
print 'sth should be called'
def test_BlockGuard(cancel_guard):
print 'test_BlockGuard'
with BlockGuard(foo) as guard:
if cancel_guard:
guard.cancel()
print 'test_BlockGuard finish'
用yield實現context manager
標準庫 contextlib 中提供了一些方法，能夠簡化我們使用 context manager，如 contextlib.contextmanager(func) 使我們
無需再去實現一個包含__enter__ __exit__方法的類。
The function being decorated must return a generator-iterator when called. This iterator must yield exactly one value, which will be bound to the targets in the with statement’s as clause, if any.
例子如下：
from contextlib import contextmanager
@contextmanager
def managed_resource(*args, **kwds):
# Code to acquire resource, e.g.:
resource = acquire_resource(*args, **kwds)
try:
yield resource
finally:
# Code to release resource, e.g.:
release_resource(resource)
>>> with managed_resource(timeout=3600) as resource:
... # Resource is released at the end of this block,
... # even if code in the block raises an exception
需要注意的是：
- 一定要寫 try finally，才能保證
release_resource邏輯一定被呼叫 - 除非特殊情況，不再 catch exception，這就跟
__exit__一般不返回True一樣
例子: no_throw
這是業務開發中的一個需求, 比如觀察者模式，不希望因為其中一個觀察者出了 trace 就影響後續的觀察者，就可以這樣做：
from contextlib import contextmanager
@contextmanager
def no_throw(*exceptions):
try:
yield
except exceptions:
pass
def notify_observers(seq):
for fn in [sum, len, max, min]:
with no_throw(Exception):
print "%s result %s" % (fn.__name__, fn(seq))
if __name__ == '__main__':
notify_observers([])
在python 3.x 的 contexlib 中，就提供了一個contextlib.suppress(*exceptions), 實現了同樣的效果。
context manager 應用場景
context manager 誕生的初衷就在於簡化 try-finally，因此就適合應用於在需要 finally 的地方，也就是需要清理的地方，比如
- 保證資源的安全釋放，如 file、lock、semaphore、network connection 等
- 臨時操作的復原，如果一段邏輯有 setup、prepare，那麼就會對應 cleanup、teardown。
對於第一種情況，網路連線釋放的例子，後面會結合 pymongo 的程式碼展示。
在這裡先來看看第二種用途：保證程式碼在一個臨時的、特殊的上下文（context）中執行，且在執行結束之後恢復到之前的上下文環境。
改變工作目錄
from contextlib import contextmanager
import os
@contextmanager
def working_directory(path):
current_dir = os.getcwd()
os.chdir(path)
try:
yield
finally:
os.chdir(current_dir)
with working_directory("data/stuff"):
pass
臨時檔案、資料夾
很多時候會產生一堆臨時檔案，比如build的中間狀態，這些臨時檔案都需要在結束之後清除。
from tempfile import mkdtemp
from shutil import rmtree
@contextmanager
def temporary_dir(*args, **kwds):
name = mkdtemp(*args, **kwds)
try:
yield name
finally:
shutil.rmtree(name)
with temporary_dir() as dirname:
pass
重定向標準輸出、標準錯誤
@contextmanager
def redirect_stdout(fileobj):
oldstdout = sys.stdout
sys.stdout = fileobj
try:
yield fieldobj
finally:
sys.stdout = oldstdout
在 python3.x 中，已經提供了 contextlib.redirect_stdout contextlib.redirect_stderr 實現上述功能
調整logging level
這個在查問題的適合非常有用，一般生產環境不會輸出 debug level 的日誌，但如果出了問題，可以臨時對某些制定的函式呼叫輸出debug 日誌
from contextlib import contextmanager
import logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)
ch = logging.StreamHandler()
ch.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
logger.addHandler(ch)
@contextmanager
def change_log_level(level):
old_level = logger.getEffectiveLevel()
try:
logger.setLevel(level)
yield
finally:
logger.setLevel(old_level)
def test_logging():
logger.debug("this is a debug message")
logger.info("this is a info message")
logger.warn("this is a warning message")
with change_log_level(logging.DEBUG):
test_logging()
pymongo中的context manager使用
在 pymongo 中，封裝了好幾個 context manager，用以
- 管理 semaphore
- 管理 connection
- 資源清理
而且，在 pymongo 中，給出了巢狀使用 context manager 的好例子，用來保證 socket 在使用完之後一定返回連線池(pool)。
# server.py
@contextlib.contextmanager
def get_socket(self, all_credentials, checkout=False):
with self.pool.get_socket(all_credentials, checkout) as sock_info:
yield sock_info
# pool.py
@contextlib.contextmanager
def get_socket(self, all_credentials, checkout=False):
sock_info = self._get_socket_no_auth()
try:
sock_info.check_auth(all_credentials)
yield sock_info
except:
# Exception in caller. Decrement semaphore.
self.return_socket(sock_info)
raise
else:
if not checkout:
self.return_socket(sock_info)
可以看到，server.get_socket 呼叫了 pool.get_socket, 使用 server.get_socket 的程式碼完全不瞭解、也完全不用關心 socket 的釋放細節，如果把 try-except-finally-else 的邏輯移到所有使用socket的地方，程式碼就會很醜、很臃腫。
比如，在mongo_client 中需要使用到 socket：
with server.get_socket(all_credentials) as sock_info:
sock_info.authenticate(credentials)