异步上下文管理器与异步迭代器

Python并发编程专题 · async with和async for的底层协议

专题：Python并发编程系统学习

关键词：Python, 并发编程, 异步上下文管理器, 异步迭代器, async with, async for, 异步生成器

一、同步上下文管理器回顾

在深入异步上下文管理器之前，我们先回顾一下Python中同步上下文管理器的基础知识。上下文管理器是Python中用于资源管理的强大工具，它确保资源在使用后能被正确释放，即使过程中发生了异常也不例外。

同步上下文管理器通过两个魔术方法实现：__enter__和__exit__。__enter__方法在进入with语句块时被调用，通常用于获取资源并返回资源对象；__exit__方法在退出with语句块时被调用，无论是否发生异常，通常用于释放资源。如果__exit__返回True，则异常会被压制（即不会继续向外传播），这在某些特定场景下非常有用。

with语句的执行流程可以概括为：首先调用__enter__获取资源，然后将返回值赋给as子句指定的变量，接着执行语句块中的代码，最后无论语句块是否正常执行完毕或抛出异常，都会调用__exit__进行清理。这种资源获取即初始化（RAII）模式在文件操作、锁管理、数据库连接等场景中得到了广泛应用。

class FileManager:
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.file.close()
        return False  # 不压制异常

# 使用示例
with FileManager('test.txt', 'w') as f:
    f.write('Hello, World!')

此外，Python标准库还提供了contextlib.contextmanager装饰器，它允许我们通过生成器函数的方式快速创建上下文管理器，而不需要编写一个完整的类。这种基于生成器的方式更加简洁，适合简单的资源管理场景。

from contextlib import contextmanager

@contextmanager
def managed_file(filename, mode):
    f = open(filename, mode)
    try:
        yield f
    finally:
        f.close()

掌握同步上下文管理器是理解异步上下文管理器的前置基础。二者的核心设计理念一致——资源的获取和释放，区别仅在于操作是否涉及异步I/O调用。当资源获取或释放操作涉及网络请求、磁盘I/O等阻塞操作时，就需要使用异步上下文管理器来避免阻塞事件循环。

二、异步上下文管理器协议

Python 3.5引入的async with语句和相应的异步上下文管理器协议，是asyncio生态系统的基石之一。异步上下文管理器的协议由两个协程方法组成：__aenter__和__aexit__，它们分别对应同步中的__enter__和__exit__。

__aenter__是一个异步方法，在进入async with语句块时被调用，返回一个awaitable对象。它通常用于执行异步的资源获取操作，比如建立数据库连接、打开HTTP会话等。__aexit__同样是一个异步方法，在退出async with语句块时被调用，接收三个参数：异常类型、异常值和回溯信息。当语句块没有发生异常时，这三个参数均为None。

"async with"的工作原理与同步的"with"非常相似，差异在于每一步都涉及await操作。Python运行时会在进入时await __aenter__的返回值，在退出时await __aexit__的返回值。这种await机制确保异步资源管理器不会阻塞事件循环。

class AsyncResource:
    async def __aenter__(self):
        await self.open()
        return self

    async def __aexit__(self, exc_type, exc, tb):
        await self.close()

使用异步上下文管理器时，需要配合async with语句，并且整个语句必须在异步函数（即async def定义的函数）中使用。不能将async with放在同步函数中执行，否则会引发语法错误。

# 正确的使用方式
async def use_resource():
    async with AsyncResource() as res:
        await res.do_something()

# 错误的用法（会报语法错误）
# def sync_func():
#     async with AsyncResource() as res:   # SyntaxError
#         pass

关键点：__aenter__必须返回一个可等待对象（awaitable），该对象通常返回资源实例（即self）。__aexit__的返回值如果是True，则异常会被压制，与同步版本的行为一致。

常见的异步资源管理场景包括：aiohttp的ClientSession管理、asyncpg的数据库连接池、aioredis的连接管理、aiofiles的异步文件操作等。在这些场景中，资源的建立和销毁都涉及网络I/O或磁盘I/O，必须使用异步方式来避免阻塞事件循环。

三、@asynccontextmanager简化实现

如果每次编写异步上下文管理器都需要定义一个类并实现__aenter__和__aexit__两个方法，代码会显得过于冗长。幸运的是，Python标准库中的contextlib模块提供了asynccontextmanager装饰器，它允许我们通过异步生成器函数快速创建异步上下文管理器。

@asynccontextmanager装饰一个异步生成器函数，该函数在yield之前是__aenter__的逻辑，yield之后是__aexit__的逻辑。如果资源获取过程中发生异常，异常会在yield语句处重新抛出，因此需要使用try/finally结构来确保资源被正确释放。

