如何用mypy为Python上下文管理器结合字符串字面量实现重载类型标注？

针对你遇到的mypy类型检查错误，核心问题是给捕获流标注了过于宽泛的Iterable[Iterable[str]]类型，而实际操作的是具备getvalue()和close()方法的io.StringIO实例。以下是适配Python3.9+的完整类型标注方案，包含函数重载以处理不同捕获场景：

完整实现代码 #

import sys
import io
from typing import overload, ContextManager, Tuple, Literal

@overload
def capture_output(
    capture_stdout: Literal[True] = True,
    capture_stderr: Literal[True] = True
) -> ContextManager[Tuple[io.StringIO, io.StringIO]]: ...

@overload
def capture_output(
    capture_stdout: Literal[True],
    capture_stderr: Literal[False]
) -> ContextManager[io.StringIO]: ...

@overload
def capture_output(
    capture_stdout: Literal[False],
    capture_stderr: Literal[True]
) -> ContextManager[io.StringIO]: ...

@overload
def capture_output(
    capture_stdout: Literal[False],
    capture_stderr: Literal[False]
) -> ContextManager[None]: ...

def capture_output(
    capture_stdout: bool = True,
    capture_stderr: bool = True
):
    original_stdout = sys.stdout
    original_stderr = sys.stderr
    stdout_capture = io.StringIO() if capture_stdout else None
    stderr_capture = io.StringIO() if capture_stderr else None

    try:
        if capture_stdout:
            sys.stdout = stdout_capture
        if capture_stderr:
            sys.stderr = stderr_capture
        # 根据捕获情况返回对应对象
        if capture_stdout and capture_stderr:
            yield stdout_capture, stderr_capture
        elif capture_stdout:
            yield stdout_capture
        elif capture_stderr:
            yield stderr_capture
        else:
            yield None
    finally:
        sys.stdout = original_stdout
        sys.stderr = original_stderr
        if stdout_capture:
            stdout_capture.close()
        if stderr_capture:
            stderr_capture.close()

关键说明 #

1. 函数重载：通过@overload定义四种场景的返回类型，让mypy能根据传入的capture_stdout和capture_stderr参数值，准确推断上下文管理器的返回值类型（元组、单个StringIO或None）。
2. 精确类型标注：使用io.StringIO而非宽泛的Iterable类型，明确告知mypy捕获流具备getvalue()和close()方法。
3. Python3.9+适配：用Literal类型限定参数的布尔字面量，让重载逻辑被mypy正确识别；若需兼容更早版本，可改用Union结合注释，但3.9+用Literal更清晰。

用法示例 #

def test_print():
    with capture_output() as (out, err):
        print("Hello stdout")
        print("Hello stderr", file=sys.stderr)
    assert out.getvalue().strip() == "Hello stdout"
    assert err.getvalue().strip() == "Hello stderr"

def test_only_stdout():
    with capture_output(capture_stderr=False) as out:
        print("Only stdout")
    assert out.getvalue().strip() == "Only stdout"

这样mypy就能正确识别out和err的类型，不会再提示getvalue()或close()不存在的错误。

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