python中的上下文管理器以及python内建模块contextlib的contextmanager方法

上下文管理器

上下文管理器是实现了上下文管理协议的对象，其特有的语法是“with …as”。主要用于保存和恢复各种全局状态，关闭文件等，并为try…except…finally提供了一个方便使用的封装。

上下文管理协议具体来说就是在类里面实现以下两个方法：

_enter_(): 从该方法进入运行时上下文，并返回当前对象或者与运行时上下文相关的其他对象。如果with语句有as关键词存在，返回值会绑定在as后的变量上。

_exit_(exc_type, exc_val, exc_tb): 退出运行时上下文，return True 如果 with 执行体有异常，则不会继续向上抛出异常；return Flase 如果 with 执行体有异常，则继续向上抛出异常。如果在执行with语句体时发生异常，那退出时参数会包括异常类型、异常值、异常追踪信息，否则，3个参数都是None。
下面，我们可以自己定义一个类来通过with进行上下文管理：

class MyContextManager:

    def __enter__(self):
        print("connect to contextmanager")
        return self  # return返回的可以是对象，会绑定给as后面的变量

    def __exit__(self, exc_type, exc_val, exc_tb):
        """__exit__方法中：
                1.如果return False：执行query_data方法有异常，则异常继续往上抛
                2.如果return True：执行query_data方法有异常，则异常不会往上抛
        """
        print('close the contextmanager')
        # return False
        return True

    def query_data(self):
        print('query data')


with MyContextManager() as m:
    m.query_data()

执行以上代码，运行结果如下：

下面，我们在类的query_data方法中增加一个报错的a，当__exit__方法 return True 时，如下：

class MyContextManager:

    def __enter__(self):
        print("connect to contextmanager")
        return self  # return返回的可以是对象，会绑定给as后面的变量

    def __exit__(self, exc_type, exc_val, exc_tb):
        """__exit__方法中：
                1.如果return False：执行query_data方法有异常，则异常继续往上抛
                2.如果return True：执行query_data方法有异常，则异常不会往上抛
        """
        print('close the contextmanager')
        # return False
        return True

    def query_data(self):
        a
        print('query data')


with MyContextManager() as m:
    m.query_data()

运行结果如下：并不会执行 query_data() 方法

当__exit__方法 return False 时，如下：

class MyContextManager:

    def __enter__(self):
        print("connect to contextmanager")
        return self  # return返回的可以是对象，会绑定给as后面的变量

    def __exit__(self, exc_type, exc_val, exc_tb):
        """__exit__方法中：
                1.如果return False：执行query_data方法有异常，则异常继续往上抛
                2.如果return True：执行query_data方法有异常，则异常不会往上抛
        """
        print('close the contextmanager')
        return False
        # return True

    def query_data(self):
        a
        print('query data')


with MyContextManager() as m:
    m.query_data()

运行结果如下：会抛异常，

实际应用示例：

创建一个上下文管理器，这个上下文管理器将会创建一个SQLite数据库连接，当任务处理完毕，将会将其关闭。

import sqlite3

class DataConn:
    def __init__(self,db_name):
        self.db_name = db_name

    def __enter__(self):
        self.conn = sqlite3.connect(self.db_name)
        return self.conn

    def __exit__(self,exc_type,exc_val,exc_tb):
        self.conn.close()
        if exc_val:
            raise

if __name__ == "__main__":
    db = "test/test.db"
    with DataConn(db) as conn:
        cursor = conn.cursor()

在上面的代码中，我们创建了一个类，获取到SQLite数据库文件的路径。__enter__方法将会自动执行，并返回数据库连接对象。现在我们已经获取到数据库连接对象，然后我们创建光标，向数据库写入数据或者对数据库进行查询。当我们退出with语句的时候，它将会调用__exit__方法用于执行和关闭这个连接。

Python的内建模块 contextlib

对于上下文的管理，python也提供了内建的模块contextlib来实现相同的机制，而且这种通过生成器和装饰器实现的上下文管理器，看起来比with语句和手动实现上下文管理协议更优雅。

from contextlib import contextmanager


class MyContextManager:
    def query_data(self):
        # a
        print('query data')


@ contextmanager  # 有了contextmanager我们就不需要手动实现enter和exit方法
def make_context_manager():
    print('connect to contextmanager')
    yield MyContextManager()  # yield就相当于return，但是不会结束函数，而是暂时挂起，如果我们后续需要的话，继续调用这个函数，会从上次挂起的地方继续执行
    print('close the contextmanager')


with make_context_manager() as mc:
    mc.query_data()