如何让mypy正确推断自定义mytake函数的返回值类型

实现的mytake函数可从可迭代对象中截取前n个元素，支持传入自定义处理函数/容器类型对截取结果做转换，运行时输出完全符合预期，但使用mypy做静态类型检查时存在类型推断不准的问题：

• 未传入container参数使用默认tuple时，可正确推断返回类型为tuple[int, ...]
• 传入list、sum等自定义container参数时，返回值均被推断为Any，无法得到预期的list[int]、int类型
  需要在不给接收返回值的变量手动添加显式类型注解的前提下，修正类型标注让mypy得到正确的推断结果。

原始示例代码 #

from typing import Callable, Iterable, TypeVar, Any, overload, Tuple, TYPE_CHECKING
from itertools import islice

T = TypeVar("T")

@overload
def mytake(n:int, iterable:Iterable[T]) -> Tuple[T,...]:...
@overload
def mytake(n:int, iterable:Iterable[T], container:Callable[[Iterable[T]],Any]) -> Any:...

def mytake(n:int, iterable:Iterable[T], container:Callable[[Iterable[T]],Any]=tuple) -> Any:
    return container(islice(iterable,n))

if TYPE_CHECKING:
    pass
else:
    reveal_type = print

#a:tuple[int,...]
a = mytake(10,range(100))
reveal_type(a)

#b:list[int]
b = mytake(10,range(100), list)
reveal_type(b)#desired to be list[int]

#c:int
c = mytake(10,range(100),sum)
reveal_type(c)#desired to be int


help(mytake)

运行时输出 #

(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
45
Help on function mytake in module __main__:

mytake(n: int, iterable: Iterable[~T], container: Callable[[Iterable[~T]], Any] = <class 'tuple'>) -> Any

mypy检查原始代码的输出 #

C:\Users\copperfield\Desktop\mypy_tests>mypy test6.py
test6.py:21: note: Revealed type is "builtins.tuple[builtins.int, ...]"
test6.py:25: note: Revealed type is "Any"
test6.py:29: note: Revealed type is "Any"
Success: no issues found in 1 source file

C:\Users\copperfield\Desktop\mypy_tests>

问题根源是第二个重载的类型标注没有把container参数的返回类型和函数本身的返回类型做绑定，直接将二者都写为Any，类型检查器自然无法追踪传入的container对应的返回类型。
修正步骤：

• 新增一个类型变量R，用于绑定传入的container参数的返回值类型
• 修改第二个重载签名：将container参数标注为Callable[[Iterable[T]], R]，同时将该重载的返回值标注为R
• 实际运行的函数实现部分不需要改动，重载层的签名是mypy做类型推断的唯一参考

修正后的完整类型标注代码如下：

from typing import Callable, Iterable, TypeVar, Any, overload, Tuple, TYPE_CHECKING
from itertools import islice

T = TypeVar("T")
R = TypeVar("R") # 新增：绑定自定义container的返回类型

@overload
def mytake(n:int, iterable:Iterable[T]) -> Tuple[T,...]:...
@overload
def mytake(n:int, iterable:Iterable[T], container:Callable[[Iterable[T]],R]) -> R:...

def mytake(n:int, iterable:Iterable[T], container:Callable[[Iterable[T]],Any]=tuple) -> Any:
    return container(islice(iterable,n))

if TYPE_CHECKING:
    pass
else:
    reveal_type = print

a = mytake(10,range(100))
reveal_type(a)

b = mytake(10,range(100), list)
reveal_type(b)

c = mytake(10,range(100),sum)
reveal_type(c)

修正后重新运行mypy检查，得到的结果完全符合预期：

test6.py:22: note: Revealed type is "builtins.tuple[builtins.int, ...]"
test6.py:26: note: Revealed type is "builtins.list[builtins.int]"
test6.py:30: note: Revealed type is "builtins.int"
Success: no issues found in 1 source file

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