Python中@contextmanager的用法

在Python中with的用法一文中已经写到了什么是上下文管理器，以及如何创建一个符合上下文协议的自定义类。那么如果我们需要将一个非自定义的类改成一个上下文管理器又该怎样实现？

我们可以为需要改写的普通类创建一个子类，在其中添加上__enter__和__exit__方法通过继承父类的方式实现上下文管理器。但是通过contextlib模块下提供的@contextmanager装饰器，我们能够更方便的将一个普通类变成上下文管理器。

@contextmanager

@contextmanager通过将一个函数变成生成器的方式来为普通类添加进入和退出时的处理代码，从而实现了将普通类变成一个上下文管理器。

示例1

from contextlib import contextmanager
class File():
    def query(self):
        print('查询文件')

@contextmanager
def open():
    print('打开文件')
    yield File()
    print('关闭文件')

with open() as f:
    f.query()
    
# 结果 打开文件 查询文件 关闭文件

执行流程

①with语句调用open函数=>②执行open中yield之前的代码(打开文件)=>③执行yield语句中的代码(File())=>④执行with语句中的代码(f.query)=>⑤执行yiled语句后的代码(关闭文件)

注意事项

• 此时的with语句后是调用被修饰的函数，而不是实例化上下文管理器对象
• with open() as f 中的f是yield语句中的内容，上述示例中即所创建的File对象
• yield前、中、后的代码都可以省略不写

示例2

如果需要改成上下文管理器的类允许被改写，或者使用问题引出中所提到的构建子类的方式，那么我们可以将@contextmanager所修饰的函数放在类的内部作为一个类方法来使用：

from contextlib import contextmanager
class DataBase():
    def query(self):
        print('写入操作')
    @contextmanager
    def open(self):
        try:
            yield
            self.commit()
        except Exception as e:
            self.rollback()
            raise e
#学习中遇到问题没人解答？小编创建了一个Python学习交流群：711312441
db = DataBase()
with db.open():
    db.query()

执行流程

①with语句调用open函数=>②执行try异常判断=>③执行with语句中的写入操作=>④执行数据库的提交操作=>⑤如果出现异常执行数据库回滚操作

总结

• 在思考过程中我们可以假设with语句中的内容整体移到了yield语句后，从而将with语句的执行简化为@contextmanager所修饰的函数的执行
• 当我们需要在一段代码的执行前后加上其他的处理代码时都可以使用@contextmanager来完成，并不需要一定是对上下文管理器的操作