WIP wagtail block strategies for hypothesis

blocks.py

from functools import singledispatch, wraps
from typing import Callable, TypeVar

from hypothesis import assume, event
from hypothesis import strategies as st
from hypothesis.extra import django as dj
from wagtail import blocks

B = TypeVar("B", bound=blocks.Block)


class UnknownBlockType(Exception):
    pass


def with_block_default_strategy(
    strategy_generator: Callable[[B], st.SearchStrategy],
) -> Callable[[B], st.SearchStrategy]:

    return strategy_generator

    # @wraps(strategy_generator)
    # def wrapper(block: B, *args, **kwargs):
    #     default_strat = st.just(block.get_default())

    #     return default_strat | strategy_generator(block, *args, **kwargs)

    # return wrapper


@singledispatch
def from_block(block, *args, **kwargs) -> st.SearchStrategy:
    raise UnknownBlockType(f"No strategy registered for arg of type {type(block)}")


@from_block.register
@with_block_default_strategy
def _(block: blocks.FieldBlock):
    """
    Generate a strategy for a FieldBlock.
    """
    return dj.from_field(block.field)


@from_block.register
@with_block_default_strategy
def _(block: blocks.ListBlock):
    """
    Generate a strategy for a ListBlock.
    """
    return st.lists(
        from_block(block.child_block),
        min_size=block.meta.min_num or 0,
        max_size=block.meta.max_num,
    ).map(lambda xs: blocks.list_block.ListValue(block, xs))


@from_block.register
@with_block_default_strategy
def _(block: blocks.StructBlock):
    """
    Generate a strategy for a StructBlock.
    """

    strat_map = {k: from_block(v) for k, v in block.child_blocks.items()}
    return st.fixed_dictionaries(strat_map).map(
        lambda d: blocks.struct_block.StructValue(block, d)
    )


NOT_SET = object()


def get_in(d: dict, key, *keys, default=NOT_SET):
    result = d
    for k in [key, *keys]:
        try:
            result = result[k]
        except KeyError:
            if default is NOT_SET:
                raise
            return default
    return result


@st.composite
def stream_values(draw, block: blocks.StreamBlock):
    """
    Generating valid values efficiently is a little tricky when considering arbitrary
    block count declarations (on the stream itself and per child block). We might be
    able to do better here (i.e. generate less rejected values) by generating values for
    blocks with strictest constraints first.
    """

    child_strategies = {name: from_block(b) for name, b in block.child_blocks.items()}
    block_counts = block.meta.block_counts

    # Permute the order so as not to bias early declarations
    order = draw(st.permutations(list(child_strategies.keys())))

    generated = [
        (name, x)
        for name in order
        for x in draw(
            st.lists(
                child_strategies[name],
                min_size=get_in(block_counts, name, "min_num", default=0),
                max_size=get_in(block_counts, name, "max_num", default=None),
            )
        )
    ]

    max_n = block.meta.max_num
    min_n = block.meta.min_num
    event(
        f"Generated stream of length {len(generated)} for {block} "
        f"having min_num={min_n} and max_num={max_n}"
    )

    if max_n is not None:
        assume(len(generated) <= max_n)
    if min_n is not None:
        assume(len(generated) >= min_n)

    return generated


@from_block.register
@with_block_default_strategy
def _(block: blocks.StreamBlock):
    """
    Generate a strategy for a StreamBlock.
    """
    return stream_values(block)