Hypothesis：自定义策略

当内置策略不能完全满足需要时，可以编写自定义策略。本页介绍具体方法。

编写辅助函数

有时，编写辅助函数可以更简洁地表达项目中的常用模式。例如，response=json() 比把整个实现直接写在使用处更容易编写，也更容易阅读：

辅助函数片段假定已导入 from hypothesis import strategies as st。这里的 json() 是生成 Python 数据结构的策略辅助函数，并非标准库的 JSON 序列化函数；当 finite_only=False 时允许的无穷值和 NaN 也不属于标准 JSON 数字。

def json(*, finite_only=True):
    """Helper function to describe JSON objects, with optional inf and nan."""
    numbers = st.floats(allow_infinity=not finite_only, allow_nan=not finite_only)
    return st.recursive(
        st.none() | st.booleans() | st.integers() | numbers | st.text(),
        extend=lambda xs: st.lists(xs) | st.dictionaries(st.text(), xs),
    )

编写自己的策略

如果 Hypothesis 中的策略不符合需要，可以自己编写。

例如，假设需要生成一个元素总和为 1 的浮点数列表。可以先用 lists(floats(0, 1)) 生成取值在 0 到 1 之间的浮点数列表。但做到这一步后，单靠现成策略就难以继续表达“总和为 1”的要求。

定义新策略的一种方式是使用 @composite 装饰器。它允许通过任意 Python 代码定义策略。上面的需求可以实现为：

from hypothesis import strategies as st

@st.composite
def sums_to_one(draw):
    l = draw(st.lists(st.floats(0, 1)))
    return [f / sum(l) for f in l]

这个初步版本还不完整：默认列表策略可能生成空列表，返回值总和便是 0；非空但全为 0 的列表会触发除零。下一个原文示例用 min_size=1 和正下界 0.001 避免这两种情况，并用近似比较处理浮点误差。

@composite 会把一个 draw 函数作为被装饰函数的第一个参数。draw 用于从其他策略中抽取随机值。sums_to_one 返回的就是所需类型的值；这里是总和为 1 的列表。

看看这个新策略如何使用：

import pytest

from hypothesis import given, strategies as st

@st.composite
def sums_to_one(draw):
    lst = draw(st.lists(st.floats(0.001, 1), min_size=1))
    return [f / sum(lst) for f in lst]

@given(sums_to_one())
def test(lst):
    # ignore floating point errors
    assert sum(lst) == pytest.approx(1)

为 @composite 添加参数

被 @composite 装饰的函数可以添加参数，包括仅限关键字的参数。它们的行为与普通函数参数相同。

例如，要把 sums_to_one 推广为 sums_to_n，可以添加参数 n：

import pytest

from hypothesis import assume, given, strategies as st

@st.composite
def sums_to_n(draw, n=1):  #  <-- changed
    lst = draw(st.lists(st.floats(0, 1), min_size=1))
    assume(sum(lst) > 0)
    return [f / sum(lst) * n for f in lst]  #  <-- changed

@given(sums_to_n(10))
def test(lst):
    assert sum(lst) == pytest.approx(10)

也可以同样简单地把 n 设为仅限关键字的参数：

import pytest

from hypothesis import assume, given, strategies as st

@st.composite
def sums_to_n(draw, *, n=1):  #  <-- changed
    lst = draw(st.lists(st.floats(0, 1), min_size=1))
    assume(sum(lst) > 0)
    return [f / sum(lst) * n for f in lst]

@given(sums_to_n(n=10))  #  <-- changed
def test(lst):
    assert sum(lst) == pytest.approx(10)

使用 @composite 生成相互依赖的值

当一个值的生成依赖另一个策略生成的值时，@composite 也很有用。例如，要生成满足 n1 <= n2 的两个整数 n1、n2：

这个片段继续使用前文导入的 st 与 given。第二次抽样把第一次得到的 n1 作为下界，因此依赖关系由策略本身保证。

@st.composite
def ordered_pairs(draw):
    n1 = draw(st.integers())
    n2 = draw(st.integers(min_value=n1))
    return (n1, n2)

@given(ordered_pairs())
def test_pairs_are_ordered(pair):
    n1, n2 = pair
    assert n1 <= n2

混合数据生成与测试代码

使用 @composite 时，必须先生成完整输入，再运行测试。但有时，希望先确认前面的断言通过，才生成后续输入；也可能因为控制流程复杂，必须在测试执行中途生成数据。

data() 可以实现这一点。它与 @composite 类似，但允许把测试代码与生成代码混在一起。

下面用 data() 重写前面的 @composite 示例：

import pytest

from hypothesis import given, strategies as st

@given(st.data())
def test(data):
    lst = data.draw(st.lists(st.floats(0.001, 1), min_size=1))
    lst = [f / sum(lst) for f in lst]
    # ignore floating point errors
    assert sum(lst) == pytest.approx(1)