如何为take2函数添加类型提示以通过mypy默认参数兼容性检查

初始示例代码 #

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


T = TypeVar('T')

def take1(n:int, iterable:Iterable[T]) -> Tuple[T,...]:  
    return tuple(islice(iterable, n))
    

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

mypy校验报错 #

test.py:12: error: Incompatible default for argument "container" (default has type "Type[Tuple[Any, ...]]", argument has type "Callable[[Iterable[Any]], T]")
Found 1 error in 1 file (checked 1 source file)

设计目标 #

take2是take1的通用实现版本：

• 默认行为与take1一致，返回截取后的元素元组
• 支持用户直接传入目标容器构造/处理函数，避免生成中间临时元组带来额外内存开销
  • 例如"".join(take1(3,"abcdefg"))与take2(3,"abcdefg","".join)效果等价，但后者不会生成take1产生的临时元组
  • 处理量级为10**10的超大迭代器场景下，take1会因构造超大元组触发内存错误，而take2直接传入sum等处理函数可正常运行
    要求不添加# type: ignore注释，正确编写take2的类型提示，既通过mypy检查，又能在调用help()函数时保留完整有效的类型提示信息。

使用@overload装饰器拆分两种调用场景的类型签名，规避默认参数的类型匹配冲突，具体实现代码如下：

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


T = TypeVar('T')

def take1(n:int, iterable:Iterable[T]) -> Tuple[T,...]:  
    return tuple(islice(iterable, n))
    
# 重载签名1：未传入container参数时，默认返回元素元组
@overload
def take2(n: int, iterable: Iterable[T]) -> Tuple[T, ...]:
    ...

# 重载签名2：传入container参数时，返回container处理后的结果类型
@overload
def take2(n: int, iterable: Iterable[Any], container: Callable[[Iterable[Any]], T]) -> T:
    ...

# 实际函数实现
def take2(n, iterable, container=tuple):
    return container(islice(iterable, n))

• 运行mypy校验无任何类型错误
• 调用help(take2)时会完整展示两个重载的调用签名，不会丢失类型提示信息
• 静态类型检查可自动识别不同调用方式的返回类型：
  • take2(3, "abc")会被识别为tuple[str, ...]类型
  • take2(3, "abc", "".join)会被识别为str类型
  • take2(3, [1,2,3,4], sum)会被识别为int类型
    完全匹配最初的设计预期。

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