Source code for schrodinger.test.hypothesis.strategies.indices

from hypothesis import strategies


def indices(ordered_collection):
    """
    Strategy that returns an index within a non-empty collection

    :param ordered_collection: Any collection that supports len and indexing
    :type ordered_collection: collections.abc.Sequence

    :rtype: int
    :return: An index
    """
    length = len(ordered_collection)
    if length == 0:
        raise ValueError("Cannot draw an index from an empty list")
    return strategies.integers(min_value=0, max_value=length - 1)


def index_lists(ordered_collection):
    """
    Strategy that returns a subset of index_lists in any collection that supports
    indexing and len

    :param ordered_collection: Any object that supports len and indexing
    :type  ordered_collection: collections.abc.Sequence

    :rtype: list(int)
    :return: A subset of index_lists in the collection
    """
    list_length = len(ordered_collection)
    if list_length == 0:
        return strategies.lists(strategies.nothing())
    return strategies.lists(strategies.integers(min_value=0,
                                                max_value=list_length - 1),
                            unique=True).map(sorted)


@strategies.composite
def slice_pairs(draw, ordered_collection):
    """
    A strategy that returns a tuple of (start, end) indices in a collection.

    :param draw: A function supplied by the hypothesis decorator
    :type draw: function

    :param ordered_collection: Any object that supports len and indexing
    :type  ordered_collection: collections.abc.Sequence

    :rtype: tuple(int, int)
    :return: A tuple of (start, end) indices in collection
    """
    start = draw(indices(ordered_collection))
    end = draw(indices(ordered_collection[start:])) + start
    return (start, end)


def sublists(elements, **kwargs):
    """
    Strategy to draw a list containing a subset of elements from `elements`.

    :param kwargs: Extra kwargs to use with the list strategy.
    """
    return strategies.lists(elements=strategies.sampled_from(elements),
                            unique=True,
                            **kwargs)