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

問(wèn)題

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

解決方案

為了讓一個(gè)對(duì)象兼容 with 語(yǔ)句，你需要實(shí)現(xiàn) __enter()__ 和 __exit__() 方法。例如，考慮如下的一個(gè)類，它能為我們創(chuàng)建一個(gè)網(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

這個(gè)類的關(guān)鍵特點(diǎn)在于它表示了一個(gè)網(wǎng)絡(luò)連接，但是初始化的時(shí)候并不會(huì)做任何事情(比如它并沒(méi)有建立一個(gè)連接)。連接的建立和關(guān)閉是使用 with 語(yǔ)句自動(dòng)完成的，例如：

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

討論

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

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

還有一個(gè)細(xì)節(jié)問(wèn)題就是 LazyConnection 類是否允許多個(gè) with 語(yǔ)句來(lái)嵌套使用連接。很顯然，上面的定義中一次只能允許一個(gè)socket連接，如果正在使用一個(gè)socket的時(shí)候又重復(fù)使用 with 語(yǔ)句，就會(huì)產(chǎn)生一個(gè)異常了。不過(guò)你可以像下面這樣修改下上面的實(shí)現(xiàn)來(lái)解決這個(gè)問(wè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

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

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

在 contextmanager 模塊中有一個(gè)標(biāo)準(zhǔn)的上下文管理方案模板，可參考9.22小節(jié)。同時(shí)在12.6小節(jié)中還有一個(gè)對(duì)本節(jié)示例程序的線程安全的修改版。