Python上下文管理器

最近学习python异步编程时，了解到异步上下文管理器，云里雾里得先一点点学，什么是上下文管理器？

作为一个程序员或者准程序员，一定听说过内存泄露，内存泄露的根本原因在于

创建了某个对象，却没有及时的释放掉，直到程序结束前，这个未被释放的对象一直占着内存。

那这样有什么问题吗？其实量少的话还好，如果量大那么就会直接把内存占满，导致程序被kill掉，这就是内存泄露。那内存泄露和上下文管理器有什么关系呢？

内存泄露和上下文管理器

首先，现在我们使用的很多高级编程语言已经不需要让我们过多的去关注内存的问题了，但是在某些情况下还是需要我们编写程序来关闭或释放某些对象。而最常见的就是文件操作。

在任何一门编程语言中，文件的输入输出、数据库的连接断开等，都是很常见的资源管理操作。但资源都是有限的，在写程序时，我们必须保证这些资源在使用过后得到释放，不然就容易造成资源泄露，轻者使得系统处理缓慢，重则会使系统崩溃。

比如下面这个例子，我们打开了1千万个文件，进行写入操作。但是没有及时的关闭文件，如果运行就会报错。

for x in range(10000000): 
    f = open('test.txt', 'w')
    f.write('后厂程序员') 
    # 没有写文件的关闭操作

这就是一个典型的资源泄露的例子。

因为程序中同时打开了太多的文件，占据了太多的资源，造成系统崩溃。

为了解决这个问题，不同的编程语言都引入了不同的机制。

而在 Python 中，对应的解决方式便是上下文管理器（context manager）。

上下文管理器，能够帮助你自动分配并且释放资源，其中最典型的应用便是 with 语句。

如果我们把上面的代码改成with语句的形式：

for x in range(10000000):
    with open('test.txt', 'w') as f:
        f.write('后厂程序员')

这样我们每次打开文件，操作完成后这个文件便会自动关闭，这样相应的资源也可以得到释放，防止资源泄露。当然with 语句的代码，也可以用下面的形式表示：

f = open('test.txt', 'w')
try:
     f.write('后厂程序员')
finally:
     f.close()

其中 finally 是哪怕在写入文件时发生错误异常，也可以保证该文件最终被关闭。不过我一般更倾向于使用 with 语句。

当然with语句的应用不仅于此，比如我想要获取一个锁，执行相应的操作，完成后再释放，那么代码就可以写成下面这样：

some_lock = threading.Lock()
with somelock:
    ...

我们可以从这两个例子中看到，with 语句的使用，可以简化了代码，有效避免资源泄露的发生。
其实，我们很多时候在学习文件操作的时候都学到了with语句的使用，但是很多课程也就仅限停留在with语句上，只知道在文件操作上使用with语句很方便，但是并不清楚with语句的实现和上下文管理器的原理。因此接下来我们通过上下文管理器的实现来更好的理解它们。

上下文管理器的实现

首先我们想要实现上下文管理器，那么我们要先知道上下文管理器协议。其实上下文管理器的协议也非常简单，就是必须在一个类中实现__enter__()和__exit__()两个方法，然后这个类的实例就是一个上下文管理器。

其中，方法__enter__()返回需要被管理的资源，方法__exit__()里通常会存在一些释放、清理资源的操作，比如关闭文件、关闭数据库连接等。

基于类的上下文管理器

# 定一个类
class FileManager:
    def __init__(self, name, mode):
        print('calling __init__ method')
        self.name = name
        self.mode = mode 
        self.file = None
    # 在类中实现__enter__，并完成文件的打开操作
    def __enter__(self):
        print('calling __enter__ method')
        self.file = open(self.name, self.mode)
        return self.file
    # 在类中实现__exit__，并完成文件的关闭操作
    def __exit__(self, exc_type, exc_val, exc_tb):
        print('calling __exit__ method')
        if self.file:
            self.file.close()

# 使用with语句来执行上下文管理器
with FileManager('test.txt', 'w') as f:
    print('ready to write to file')
    f.write('hello world')

