9.22 定義上下文管理器的簡(jiǎn)單方法

問(wèn)題

你想自己去實(shí)現(xiàn)一個(gè)新的上下文管理器，以便使用with語(yǔ)句。

解決方案

實(shí)現(xiàn)一個(gè)新的上下文管理器的最簡(jiǎn)單的方法就是使用?<span class="pre" style="box-sizing: border-box;">contexlib</span>?模塊中的?<span class="pre" style="box-sizing: border-box;">@contextmanager</span>?裝飾器。 下面是一個(gè)實(shí)現(xiàn)了代碼塊計(jì)時(shí)功能的上下文管理器例子：

import time
from contextlib import contextmanager

@contextmanager
def timethis(label):
    start = time.time()
    try:
        yield
    finally:
        end = time.time()
        print('{}: {}'.format(label, end - start))

# Example use
with timethis('counting'):
    n = 10000000
    while n > 0:
        n -= 1

在函數(shù)?<span class="pre" style="box-sizing: border-box;">timethis()</span>?中，<span class="pre" style="box-sizing: border-box;">yield</span>?之前的代碼會(huì)在上下文管理器中作為?<span class="pre" style="box-sizing: border-box;">__enter__()</span>?方法執(zhí)行， 所有在?<span class="pre" style="box-sizing: border-box;">yield</span>?之后的代碼會(huì)作為?<span class="pre" style="box-sizing: border-box;">__exit__()</span>?方法執(zhí)行。 如果出現(xiàn)了異常，異常會(huì)在yield語(yǔ)句那里拋出。

下面是一個(gè)更加高級(jí)一點(diǎn)的上下文管理器，實(shí)現(xiàn)了列表對(duì)象上的某種事務(wù)：

@contextmanager
def list_transaction(orig_list):
    working = list(orig_list)
    yield working
    orig_list[:] = working

這段代碼的作用是任何對(duì)列表的修改只有當(dāng)所有代碼運(yùn)行完成并且不出現(xiàn)異常的情況下才會(huì)生效。 下面我們來(lái)演示一下：

>>> items = [1, 2, 3]
>>> with list_transaction(items) as working:
...     working.append(4)
...     working.append(5)
...
>>> items
[1, 2, 3, 4, 5]
>>> with list_transaction(items) as working:
...     working.append(6)
...     working.append(7)
...     raise RuntimeError('oops')
...
Traceback (most recent call last):
    File "<stdin>", line 4, in <module>
RuntimeError: oops
>>> items
[1, 2, 3, 4, 5]
>>>

討論

通常情況下，如果要寫(xiě)一個(gè)上下文管理器，你需要定義一個(gè)類(lèi)，里面包含一個(gè)?<span class="pre" style="box-sizing: border-box;">__enter__()</span>?和一個(gè)<span class="pre" style="box-sizing: border-box;">__exit__()</span>?方法，如下所示：

import time

class timethis:
    def __init__(self, label):
        self.label = label

    def __enter__(self):
        self.start = time.time()

    def __exit__(self, exc_ty, exc_val, exc_tb):
        end = time.time()
        print('{}: {}'.format(self.label, end - self.start))

盡管這個(gè)也不難寫(xiě)，但是相比較寫(xiě)一個(gè)簡(jiǎn)單的使用?<span class="pre" style="box-sizing: border-box;">@contextmanager</span>?注解的函數(shù)而言還是稍顯乏味。

<span class="pre" style="box-sizing: border-box;">@contextmanager</span>?應(yīng)該僅僅用來(lái)寫(xiě)自包含的上下文管理函數(shù)。 如果你有一些對(duì)象(比如一個(gè)文件、網(wǎng)絡(luò)連接或鎖)，需要支持?<span class="pre" style="box-sizing: border-box;">with</span>?語(yǔ)句，那么你就需要單獨(dú)實(shí)現(xiàn)?<span class="pre" style="box-sizing: border-box;">__enter__()</span>?方法和?<span class="pre" style="box-sizing: border-box;">__exit__()</span>?方法。