Source code for hypothesis.strategies._internal.collections

# This file is part of Hypothesis, which may be found at
# https://github.com/HypothesisWorks/hypothesis/
#
# Copyright the Hypothesis Authors.
# Individual contributors are listed in AUTHORS.rst and the git log.
#
# This Source Code Form is subject to the terms of the Mozilla Public License,
# v. 2.0. If a copy of the MPL was not distributed with this file, You can
# obtain one at https://mozilla.org/MPL/2.0/.

import copy
from typing import Any, Iterable, Tuple, overload

from hypothesis.errors import InvalidArgument
from hypothesis.internal.conjecture import utils as cu
from hypothesis.internal.conjecture.junkdrawer import LazySequenceCopy
from hypothesis.internal.conjecture.utils import combine_labels
from hypothesis.internal.reflection import is_identity_function
from hypothesis.strategies._internal.strategies import (
    T3,
    T4,
    T5,
    Ex,
    MappedSearchStrategy,
    SearchStrategy,
    T,
    check_strategy,
    filter_not_satisfied,
)
from hypothesis.strategies._internal.utils import cacheable, defines_strategy


class TupleStrategy(SearchStrategy):
    """A strategy responsible for fixed length tuples based on heterogeneous
    strategies for each of their elements."""

    def __init__(self, strategies: Iterable[SearchStrategy[Any]]):
        super().__init__()
        self.element_strategies = tuple(strategies)

    def do_validate(self):
        for s in self.element_strategies:
            s.validate()

    def calc_label(self):
        return combine_labels(
            self.class_label, *(s.label for s in self.element_strategies)
        )

    def __repr__(self):
        tuple_string = ", ".join(map(repr, self.element_strategies))
        return f"TupleStrategy(({tuple_string}))"

    def calc_has_reusable_values(self, recur):
        return all(recur(e) for e in self.element_strategies)

    def do_draw(self, data):
        return tuple(data.draw(e) for e in self.element_strategies)

    def calc_is_empty(self, recur):
        return any(recur(e) for e in self.element_strategies)


@overload
def tuples() -> SearchStrategy[Tuple[()]]:  # pragma: no cover
    ...


@overload  # noqa: F811
def tuples(__a1: SearchStrategy[Ex]) -> SearchStrategy[Tuple[Ex]]:  # pragma: no cover
    ...


@overload  # noqa: F811
def tuples(
    __a1: SearchStrategy[Ex], __a2: SearchStrategy[T]
) -> SearchStrategy[Tuple[Ex, T]]:  # pragma: no cover
    ...


@overload  # noqa: F811
def tuples(
    __a1: SearchStrategy[Ex], __a2: SearchStrategy[T], __a3: SearchStrategy[T3]
) -> SearchStrategy[Tuple[Ex, T, T3]]:  # pragma: no cover
    ...


@overload  # noqa: F811
def tuples(
    __a1: SearchStrategy[Ex],
    __a2: SearchStrategy[T],
    __a3: SearchStrategy[T3],
    __a4: SearchStrategy[T4],
) -> SearchStrategy[Tuple[Ex, T, T3, T4]]:  # pragma: no cover
    ...


@overload  # noqa: F811
def tuples(
    __a1: SearchStrategy[Ex],
    __a2: SearchStrategy[T],
    __a3: SearchStrategy[T3],
    __a4: SearchStrategy[T4],
    __a5: SearchStrategy[T5],
) -> SearchStrategy[Tuple[Ex, T, T3, T4, T5]]:  # pragma: no cover
    ...


@overload  # noqa: F811
def tuples(
    *args: SearchStrategy[Any],
) -> SearchStrategy[Tuple[Any, ...]]:  # pragma: no cover
    ...


