8.3 讓對象支持上下文管理協(xié)議

問題

你想讓你的對象支持上下文管理協(xié)議(with語句)。

解決方案

為了讓一個對象兼容 with 語句，你需要實現(xiàn) __enter()__ 和 __exit__() 方法。例如，考慮如下的一個類，它能為我們創(chuàng)建一個網(wǎng)絡連接：

from socket import socket, AF_INET, SOCK_STREAM

class LazyConnection:
    def __init__(self, address, family=AF_INET, type=SOCK_STREAM):
        self.address = address
        self.family = family
        self.type = type
        self.sock = None

    def __enter__(self):
        if self.sock is not None:
            raise RuntimeError('Already connected')
        self.sock = socket(self.family, self.type)
        self.sock.connect(self.address)
        return self.sock

    def __exit__(self, exc_ty, exc_val, tb):
        self.sock.close()
        self.sock = None

這個類的關鍵特點在于它表示了一個網(wǎng)絡連接，但是初始化的時候并不會做任何事情(比如它并沒有建立一個連接)。連接的建立和關閉是使用 with 語句自動完成的，例如：

from functools import partial

conn = LazyConnection(('www.python.org', 80))
# Connection closed
with conn as s:
    # conn.__enter__() executes: connection open
    s.send(b'GET /index.html HTTP/1.0\r\n')
    s.send(b'Host: www.python.org\r\n')
    s.send(b'\r\n')
    resp = b''.join(iter(partial(s.recv, 8192), b''))
    # conn.__exit__() executes: connection closed

討論

編寫上下文管理器的主要原理是你的代碼會放到 with 語句塊中執(zhí)行。當出現(xiàn) with 語句的時候，對象的 __enter__() 方法被觸發(fā)，它返回的值(如果有的話)會被賦值給 as 聲明的變量。然后，with 語句塊里面的代碼開始執(zhí)行。最后，__exit__() 方法被觸發(fā)進行清理工作。

不管 with 代碼塊中發(fā)生什么，上面的控制流都會執(zhí)行完，就算代碼塊中發(fā)生了異常也是一樣的。事實上，__exit__() 方法的第三個參數(shù)包含了異常類型、異常值和追溯信息(如果有的話)。__exit__() 方法能自己決定怎樣利用這個異常信息，或者忽略它并返回一個None值。如果 __exit__() 返回 True ，那么異常會被清空，就好像什么都沒發(fā)生一樣，with 語句后面的程序繼續(xù)在正常執(zhí)行。

還有一個細節(jié)問題就是 LazyConnection 類是否允許多個 with 語句來嵌套使用連接。很顯然，上面的定義中一次只能允許一個socket連接，如果正在使用一個socket的時候又重復使用 with 語句，就會產(chǎn)生一個異常了。不過你可以像下面這樣修改下上面的實現(xiàn)來解決這個問題：

from socket import socket, AF_INET, SOCK_STREAM

class LazyConnection:
    def __init__(self, address, family=AF_INET, type=SOCK_STREAM):
        self.address = address
        self.family = family
        self.type = type
        self.connections = []

    def __enter__(self):
        sock = socket(self.family, self.type)
        sock.connect(self.address)
        self.connections.append(sock)
        return sock

    def __exit__(self, exc_ty, exc_val, tb):
        self.connections.pop().close()

# Example use
from functools import partial

conn = LazyConnection(('www.python.org', 80))
with conn as s1:
    pass
    with conn as s2:
        pass
        # s1 and s2 are independent sockets

在第二個版本中，LazyConnection 類可以被看做是某個連接工廠。在內(nèi)部，一個列表被用來構造一個棧。每次 __enter__() 方法執(zhí)行的時候，它復制創(chuàng)建一個新的連接并將其加入到棧里面。__exit__() 方法簡單的從棧中彈出最后一個連接并關閉它。這里稍微有點難理解，不過它能允許嵌套使用 with 語句創(chuàng)建多個連接，就如上面演示的那樣。

在需要管理一些資源比如文件、網(wǎng)絡連接和鎖的編程環(huán)境中，使用上下文管理器是很普遍的。這些資源的一個主要特征是它們必須被手動的關閉或釋放來確保程序的正確運行。例如，如果你請求了一個鎖，那么你必須確保之后釋放了它，否則就可能產(chǎn)生死鎖。通過實現(xiàn) __enter__() 和 __exit__() 方法并使用 with 語句可以很容易的避免這些問題，因為 __exit__() 方法可以讓你無需擔心這些了。

在 contextmanager 模塊中有一個標準的上下文管理方案模板，可參考9.22小節(jié)。同時在12.6小節(jié)中還有一個對本節(jié)示例程序的線程安全的修改版。