1# This file is part of Hypothesis, which may be found at
2# https://github.com/HypothesisWorks/hypothesis/
3#
4# Copyright the Hypothesis Authors.
5# Individual contributors are listed in AUTHORS.rst and the git log.
6#
7# This Source Code Form is subject to the terms of the Mozilla Public License,
8# v. 2.0. If a copy of the MPL was not distributed with this file, You can
9# obtain one at https://mozilla.org/MPL/2.0/.
10
11import copy
12from typing import Any, Iterable, Tuple, overload
13
14from hypothesis.errors import InvalidArgument
15from hypothesis.internal.conjecture import utils as cu
16from hypothesis.internal.conjecture.engine import BUFFER_SIZE
17from hypothesis.internal.conjecture.junkdrawer import LazySequenceCopy
18from hypothesis.internal.conjecture.utils import combine_labels
19from hypothesis.internal.filtering import get_integer_predicate_bounds
20from hypothesis.internal.reflection import is_identity_function
21from hypothesis.strategies._internal.strategies import (
22 T3,
23 T4,
24 T5,
25 Ex,
26 MappedStrategy,
27 SearchStrategy,
28 T,
29 check_strategy,
30 filter_not_satisfied,
31)
32from hypothesis.strategies._internal.utils import cacheable, defines_strategy
33
34
35class TupleStrategy(SearchStrategy):
36 """A strategy responsible for fixed length tuples based on heterogeneous
37 strategies for each of their elements."""
38
39 def __init__(self, strategies: Iterable[SearchStrategy[Any]]):
40 super().__init__()
41 self.element_strategies = tuple(strategies)
42
43 def do_validate(self):
44 for s in self.element_strategies:
45 s.validate()
46
47 def calc_label(self):
48 return combine_labels(
49 self.class_label, *(s.label for s in self.element_strategies)
50 )
51
52 def __repr__(self):
53 tuple_string = ", ".join(map(repr, self.element_strategies))
54 return f"TupleStrategy(({tuple_string}))"
55
56 def calc_has_reusable_values(self, recur):
57 return all(recur(e) for e in self.element_strategies)
58
59 def do_draw(self, data):
60 return tuple(data.draw(e) for e in self.element_strategies)
61
62 def calc_is_empty(self, recur):
63 return any(recur(e) for e in self.element_strategies)
64
65
66@overload
67def tuples() -> SearchStrategy[Tuple[()]]: # pragma: no cover
68 ...
69
70
71@overload
72def tuples(__a1: SearchStrategy[Ex]) -> SearchStrategy[Tuple[Ex]]: # pragma: no cover
73 ...
74
75
76@overload
77def tuples(
78 __a1: SearchStrategy[Ex], __a2: SearchStrategy[T]
79) -> SearchStrategy[Tuple[Ex, T]]: # pragma: no cover
80 ...
81
82
83@overload
84def tuples(
85 __a1: SearchStrategy[Ex], __a2: SearchStrategy[T], __a3: SearchStrategy[T3]
86) -> SearchStrategy[Tuple[Ex, T, T3]]: # pragma: no cover
87 ...
88
89
90@overload
91def tuples(
92 __a1: SearchStrategy[Ex],
93 __a2: SearchStrategy[T],
94 __a3: SearchStrategy[T3],
95 __a4: SearchStrategy[T4],
96) -> SearchStrategy[Tuple[Ex, T, T3, T4]]: # pragma: no cover
97 ...
98
99
100@overload
101def tuples(
102 __a1: SearchStrategy[Ex],
103 __a2: SearchStrategy[T],
104 __a3: SearchStrategy[T3],
105 __a4: SearchStrategy[T4],
106 __a5: SearchStrategy[T5],
107) -> SearchStrategy[Tuple[Ex, T, T3, T4, T5]]: # pragma: no cover
108 ...
109
110
111@overload
112def tuples(
113 *args: SearchStrategy[Any],
114) -> SearchStrategy[Tuple[Any, ...]]: # pragma: no cover
115 ...
116
117
118@cacheable
119@defines_strategy()
120def tuples(*args: SearchStrategy[Any]) -> SearchStrategy[Tuple[Any, ...]]:
121 """Return a strategy which generates a tuple of the same length as args by
122 generating the value at index i from args[i].
123
124 e.g. tuples(integers(), integers()) would generate a tuple of length
125 two with both values an integer.
126
127 Examples from this strategy shrink by shrinking their component parts.
128 """
129 for arg in args:
130 check_strategy(arg)
131
132 return TupleStrategy(args)
133
134
135class ListStrategy(SearchStrategy):
136 """A strategy for lists which takes a strategy for its elements and the
137 allowed lengths, and generates lists with the correct size and contents."""
138
139 _nonempty_filters: tuple = (bool, len, tuple, list)
140
141 def __init__(self, elements, min_size=0, max_size=float("inf")):
142 super().__init__()
143 self.min_size = min_size or 0
144 self.max_size = max_size if max_size is not None else float("inf")
145 assert 0 <= self.min_size <= self.max_size
146 self.average_size = min(
147 max(self.min_size * 2, self.min_size + 5),
148 0.5 * (self.min_size + self.max_size),
149 )
150 self.element_strategy = elements
151 if min_size > BUFFER_SIZE:
152 raise InvalidArgument(
153 f"{self!r} can never generate an example, because min_size is larger "
154 "than Hypothesis supports. Including it is at best slowing down your "
155 "tests for no benefit; at worst making them fail (maybe flakily) with "
156 "a HealthCheck error."
157 )
158
159 def calc_label(self):
160 return combine_labels(self.class_label, self.element_strategy.label)
161
162 def do_validate(self):
163 self.element_strategy.validate()
164 if self.is_empty:
165 raise InvalidArgument(
166 "Cannot create non-empty lists with elements drawn from "
167 f"strategy {self.element_strategy!r} because it has no values."
168 )
169 if self.element_strategy.is_empty and 0 < self.max_size < float("inf"):
170 raise InvalidArgument(
171 f"Cannot create a collection of max_size={self.max_size!r}, "
172 "because no elements can be drawn from the element strategy "
173 f"{self.element_strategy!r}"
174 )
175
176 def calc_is_empty(self, recur):
177 if self.min_size == 0:
178 return False
179 else:
180 return recur(self.element_strategy)
181
182 def do_draw(self, data):
183 if self.element_strategy.is_empty:
184 assert self.min_size == 0
185 return []
186
187 elements = cu.many(
188 data,
189 min_size=self.min_size,
190 max_size=self.max_size,
191 average_size=self.average_size,
192 )
193 result = []
194 while elements.more():
195 result.append(data.draw(self.element_strategy))
196 return result
197
198 def __repr__(self):
199 return (
200 f"{self.__class__.__name__}({self.element_strategy!r}, "
201 f"min_size={self.min_size:_}, max_size={self.max_size:_})"
202 )
203
204 def filter(self, condition):
205 if condition in self._nonempty_filters or is_identity_function(condition):
206 assert self.max_size >= 1, "Always-empty is special cased in st.lists()"
207 if self.min_size >= 1:
208 return self
209 new = copy.copy(self)
210 new.min_size = 1
211 return new
212
213 kwargs, pred = get_integer_predicate_bounds(condition)
214 if kwargs.get("len") and ("min_value" in kwargs or "max_value" in kwargs):
215 new = copy.copy(self)
216 new.min_size = max(self.min_size, kwargs.get("min_value", self.min_size))
217 new.max_size = min(self.max_size, kwargs.get("max_value", self.max_size))
218 # Unsatisfiable filters are easiest to understand without rewriting.
219 if new.min_size > new.max_size:
220 return SearchStrategy.filter(self, condition)
221 # Recompute average size; this is cheaper than making it into a property.
222 new.average_size = min(
223 max(new.min_size * 2, new.min_size + 5),
224 0.5 * (new.min_size + new.max_size),
225 )
226 if pred is None:
227 return new
228 return SearchStrategy.filter(new, condition)
229
230 return SearchStrategy.filter(self, condition)
231
232
233class UniqueListStrategy(ListStrategy):
234 def __init__(self, elements, min_size, max_size, keys, tuple_suffixes):
235 super().__init__(elements, min_size, max_size)
236 self.keys = keys
237 self.tuple_suffixes = tuple_suffixes
238
239 def do_draw(self, data):
240 if self.element_strategy.is_empty:
241 assert self.min_size == 0
242 return []
243
244 elements = cu.many(
245 data,
246 min_size=self.min_size,
247 max_size=self.max_size,
248 average_size=self.average_size,
249 )
250 seen_sets = tuple(set() for _ in self.keys)
251 result = []
252
253 # We construct a filtered strategy here rather than using a check-and-reject
254 # approach because some strategies have special logic for generation under a
255 # filter, and FilteredStrategy can consolidate multiple filters.
256 def not_yet_in_unique_list(val):
257 return all(key(val) not in seen for key, seen in zip(self.keys, seen_sets))
258
259 filtered = self.element_strategy._filter_for_filtered_draw(
260 not_yet_in_unique_list
261 )
262 while elements.more():
263 value = filtered.do_filtered_draw(data)
264 if value is filter_not_satisfied:
265 elements.reject(f"Aborted test because unable to satisfy {filtered!r}")
266 else:
267 for key, seen in zip(self.keys, seen_sets):
268 seen.add(key(value))
269 if self.tuple_suffixes is not None:
270 value = (value, *data.draw(self.tuple_suffixes))
271 result.append(value)
272 assert self.max_size >= len(result) >= self.min_size
273 return result
274
275
276class UniqueSampledListStrategy(UniqueListStrategy):
277 def do_draw(self, data):
278 should_draw = cu.many(
279 data,
280 min_size=self.min_size,
281 max_size=self.max_size,
282 average_size=self.average_size,
283 )
284 seen_sets = tuple(set() for _ in self.keys)
285 result = []
286
287 remaining = LazySequenceCopy(self.element_strategy.elements)
288
289 while remaining and should_draw.more():
290 j = data.draw_integer(0, len(remaining) - 1)
291 value = self.element_strategy._transform(remaining.pop(j))
292 if value is not filter_not_satisfied and all(
293 key(value) not in seen for key, seen in zip(self.keys, seen_sets)
294 ):
295 for key, seen in zip(self.keys, seen_sets):
296 seen.add(key(value))
297 if self.tuple_suffixes is not None:
298 value = (value, *data.draw(self.tuple_suffixes))
299 result.append(value)
300 else:
301 should_draw.reject(
302 "UniqueSampledListStrategy filter not satisfied or value already seen"
303 )
304 assert self.max_size >= len(result) >= self.min_size
305 return result
306
307
308class FixedKeysDictStrategy(MappedStrategy):
309 """A strategy which produces dicts with a fixed set of keys, given a
310 strategy for each of their equivalent values.
311
312 e.g. {'foo' : some_int_strategy} would generate dicts with the single
313 key 'foo' mapping to some integer.
314 """
315
316 def __init__(self, strategy_dict):
317 dict_type = type(strategy_dict)
318 self.keys = tuple(strategy_dict.keys())
319 super().__init__(
320 strategy=TupleStrategy(strategy_dict[k] for k in self.keys),
321 pack=lambda value: dict_type(zip(self.keys, value)),
322 )
323
324 def calc_is_empty(self, recur):
325 return recur(self.mapped_strategy)
326
327 def __repr__(self):
328 return f"FixedKeysDictStrategy({self.keys!r}, {self.mapped_strategy!r})"
329
330
331class FixedAndOptionalKeysDictStrategy(SearchStrategy):
332 """A strategy which produces dicts with a fixed set of keys, given a
333 strategy for each of their equivalent values.
334
335 e.g. {'foo' : some_int_strategy} would generate dicts with the single
336 key 'foo' mapping to some integer.
337 """
338
339 def __init__(self, strategy_dict, optional):
340 self.required = strategy_dict
341 self.fixed = FixedKeysDictStrategy(strategy_dict)
342 self.optional = optional
343
344 def calc_is_empty(self, recur):
345 return recur(self.fixed)
346
347 def __repr__(self):
348 return f"FixedAndOptionalKeysDictStrategy({self.required!r}, {self.optional!r})"
349
350 def do_draw(self, data):
351 result = data.draw(self.fixed)
352 remaining = [k for k, v in self.optional.items() if not v.is_empty]
353 should_draw = cu.many(
354 data, min_size=0, max_size=len(remaining), average_size=len(remaining) / 2
355 )
356 while should_draw.more():
357 j = data.draw_integer(0, len(remaining) - 1)
358 remaining[-1], remaining[j] = remaining[j], remaining[-1]
359 key = remaining.pop()
360 result[key] = data.draw(self.optional[key])
361 return result