9.22 定義上下文管理器的簡單方法

問題

你想自己去實現一個新的上下文管理器，以便使用with語句。

解決方案

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

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

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

下面是一個更加高級一點的上下文管理器，實現了列表對象上的某種事務：

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

這段代碼的作用是任何對列表的修改只有當所有代碼運行完成并且不出現異常的情況下才會生效。 下面我們來演示一下：

>>> 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]
>>>

討論

通常情況下，如果要寫一個上下文管理器，你需要定義一個類，里面包含一個?<span class="pre" style="box-sizing: border-box;">__enter__()</span>?和一個<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))

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

<span class="pre" style="box-sizing: border-box;">@contextmanager</span>?應該僅僅用來寫自包含的上下文管理函數。 如果你有一些對象(比如一個文件、網絡連接或鎖)，需要支持?<span class="pre" style="box-sizing: border-box;">with</span>?語句，那么你就需要單獨實現?<span class="pre" style="box-sizing: border-box;">__enter__()</span>?方法和?<span class="pre" style="box-sizing: border-box;">__exit__()</span>?方法。