如何为Hypothesis组合策略添加正确的类型提示？

问题出在对@composite装饰器的类型逻辑理解偏差上：装饰器会把你编写的、返回具体值的函数，包装成一个返回SearchStrategy[T]的函数。所以你需要调整注解的位置和目标：

正确步骤 #

1. 导入必要的类型：从hypothesis.strategies导入DrawFn和SearchStrategy。
2. 注解composite函数本身：给被装饰的函数标注返回具体值的类型（比如int），同时给draw参数标注DrawFn类型。
3. 注解composite函数的调用结果：当你调用这个被装饰后的函数时，它的返回值才是SearchStrategy[T]，可以给这个变量添加类型注解。

示例代码 #

from hypothesis import composite
from hypothesis.strategies import SearchStrategy, integers, DrawFn

# 被@composite装饰的函数：返回具体的int，参数draw标注为DrawFn
@composite
def int_strategy(draw: DrawFn) -> int:
    # 通过draw获取基础策略生成的值
    base_num = draw(integers(min_value=0, max_value=100))
    # 返回具体的int值，装饰器会自动包装为SearchStrategy[int]
    return base_num * 2

# 调用装饰后的函数，结果类型是SearchStrategy[int]
my_int_strategy: SearchStrategy[int] = int_strategy()

为什么之前会报错？ #

你之前给int_strategy标注返回SearchStrategy[int]，但函数体实际返回的是int，mypy会检测到返回值类型不匹配。而@composite装饰器的作用就是将(DrawFn) -> T的函数，转换为() -> SearchStrategy[T]的函数，所以你只需要给原函数标注返回T即可。

更精确的类型（可选） #

如果需要更精确地标注draw参数的类型（比如明确它能绘制的策略类型），可以使用DrawFn[int]：

@composite
def int_strategy(draw: DrawFn[int]) -> int:
    base_num = draw(integers(min_value=0, max_value=100))
    return base_num * 2

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