9.22 定義上下文管理器的簡單方法 · Python3 CookBook中文版

## 問題你想自己去實現一個新的上下文管理器，以便使用with語句。 ## 解決方案實現一個新的上下文管理器的最簡單的方法就是使用?`contexlib`?模塊中的?`@contextmanager`?裝飾器。下面是一個實現了代碼塊計時功能的上下文管理器例子： ~~~ 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 ~~~ 在函數?`timethis()`?中，`yield`?之前的代碼會在上下文管理器中作為?`__enter__()`?方法執行，所有在?`yield`?之后的代碼會作為?`__exit__()`?方法執行。如果出現了異常，異常會在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] >>> ~~~ ## 討論通常情況下，如果要寫一個上下文管理器，你需要定義一個類，里面包含一個?`__enter__()`?和一個`__exit__()`?方法，如下所示： ~~~ 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)) ~~~ 盡管這個也不難寫，但是相比較寫一個簡單的使用?`@contextmanager`?注解的函數而言還是稍顯乏味。 `@contextmanager`?應該僅僅用來寫自包含的上下文管理函數。如果你有一些對象(比如一個文件、網絡連接或鎖)，需要支持?`with`?語句，那么你就需要單獨實現?`__enter__()`?方法和?`__exit__()`?方法。