mypy在函数定义内未应用类型窄化？bar函数类型检查报错求助

问题重现 #

以下代码中，foo函数能被mypy正确识别为类型合法，但bar函数会抛出类型错误：

from typing import Callable


def foo(a: str | int) -> list[str]:
    x: list[str] = ["abc", "def"]
    if isinstance(a, int):
        x.insert(a, "ghi")
    elif isinstance(a, str):
        x.insert(0, a)
    return x


def bar(a: str | int) -> Callable[[list[str]], list[str]]:
    if isinstance(a, int):
        def modify(x: list[str]) -> list[str]:
            x.insert(a, "ghi")
            return x
    elif isinstance(a, str):
        def modify(x: list[str]) -> list[str]:
            x.insert(0, a)
            return x
    return modify

mypy报错信息：

16: error: Argument 1 to "insert" of "list" has incompatible type "Union[str, int]"; expected "SupportsIndex"
20: error: Argument 2 to "insert" of "list" has incompatible type "Union[str, int]"; expected "str"

而Pyright对这段代码无类型报错，仅当错误混用参数位置时才会触发错误。

原因分析 #

这不是mypy的bug，而是它的类型推断规则限制：mypy不会自动将外部作用域中通过条件分支窄化的类型信息传递给嵌套函数。嵌套函数捕获外部变量时，mypy默认假设该变量可能在函数定义后被修改，因此保留了变量的原始联合类型（str | int），而非分支中窄化后的具体类型。

解决方案 #

方法1：使用typing.cast强制指定类型

在嵌套函数内部，通过cast明确变量的窄化类型，帮助mypy完成类型推断：

from typing import Callable, cast


def bar(a: str | int) -> Callable[[list[str]], list[str]]:
    if isinstance(a, int):
        def modify(x: list[str]) -> list[str]:
            a_int = cast(int, a)
            x.insert(a_int, "ghi")
            return x
    elif isinstance(a, str):
        def modify(x: list[str]) -> list[str]:
            a_str = cast(str, a)
            x.insert(0, a_str)
            return x
    return modify

方法2：将变量作为参数传递给嵌套函数

避免捕获外部变量，而是将窄化后的a作为参数传入嵌套函数，让mypy直接推断参数类型：

from typing import Callable


def bar(a: str | int) -> Callable[[list[str]], list[str]]:
    def modify_int(idx: int, x: list[str]) -> list[str]:
        x.insert(idx, "ghi")
        return x
    
    def modify_str(s: str, x: list[str]) -> list[str]:
        x.insert(0, s)
        return x
    
    if isinstance(a, int):
        return lambda x: modify_int(a, x)
    elif isinstance(a, str):
        return lambda x: modify_str(a, x)

方法3：使用闭包工厂函数封装分支逻辑

通过工厂函数生成对应类型的闭包，让mypy能明确每个分支返回的函数类型：

from typing import Callable


def create_modifier_int(idx: int) -> Callable[[list[str]], list[str]]:
    def modify(x: list[str]) -> list[str]:
        x.insert(idx, "ghi")
        return x
    return modify


def create_modifier_str(s: str) -> Callable[[list[str]], list[str]]:
    def modify(x: list[str]) -> list[str]:
        x.insert(0, s)
        return x
    return modify


def bar(a: str | int) -> Callable[[list[str]], list[str]]:
    if isinstance(a, int):
        return create_modifier_int(a)
    elif isinstance(a, str):
        return create_modifier_str(a)

补充说明 #

Pyright的类型推断引擎对闭包捕获变量的追踪能力更强，能识别出嵌套函数定义时外部变量的窄化状态，因此不会触发此类错误。

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