Python高级技术之:`Python`的`async` `context manager`:`__aenter__`和`__aexit__`的用法。

各位老铁,早上好!今天咱们聊聊Python异步上下文管理器,保证让你们听完之后,感觉自己又变强了。

啥是异步上下文管理器?简单来说,就是能用 async with 语句管理资源的家伙。这玩意儿在处理异步编程中的资源获取和释放时特别好使,比如异步数据库连接、异步文件操作等等。

咱们今天要深入探讨的就是实现异步上下文管理器的关键:__aenter____aexit__ 这两个魔法方法。

1. 铺垫:先回忆一下同步上下文管理器

在深入异步之前,咱们先来回忆一下同步上下文管理器,也就是我们平时常用的 with 语句。

class MyContextManager:
    def __enter__(self):
        print("Entering the context")
        return self  # 可以返回任何东西,这里返回self
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Exiting the context")
        if exc_type:
            print(f"An exception occurred: {exc_type}, {exc_val}")
        return False # 如果返回True,则会阻止异常传播

with MyContextManager() as obj:
    print("Inside the context")
    # raise ValueError("Something went wrong") # 取消注释看看效果
print("Outside the context")

运行结果大概是这样:

Entering the context
Inside the context
Exiting the context
Outside the context

如果把 raise ValueError("Something went wrong") 那行取消注释,结果会变成:

Entering the context
Inside the context
Exiting the context
An exception occurred: <class 'ValueError'>, Something went wrong
Traceback (most recent call last):
  File "...", line ..., in <module>
    raise ValueError("Something went wrong") # 取消注释看看效果
ValueError: Something went wrong
Outside the context

可以看到,即使发生了异常,__exit__ 方法仍然会被执行,保证资源得到清理。如果 __exit__ 返回 True,异常就会被吞掉,不会向上抛出(通常不建议这么做)。

2. 异步上下文管理器:__aenter____aexit__ 登场

异步上下文管理器和同步上下文管理器类似,只不过需要使用 async def 定义 __aenter____aexit__ 方法。

  • __aenter__(self): 进入上下文时被调用,必须是一个 async 函数。它的返回值会赋值给 async with 语句中的 as 变量。
  • __aexit__(self, exc_type, exc_val, exc_tb): 退出上下文时被调用,也必须是一个 async 函数。它接收三个参数:异常类型、异常值和异常回溯信息。如果没有异常发生,这三个参数都为 None

下面是一个简单的例子:

import asyncio

class AsyncContextManager:
    async def __aenter__(self):
        print("Entering the async context")
        await asyncio.sleep(1) # 模拟一些异步操作
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        print("Exiting the async context")
        if exc_type:
            print(f"An exception occurred: {exc_type}, {exc_val}")
        await asyncio.sleep(1) # 模拟一些异步清理操作

async def main():
    async with AsyncContextManager() as obj:
        print("Inside the async context")
        await asyncio.sleep(0.5)
        # raise ValueError("Async context error") # 取消注释看看效果
    print("Outside the async context")

if __name__ == "__main__":
    asyncio.run(main())

运行结果:

Entering the async context
Inside the async context
Exiting the async context
Outside the async context

如果把 raise ValueError("Async context error") 那行取消注释,结果会变成:

Entering the async context
Inside the async context
Exiting the async context
An exception occurred: <class 'ValueError'>, Async context error
Traceback (most recent call last):
  File "...", line ..., in main
    raise ValueError("Async context error") # 取消注释看看效果
ValueError: Async context error
Outside the async context

3. 实战:异步文件操作

我们来创建一个异步文件读写的上下文管理器。

import asyncio
import aiofiles # 需要安装:pip install aiofiles

class AsyncFile:
    def __init__(self, filename, mode='r'):
        self.filename = filename
        self.mode = mode
        self.file = None

    async def __aenter__(self):
        self.file = await aiofiles.open(self.filename, self.mode)
        return self.file

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.file:
            await self.file.close()
        if exc_type:
            print(f"File operation error: {exc_type}, {exc_val}")

async def main():
    async with AsyncFile('my_async_file.txt', 'w') as f:
        await f.write("Hello, async world!n")
        await f.write("This is a test.n")

    async with AsyncFile('my_async_file.txt', 'r') as f:
        contents = await f.read()
        print(f"File contents:n{contents}")

if __name__ == "__main__":
    asyncio.run(main())

