問題

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

解決方案

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

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

這個類的關(guān)鍵特點在于它表示了一個網(wǎng)絡(luò)連接，但是初始化的時候并不會做任何事情(比如它并沒有建立一個連接)。連接的建立和關(guān)閉是使用 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的時候又重復(fù)使用 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)部，一個列表被用來構(gòu)造一個棧。每次 __enter__() 方法執(zhí)行的時候，它復(fù)制創(chuàng)建一個新的連接并將其加入到棧里面。__exit__() 方法簡單的從棧中彈出最后一個連接并關(guān)閉它。這里稍微有點難理解，不過它能允許嵌套使用 with 語句創(chuàng)建多個連接，就如上面演示的那樣。

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

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