# 当我们用 with 语句，执行这个上下文管理器时：
# 1. 方法`__init__()`被调用，程序初始化对象 FileManager，使得文件名（name）是"test.txt"，文件模式 (mode) 是'w'；
# 2. 方法`__enter__()`被调用，文件“test.txt”以写入的模式被打开，并且返回 FileManager 对象赋予变量 f；
# 3. 字符串“hello world”被写入文件“test.txt”；
# 4. 方法`__exit__()`被调用，负责关闭之前打开的文件流。

# 最终的输出结果：
calling __init__ method
calling __enter__ method
ready to write to file
calling __exit__ meth

另外，exit()方法中的参数“exc_type, exc_val, exc_tb”，分别表示 exception_type、exception_value 和 traceback。当我们执行含有上下文管理器的 with 语句时，如果有异常抛出，异常的信息就会包含在这三个变量中，传入方法__exit__()。

比如像下面这样：

class Foo:
    def __init__(self):
        print('__init__ called')        

    def __enter__(self):
        print('__enter__ called')
        return self
    # 在__exit__方法中捕获并输出异常信息 
    def __exit__(self, exc_type, exc_value, exc_tb):
        print('__exit__ called')
        if exc_type:
            print(f'exc_type: {exc_type}')
            print(f'exc_value: {exc_value}')
            print(f'exc_traceback: {exc_tb}')
            print('exception handled')
        return True # 异常处理后必须返回True

# 调用并手动抛出异常
with Foo() as obj:
    raise Exception('exception raised').with_traceback(None)

# 输出
# __init__ called
# __enter__ called
# __exit__ called
# exc_type: 
# exc_value: exception raised
# exc_traceback: 
# exception handled

需要注意，如果方法__exit__()没有返回 True，异常仍然会被抛出。因此，如果异常已经被处理了则必须在__exit__()方法中返回True。

同样的，数据库的连接操作，也可以用上下文管理器来表示：

class DBCM: 
   # 负责对数据库进行初始化，也就是将主机名、接口（这里是 localhost 和 8080）分别赋予变量 hostname 和 port；
    def __init__(self, hostname, port): 
        self.hostname = hostname 
        self.port = port 
        self.connection = None
   # 连接数据库，并且返回对象 DBCM；
    def __enter__(self): 
        self.connection = DBClient(self.hostname, self.port) 
        return self
   # 负责关闭数据库的连接
    def __exit__(self, exc_type, exc_val, exc_tb): 
        self.connection.close() 
  
with DBCM('localhost', '8080') as db_client: 
    ....

这样只要你写完了 DBCM 这个类，那么在程序每次连接数据库时，我们都只需要简单地调用 with 语句即可，并不需要关心数据库的关闭、异常等等，大大提高了开发的效率。

基于生成器的上下文管理器
Python中的上下文管理器除了基于类，还可以基于生成器来实现。

使用装饰器 contextlib.contextmanager，来自定义基于生成器的上下文管理器，用以支持 with 语句。

还是拿前面的类上下文管理器 FileManager 来说，我们也可以用下面形式来表示：

from contextlib import contextmanager

@contextmanager
def file_manager(name, mode):
    try:
        f = open(name, mode)
        yield f
    finally:
        f.close()
        
with file_manager('test.txt', 'w') as f:
    f.write('hello world')

这段代码中，函数 file_manager() 是一个生成器，当我们执行 with 语句时，便会打开文件，并返回文件对象 f；

当 with 语句执行完后，finally block 中的关闭文件操作便会执行。

注意：基于生成器定义的上下文管理需要使用装饰器 @contextmanager，不再使用生成器协议方法。

基于类的上下文管理器和基于生成器的上下文管理器，这两者在功能上是一致的。

这样小伙伴是不是对python中的with语句有了更深的认识，并且在了解了上下文管理器后，未来在开发中也可以自定义上下文管理器来实现资源的上下文管理了。

下面附带一个最近学的异步上下文管理器test：

import asyncio
# 此种对象通过定义__aenter__()和__aexit__()方法来对async with语句中的环境进行控制。

class AsyncContextManager:
    def __init__(self, conn):
        self.conn = conn

    async def do_something(self):
        # 异步操作数据库
        return 666

    async def __aenter__(self):
        # 异步链接数据库
        self.conn = await asyncio.sleep(1)
        return self

    async def __aexit__(self, exc_type, exc, tb):
        await asyncio.sleep(1)

async def fun():
    async with AsyncContextManager as f:
        result = await f.do_something()
        print(result)
asyncio.run(fun())