这个例子使用了 aiofiles 库,它提供了异步的文件操作。 __aenter__ 方法打开文件,__aexit__ 方法关闭文件,保证文件资源在使用完毕后被释放。

4. 更复杂的例子:异步数据库连接池

接下来,我们来一个更复杂的例子,创建一个异步数据库连接池的上下文管理器。 这里我们使用 asyncpg 库,它是一个高性能的 PostgreSQL 异步客户端。

import asyncio
import asyncpg # 需要安装:pip install asyncpg

class AsyncConnectionPool:
    def __init__(self, dsn, min_size=1, max_size=10):
        self.dsn = dsn
        self.min_size = min_size
        self.max_size = max_size
        self.pool = None

    async def __aenter__(self):
        self.pool = await asyncpg.create_pool(self.dsn, min_size=self.min_size, max_size=self.max_size)
        return self.pool

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.pool:
            await self.pool.close()
        if exc_type:
            print(f"Database error: {exc_type}, {exc_val}")

async def main():
    dsn = 'postgresql://user:password@host:port/database' # 替换成你的数据库连接信息

    async with AsyncConnectionPool(dsn) as pool:
        async with pool.acquire() as conn:
            # 执行一些数据库操作
            result = await conn.fetchval("SELECT 1 + 1")
            print(f"Result of 1 + 1: {result}")

            # 插入数据
            await conn.execute("INSERT INTO my_table (name) VALUES ($1)", "Async Context Test")

            # 查询数据
            rows = await conn.fetch("SELECT id, name FROM my_table")
            for row in rows:
                print(f"ID: {row['id']}, Name: {row['name']}")

if __name__ == "__main__":
    asyncio.run(main())

这个例子创建了一个异步连接池,__aenter__ 方法创建连接池,__aexit__ 方法关闭连接池。在 async with 语句块中,我们可以安全地获取连接并执行数据库操作,不用担心连接泄漏。

5. 异常处理的细节

__aexit__ 方法的参数 exc_type, exc_val, exc_tb 包含了异常的信息。

  • exc_type: 异常类型,例如 ValueError, TypeError 等。
  • exc_val: 异常实例,例如 ValueError("Something went wrong")
  • exc_tb: 异常回溯信息,是一个 traceback 对象。

__aexit__ 方法的返回值也很重要:

  • 如果返回 True (或者任何 truthy 值),则表示异常已经被处理,不会向上抛出。
  • 如果返回 False (或者任何 falsy 值,包括 None),则表示异常没有被处理,会向上抛出。

注意: 除非你真的知道自己在做什么,否则最好不要在 __aexit__ 中吞掉异常,让异常继续传播,方便调试和错误处理。

6. 总结:__aenter____aexit__ 的用法

咱们用一个表格来总结一下 __aenter____aexit__ 的用法:

方法 作用 参数 返回值
__aenter__ 进入异步上下文时被调用,负责获取资源(例如打开文件、获取数据库连接等)。 self 返回值会被赋值给 async with 语句中的 as 变量。通常返回 self,也可以返回其他对象。
__aexit__ 退出异步上下文时被调用,负责释放资源(例如关闭文件、释放数据库连接等)。无论是否发生异常,都会被执行。 self, exc_type, exc_val, exc_tbexc_type, exc_val, exc_tb 包含了异常信息。如果没有异常发生,这三个参数都为 None 如果返回 True (或任何 truthy 值),则表示异常已经被处理,不会向上抛出。 如果返回 False (或任何 falsy 值,包括 None),则表示异常没有被处理,会向上抛出。

7. 最佳实践

  • 始终确保资源被释放: 无论是否发生异常,__aexit__ 都应该被执行,所以务必在 __aexit__ 中释放资源。
  • 小心处理异常: 除非你真的想吞掉异常,否则不要在 __aexit__ 中返回 True
  • 使用合适的库: 对于异步文件操作,可以使用 aiofiles。对于异步数据库操作,可以使用 asyncpg, aiosqlite 等。
  • 避免阻塞操作: __aenter____aexit__ 内部应该避免执行任何阻塞操作,否则会影响事件循环的性能。

8. 总结

异步上下文管理器是 Python 异步编程中一个非常重要的概念,它可以帮助我们更方便地管理资源,提高代码的可读性和可维护性。 __aenter____aexit__ 是实现异步上下文管理器的关键,掌握它们,你就能写出更优雅、更健壮的异步代码。

好了,今天的分享就到这里。希望大家能够有所收获,下次再见!

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注