需要注意的是，异步生成器中的yield只能出现一次。如果多次yield会导致运行时抛出StopIteration异常。此外，yield语句本身不是异步的，不能直接在其前面加await。

from contextlib import asynccontextmanager

@asynccontextmanager
async def managed_resource():
    resource = await open_resource()
    try:
        yield resource
    finally:
        await resource.close()

以下是一个使用@asynccontextmanager管理aiohttp会话的实际示例：

import aiohttp
from contextlib import asynccontextmanager

@asynccontextmanager
async def http_session(base_url=''):
    session = aiohttp.ClientSession(base_url=base_url)
    try:
        yield session
    finally:
        await session.close()

async def fetch_data():
    async with http_session() as session:
        async with session.get('https://api.example.com/data') as resp:
            return await resp.json()

技巧：使用@asynccontextmanager时，务必使用try/finally块包装yield语句。这确保了即使在异步操作中发生异常，资源也能被正确释放。这是编写健壮异步代码的重要实践。

四、异步迭代器协议

异步迭代器是Python异步编程中的另一个核心协议。与普通的同步迭代器类似，异步迭代器通过__aiter__和__anext__两个方法来定义。其中__aiter__返回异步迭代器对象自身，而__anext__返回一个可等待对象（awaitable），当没有更多元素时抛出StopAsyncIteration异常。

异步迭代器的主要应用场景是那些需要从异步数据源（如网络流、数据库游标、消息队列等）逐个获取数据的情况。在这些场景中，获取下一个数据元素可能涉及I/O等待，因此需要异步操作来避免阻塞。

import asyncio

class AsyncCounter:
    def __init__(self, limit):
        self.limit = limit
        self.current = 0

    def __aiter__(self):
        return self

    async def __anext__(self):
        if self.current >= self.limit:
            raise StopAsyncIteration
        await asyncio.sleep(0.1)
        value = self.current
        self.current += 1
        return value

在使用async for循环遍历异步迭代器时，Python会在每次迭代中自动await __anext__方法的返回值，直到它抛出StopAsyncIteration异常为止。这个过程完全由Python运行时管理，开发者无需手动处理异常或await操作。

# 使用 async for 遍历异步迭代器
async def main():
    async for num in AsyncCounter(5):
        print(num)

# 等价的手动实现方式
async def manual_iteration():
    it = AsyncCounter(5).__aiter__()
    try:
        while True:
            num = await it.__anext__()
            print(num)
    except StopAsyncIteration:
        pass

特性	同步迭代器	异步迭代器
获取迭代器	`__iter__()`	`__aiter__()`
获取下一项	`__next__()`	`__anext__()`（返回awaitable）
结束信号	`StopIteration`	`StopAsyncIteration`
遍历语法	`for x in iterable`	`async for x in async_iterable`

关键区别：异步迭代器的__anext__方法必须是一个协程（即使用async def定义或返回一个awaitable对象），因为获取下一个元素可能涉及I/O等待。而__aiter__方法可以是普通方法，因为它通常只是返回self，不涉及任何异步操作。

五、异步生成器

异步生成器是Python 3.6引入的强大特性，它结合了异步函数和生成器的能力。通过使用async def和yield的组合，我们可以定义一个异步生成器函数，它会在每次产生值时可能进行await操作。异步生成器自动实现了异步迭代器协议，因此可以直接在async for循环中使用。

异步生成器的核心优势在于：它让我们能够以近乎同步代码的简洁风格编写异步数据生成逻辑，而无需手动实现__aiter__和__anext__方法。Python运行时会自动为异步生成器创建必要的迭代器协议方法。

async def async_range(n):
    for i in range(n):
        await asyncio.sleep(0.1)
        yield i

以下是一个更实用的示例——异步读取并逐行处理大型文件的生成器：

import aiofiles

async def read_lines(filepath):
    async with aiofiles.open(filepath, mode='r') as f:
        async for line in f:
            await asyncio.sleep(0)  # 让出控制权
            yield line.strip()

async def process_file():
    async for line in read_lines('data.txt'):
        if line:
            print(f'处理行: {line}')

异步生成器的生命周期管理需要特别注意三点：第一，异步生成器对象如果未被完全消费完毕就被垃圾回收，其aclose()方法会被自动调用，以释放资源；第二，可以通过asend()、athrow()和aclose()方法与异步生成器进行交互，这些方法都返回awaitable对象；第三，异步生成器不能与同步的yield from混用，但可以使用async for在另一个异步生成器中消费它。

最佳实践：在设计异步API时，优先考虑使用异步生成器来替代手动实现的异步迭代器类。异步生成器的代码更简洁、更易读，且自动处理了迭代器协议的大部分细节。当迭代逻辑本身不复杂但涉及I/O操作时，异步生成器是最佳选择。

六、async for循环的本质

