mypy检查复合with表达式报object无__enter__属性问题咨询

问题背景 #

需要根据运行时条件，选择with语句管理的io.TextIOWrapper对象：既支持直接使用open()返回的文件对象，也支持使用自定义上下文管理器产出的io.TextIOWrapper对象，初始实现代码如下：

import contextlib
import io
import typing


def condition() -> bool:
    return True


@contextlib.contextmanager
def foo() -> typing.Generator[io.TextIOWrapper, None, None]:
    with open("README.md", "rb") as f:
        yield io.TextIOWrapper(f)


with foo() if condition() else open("README.md") as f:
    for line in f:
        print(line, end="")

报错现象 #

上述代码运行时无任何异常，但执行mypy类型检查时会抛出如下错误：

foo.py:16: error: "object" has no attribute "__enter__"
foo.py:16: error: "object" has no attribute "__exit__"

单独使用两个分支的任意一个表达式（例如单独编写with foo() as f:）时，mypy检查都可以正常通过，只有两个分支通过条件表达式组合时才会触发上述报错。

根因分析 #

问题出在foo函数的返回值类型标注错误：

• 被@contextlib.contextmanager装饰的生成器函数，调用后实际返回的是包装后的上下文管理器对象，而非原生生成器
• 原标注将foo的返回值标为typing.Generator[io.TextIOWrapper, None, None]，而open()返回值类型为io.TextIOWrapper，mypy会推断二者的公共超类为object，自然无法识别出对象实现了上下文管理器必需的__enter__、__exit__方法

之前使用typing.cast强制转换的方案属于临时规避，没有从根源解决类型标注错误的问题，临时方案代码如下：

with (typing.cast(io.TextIOWrapper, foo()) if condition() else
      open("README.md")) as f:

最优解决方案 #

直接修正foo函数的返回值类型标注，将返回值声明为typing.ContextManager[io.TextIOWrapper]（Python 3.9+也可使用contextlib.AbstractContextManager[io.TextIOWrapper]），明确告知类型检查器该函数返回的是包裹io.TextIOWrapper的上下文管理器对象。
修正后的foo函数定义如下：

@contextlib.contextmanager
def foo() -> typing.ContextManager[io.TextIOWrapper]:
    with open("README.md", "rb") as f:
        yield io.TextIOWrapper(f)

修正后条件表达式的两个分支会被mypy统一推断为ContextManager[io.TextIOWrapper]类型，既符合实际运行时的对象结构，也不会再触发类型检查错误，不需要额外添加类型强制转换。

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