使用生成器語法編寫自己的上下文管理器

由於 contextlib.contextmanager 裝飾器,也可以使用生成器語法編寫上下文管理器:

import contextlib

@contextlib.contextmanager
def context_manager(num):
    print('Enter')
    yield num + 1
    print('Exit')

with context_manager(2) as cm:
    # the following instructions are run when the 'yield' point of the context
    # manager is reached.
    # 'cm' will have the value that was yielded
    print('Right in the middle with cm = {}'.format(cm))

生產:

Enter
Right in the middle with cm = 3
Exit

裝飾器通過將生成器轉換為一個來簡化編寫上下文管理器的任務。yield 表示式之前的所有內容都成為 __enter__ 方法,得到的值變為生成器返回的值(可以繫結到 with 語句中的變數),yield 表示式之後的所有內容都變為 __exit__ 方法。

如果需要由上下文管理器處理異常,則可以在生成器中寫入 try..except..finally 塊,並且 with 塊中引發的任何異常都將由此異常塊處理。

@contextlib.contextmanager
def error_handling_context_manager(num):
    print("Enter")
    try:
        yield num + 1
    except ZeroDivisionError:
        print("Caught error")
    finally:
        print("Cleaning up")
    print("Exit")

with error_handling_context_manager(-1) as cm:
    print("Dividing by cm = {}".format(cm))
    print(2 / cm)

這會產生:

Enter
Dividing by cm = 0
Caught error
Cleaning up
Exit