[docs]@cacheable @defines_strategy() def tuples(*args: SearchStrategy[Any]) -> SearchStrategy[Tuple[Any, ...]]: # noqa: F811 """Return a strategy which generates a tuple of the same length as args by generating the value at index i from args[i]. e.g. tuples(integers(), integers()) would generate a tuple of length two with both values an integer. Examples from this strategy shrink by shrinking their component parts. """ for arg in args: check_strategy(arg) return TupleStrategy(args)
class ListStrategy(SearchStrategy): """A strategy for lists which takes a strategy for its elements and the allowed lengths, and generates lists with the correct size and contents.""" _nonempty_filters: tuple = (bool, len, tuple, list) def __init__(self, elements, min_size=0, max_size=float("inf")): super().__init__() self.min_size = min_size or 0 self.max_size = max_size if max_size is not None else float("inf") assert 0 <= self.min_size <= self.max_size self.average_size = min( max(self.min_size * 2, self.min_size + 5), 0.5 * (self.min_size + self.max_size), ) self.element_strategy = elements def calc_label(self): return combine_labels(self.class_label, self.element_strategy.label) def do_validate(self): self.element_strategy.validate() if self.is_empty: raise InvalidArgument( "Cannot create non-empty lists with elements drawn from " f"strategy {self.element_strategy!r} because it has no values." ) if self.element_strategy.is_empty and 0 < self.max_size < float("inf"): raise InvalidArgument( f"Cannot create a collection of max_size={self.max_size!r}, " "because no elements can be drawn from the element strategy " f"{self.element_strategy!r}" ) def calc_is_empty(self, recur): if self.min_size == 0: return False else: return recur(self.element_strategy) def do_draw(self, data): if self.element_strategy.is_empty: assert self.min_size == 0 return [] elements = cu.many( data, min_size=self.min_size, max_size=self.max_size, average_size=self.average_size, ) result = [] while elements.more(): result.append(data.draw(self.element_strategy)) return result def __repr__(self): return "{}({!r}, min_size={!r}, max_size={!r})".format( self.__class__.__name__, self.element_strategy, self.min_size, self.max_size ) def filter(self, condition): if condition in self._nonempty_filters or is_identity_function(condition): assert self.max_size >= 1, "Always-empty is special cased in st.lists()" if self.min_size >= 1: return self new = copy.copy(self) new.min_size = 1 return new return super().filter(condition) class UniqueListStrategy(ListStrategy): def __init__(self, elements, min_size, max_size, keys, tuple_suffixes): super().__init__(elements, min_size, max_size) self.keys = keys self.tuple_suffixes = tuple_suffixes def do_draw(self, data): if self.element_strategy.is_empty: assert self.min_size == 0 return [] elements = cu.many( data, min_size=self.min_size, max_size=self.max_size, average_size=self.average_size, ) seen_sets = tuple(set() for _ in self.keys) result = [] # We construct a filtered strategy here rather than using a check-and-reject # approach because some strategies have special logic for generation under a # filter, and FilteredStrategy can consolidate multiple filters. def not_yet_in_unique_list(val): return all(key(val) not in seen for key, seen in zip(self.keys, seen_sets)) filtered = self.element_strategy._filter_for_filtered_draw( not_yet_in_unique_list ) while elements.more(): value = filtered.do_filtered_draw(data) if value is filter_not_satisfied: elements.reject() else: for key, seen in zip(self.keys, seen_sets): seen.add(key(value)) if self.tuple_suffixes is not None: value = (value,) + data.draw(self.tuple_suffixes) result.append(value) assert self.max_size >= len(result) >= self.min_size return result class UniqueSampledListStrategy(UniqueListStrategy): def do_draw(self, data): should_draw = cu.many( data, min_size=self.min_size, max_size=self.max_size, average_size=self.average_size, ) seen_sets = tuple(set() for _ in self.keys) result = [] remaining = LazySequenceCopy(self.element_strategy.elements) while remaining and should_draw.more(): i = len(remaining) - 1 j = cu.integer_range(data, 0, i) if j != i: remaining[i], remaining[j] = remaining[j], remaining[i] value = self.element_strategy._transform(remaining.pop()) if value is not filter_not_satisfied and all( key(value) not in seen for key, seen in zip(self.keys, seen_sets) ): for key, seen in zip(self.keys, seen_sets): seen.add(key(value)) if self.tuple_suffixes is not None: value = (value,) + data.draw(self.tuple_suffixes) result.append(value) else: should_draw.reject() assert self.max_size >= len(result) >= self.min_size return result class FixedKeysDictStrategy(MappedSearchStrategy): """A strategy which produces dicts with a fixed set of keys, given a strategy for each of their equivalent values. e.g. {'foo' : some_int_strategy} would generate dicts with the single key 'foo' mapping to some integer. """ def __init__(self, strategy_dict): self.dict_type = type(strategy_dict) self.keys = tuple(strategy_dict.keys()) super().__init__(strategy=TupleStrategy(strategy_dict[k] for k in self.keys)) def calc_is_empty(self, recur): return recur(self.mapped_strategy) def __repr__(self): return f"FixedKeysDictStrategy({self.keys!r}, {self.mapped_strategy!r})" def pack(self, value): return self.dict_type(zip(self.keys, value)) class FixedAndOptionalKeysDictStrategy(SearchStrategy): """A strategy which produces dicts with a fixed set of keys, given a strategy for each of their equivalent values. e.g. {'foo' : some_int_strategy} would generate dicts with the single key 'foo' mapping to some integer. """ def __init__(self, strategy_dict, optional): self.required = strategy_dict self.fixed = FixedKeysDictStrategy(strategy_dict) self.optional = optional def calc_is_empty(self, recur): return recur(self.fixed) def __repr__(self): return f"FixedAndOptionalKeysDictStrategy({self.required!r}, {self.optional!r})" def do_draw(self, data): result = data.draw(self.fixed) remaining = [k for k, v in self.optional.items() if not v.is_empty] should_draw = cu.many( data, min_size=0, max_size=len(remaining), average_size=len(remaining) / 2 ) while should_draw.more(): j = cu.integer_range(data, 0, len(remaining) - 1) remaining[-1], remaining[j] = remaining[j], remaining[-1] key = remaining.pop() result[key] = data.draw(self.optional[key]) return result