async for是Python为异步迭代提供的语法糖，它在底层完成了一系列复杂的异步操作。理解async for的本质对于编写正确的异步代码至关重要。

当Python遇到async for x in async_iterable: ...语句时，它执行的操作序列如下：第一步，调用async_iterable.__aiter__()获取异步迭代器对象；第二步，在每次循环迭代中，await iterator.__anext__()获取下一个值；第三步，当捕获到StopAsyncIteration异常时，正常结束循环。这个过程与同步for循环在结构上完全对称，唯一的区别在于迭代的每一步都涉及await操作。

需要特别注意的是，async for只能在async def定义的协程函数中使用。如果尝试在普通同步函数中使用async for，Python会报SyntaxError。这是因为async for中的await操作需要一个事件循环上下文。

# async for 的等价手动实现
async def async_for_demo():
    # async for num in AsyncCounter(5):
    #     print(num)
    #
    # 等价于：
    _iterator = AsyncCounter(5).__aiter__()
    while True:
        try:
            num = await _iterator.__anext__()
        except StopAsyncIteration:
            break
        else:
            print(num)

此外，async for还支持与async with类似的异常传播机制。如果在async for循环体内部抛出了异常，循环会立即终止，并且异常会向上传播。异步迭代器的资源清理（如关闭文件句柄或网络连接）通常通过实现__aenter__/__aexit__协议来配合async with使用，或者由异步生成器的垃圾回收机制自动处理。

# async for 与 async with 的组合使用
async def process_data():
    async with database.connect() as conn:
        async for row in conn.query('SELECT * FROM users'):
            process(row)

核心要点：async for的每一步迭代都涉及await操作。如果在迭代过程中其他协程正在运行，它们会在await点获得执行机会。这正是异步并发能够实现的关键——通过在每次迭代中主动让出控制权，事件循环得以调度其他任务。

七、实际应用场景

异步上下文管理器和异步迭代器在实际项目中的应用非常广泛。以下是几个典型的使用场景，展示了这些协议在实际开发中的强大能力。

1. aiohttp session管理

在使用aiohttp进行HTTP请求时，每个应用程序通常只需要创建一个ClientSession实例，并在整个应用生命周期中复用。使用async with管理session的创建和关闭是最佳实践，确保网络连接被正确释放。

import aiohttp
import asyncio

async def fetch_urls(urls):
    async with aiohttp.ClientSession() as session:
        for url in urls:
            async with session.get(url) as resp:
                yield await resp.text()

2. 异步数据库连接池

各种异步数据库驱动（如asyncpg、aiomysql、aioredis）都提供了基于async with的连接管理API。从连接池中获取连接和执行查询都通过异步上下文管理器进行，确保连接被归还到池中。

import asyncpg

async def query_database():
    pool = await asyncpg.create_pool(
        user='user', password='pass',
        database='db', host='localhost'
    )
    async with pool.acquire() as conn:
        async for row in conn.cursor('SELECT id, name FROM users'):
            print(f'用户: {row["name"]}')

3. 流式数据处理

在WebSocket通信、大文件上传/下载、实时数据流等场景中，数据是分批到达的。使用异步迭代器和异步生成器可以优雅地处理流式数据，以推送（push）或拉取（pull）模式逐块处理数据。

# WebSocket流式消息处理
async def handle_websocket(websocket):
    async for message in websocket:
        print(f'收到消息: {message}')
        await process_message(message)
        await websocket.send(f'已确认: {message}')

4. 自定义异步上下文管理器

在开发异步库或框架时，经常会需要实现自定义的资源管理器。以下是一个带超时控制的异步锁管理器示例：

class TimedLock:
    def __init__(self, timeout=5.0):
        self._lock = asyncio.Lock()
        self.timeout = timeout

    async def __aenter__(self):
        try:
            await asyncio.wait_for(
                self._lock.acquire(),
                timeout=self.timeout
            )
        except asyncio.TimeoutError:
            raise TimeoutError(
                f'获取锁超时（{self.timeout}秒）'
            )
        return self

    async def __aexit__(self, exc_type, exc, tb):
        self._lock.release()

实践建议：在编写异步代码时，应当将"使用async with管理资源"作为默认模式。无论是网络连接、文件句柄还是锁资源，async with都能确保即使在异常情况下资源也能被正确释放。结合异步生成器进行流式数据处理，可以最大化代码的简洁性和可维护性。

通过对比同步和异步版本的代码可以发现，异步上下文管理器和异步迭代器在设计模式上与它们的同步版本保持高度一致。这种一致性降低了学习曲线，让开发者能够将已有的同步编程经验迁移到异步领域。掌握这两个协议后，就可以阅读和使用绝大多数基于asyncio的第三方库，并编写自己的异步资源管理代码。