使用生成器语法编写自己的上下文管理器

由于 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