Python异步上下文管理器管理连接时是否需要并行任务？

我正在开发一个Python项目，负责管理供暖系统的连接。我有一个继承自HeatingSystemStore的ControllableHeatingSystemStore类，其中的异步上下文管理器open_connection会打开连接、启动后台任务获取供暖系统控制权，关闭连接时取消任务。

相关代码如下：

class ControllableHeatingSystemStore(
    HeatingSystemStore[HubControlConnection, ControllableHeatingSystem]
):
    @asynccontextmanager
    async def open_connection(
        self, connection: HubControlConnection, building_id: str
    ) -> AsyncGenerator[HubControlConnection, None]:
        async with super().open_connection(connection, building_id) as conn:

            async def await_acquire_control() -> None:
                await conn.acquire_control()
                self.add_heating_system(
                    building_id, CloudControllableHeatingSystem(conn)
                )

            async with asyncio.TaskGroup() as tg:
                t = tg.create_task(await_acquire_control())
                yield conn
                t.cancel()

我在代码库中频繁使用该模式，担心是否过于复杂。我知道可以将其封装为工具上下文管理器（如async with BackgroundTask(await_acquire_control())），但能否用单个任务处理？有没有更高效、更符合Python风格的方式来管理这些连接，无需将后台任务生命周期与上下文管理器绑定？希望得到相关见解与建议。

1. 用单个任务替代TaskGroup #

你当前用TaskGroup的写法没问题，但如果只想用单个任务处理，可以直接创建任务，在上下文退出时手动取消并等待，不需要TaskGroup：

@asynccontextmanager
async def open_connection(
    self, connection: HubControlConnection, building_id: str
) -> AsyncGenerator[HubControlConnection, None]:
    async with super().open_connection(connection, building_id) as conn:
        async def await_acquire_control() -> None:
            await conn.acquire_control()
            self.add_heating_system(building_id, CloudControllableHeatingSystem(conn))

        task = asyncio.create_task(await_acquire_control())
        try:
            yield conn
        finally:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass

这种写法更轻量化，直接管理单个任务的生命周期，避免了TaskGroup的额外层级。

2. 封装通用后台任务上下文管理器 #

既然你频繁使用这个模式，最好封装一个可复用的BackgroundTask上下文管理器，把任务启动、取消的逻辑抽离出来：

@asynccontextmanager
async def background_task(coro):
    task = asyncio.create_task(coro)
    try:
        yield
    finally:
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass

然后在你的代码中直接复用：

@asynccontextmanager
async def open_connection(
    self, connection: HubControlConnection, building_id: str
) -> AsyncGenerator[HubControlConnection, None]:
    async with super().open_connection(connection, building_id) as conn:
        async def await_acquire_control() -> None:
            await conn.acquire_control()
            self.add_heating_system(building_id, CloudControllableHeatingSystem(conn))

        async with background_task(await_acquire_control()):
            yield conn

这样代码更简洁，重复逻辑只需要维护一处，符合DRY原则。

3. 解耦任务与上下文管理器的生命周期 #

如果想完全解耦后台任务和上下文管理器，有两种可行思路：

• 连接对象内置任务管理：让HubControlConnection自身负责启动和停止acquire_control相关任务，比如给连接添加start_control()和stop_control()方法，上下文管理器只需要调用这些方法：

  @asynccontextmanager
  async def open_connection(
      self, connection: HubControlConnection, building_id: str
  ) -> AsyncGenerator[HubControlConnection, None]:
      async with super().open_connection(connection, building_id) as conn:
          await conn.start_control(building_id, self.add_heating_system)
          try:
              yield conn
          finally:
              await conn.stop_control()
  

  把任务管理逻辑内聚到连接对象中，让上下文管理器只专注于连接的打开/关闭职责。

• 用任务集合管理生命周期：维护一个实例级的任务集合（比如self._active_tasks = set()），当连接建立时添加任务，连接关闭时从集合中移除并取消。这种方式需要注意协程安全，适合多连接并发的复杂场景。

4. 细节优化：提升代码可读性 #

把内部的await_acquire_control函数改成类的辅助方法，避免在上下文管理器中嵌套定义函数，提升代码可读性：

class ControllableHeatingSystemStore(
    HeatingSystemStore[HubControlConnection, ControllableHeatingSystem]
):
    async def _acquire_control_and_register(self, conn: HubControlConnection, building_id: str) -> None:
        await conn.acquire_control()
        self.add_heating_system(building_id, CloudControllableHeatingSystem(conn))

    @asynccontextmanager
    async def open_connection(
        self, connection: HubControlConnection, building_id: str
    ) -> AsyncGenerator[HubControlConnection, None]:
        async with super().open_connection(connection, building_id) as conn:
            async with background_task(self._acquire_control_and_register(conn, building_id)):
                yield conn

内容的提问来源于stack exchange，提问作者Sam Pattuzzi