如何为Python上下文管理器添加类型提示注解？选Generator还是ContextManager？

如果你在用@contextlib.contextmanager写上下文管理器函数，别直接用typing.ContextManager当返回类型，正确的做法是用typing.Generator（或者Python 3.9+的collections.abc.Generator）来标注返回值类型。

为什么选Generator而不是ContextManager？ #

原因很简单：@contextlib.contextmanager的作用是把一个生成器函数包装成上下文管理器对象，但你的函数本身的返回值本质上是生成器。类型检查器需要识别原函数的返回类型（生成器），才能正确理解装饰器的转换逻辑；如果直接标注ContextManager，类型检查器会报错，因为原函数实际返回的是生成器，和标注的类型不匹配。

具体怎么写？ #

Generator有三个类型参数，对应生成器的三个特性：

• 第一个参数：yield语句产出的值的类型（如果你的生成器不产出值，就写None）
• 第二个参数：可以发送给生成器的值的类型（上下文管理器场景下几乎用不到，写None即可）
• 第三个参数：生成器最终返回的值的类型（上下文管理器里一般也不用，写None）

对应你给出的示例代码，正确的写法是：

import contextlib
from typing import Generator

@contextlib.contextmanager
def foo() -> Generator[None, None, None]:
    yield

如果你的上下文管理器需要产出某个值（比如临时文件路径），就修改第一个参数：

import contextlib
import os
from typing import Generator

@contextlib.contextmanager
def temp_file() -> Generator[str, None, None]:
    temp_path = "/tmp/temp_file.txt"
    try:
        # 前置准备：创建临时文件
        open(temp_path, "a").close()
        yield temp_path  # 产出临时文件路径给with语句使用
    finally:
        # 清理操作：删除临时文件
        if os.path.exists(temp_path):
            os.remove(temp_path)

版本兼容提示 #

• Python 3.9及以上：可以直接用标准库collections.abc里的Generator，不用再从typing导入
• Python 3.10+：还可以用typing.Generator的简写语法，不过和之前的写法效果一致

虽然contextlib官方文档对类型提示的提及不多，但typing.Generator的示例里明确覆盖了上下文管理器的场景，这也是官方推荐的标注方式。

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