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
12import math
13from collections.abc import Callable, Iterable
14from typing import Any, TypeGuard, overload
15
16from hypothesis import strategies as st
17from hypothesis.control import current_build_context
18from hypothesis.errors import InvalidArgument
19from hypothesis.internal.conjecture import utils as cu
20from hypothesis.internal.conjecture.data import ConjectureData
21from hypothesis.internal.conjecture.engine import BUFFER_SIZE
22from hypothesis.internal.conjecture.junkdrawer import LazySequenceCopy
23from hypothesis.internal.conjecture.utils import combine_labels
24from hypothesis.internal.filtering import get_integer_predicate_bounds
25from hypothesis.internal.reflection import is_identity_function
26from hypothesis.strategies._internal.strategies import (
27 T3,
28 T4,
29 T5,
30 Ex,
31 FilteredStrategy,
32 RecurT,
33 SampledFromStrategy,
34 SearchStrategy,
35 T,
36 check_strategy,
37 filter_not_satisfied,
38)
39from hypothesis.strategies._internal.utils import cacheable, defines_strategy
40from hypothesis.utils.conventions import UniqueIdentifier
41from hypothesis.vendor.pretty import (
42 ArgLabelsT,
43 IDKey,
44 _fixeddict_pprinter,
45 _tuple_pprinter,
46)
47
48
49class TupleStrategy(SearchStrategy[tuple[Ex, ...]]):
50 """A strategy responsible for fixed length tuples based on heterogeneous
51 strategies for each of their elements."""
52
53 def __init__(self, strategies: Iterable[SearchStrategy[Any]]):
54 super().__init__()
55 self.element_strategies = tuple(strategies)
56
57 def do_validate(self) -> None:
58 for s in self.element_strategies:
59 s.validate()
60
61 def calc_label(self) -> int:
62 return combine_labels(
63 self.class_label, *(s.label for s in self.element_strategies)
64 )
65
66 def __repr__(self) -> str:
67 tuple_string = ", ".join(map(repr, self.element_strategies))
68 return f"TupleStrategy(({tuple_string}))"
69
70 def calc_has_reusable_values(self, recur: RecurT) -> bool:
71 return all(recur(e) for e in self.element_strategies)
72
73 def do_draw(self, data: ConjectureData) -> tuple[Ex, ...]:
74 context = current_build_context()
75 arg_labels: ArgLabelsT = {}
76 result = []
77 for i, strategy in enumerate(self.element_strategies):
78 with context.track_arg_label(f"arg[{i}]") as arg_label:
79 result.append(data.draw(strategy))
80 arg_labels |= arg_label
81
82 result = tuple(result)
83 if arg_labels:
84 context.known_object_printers[IDKey(result)].append(
85 _tuple_pprinter(arg_labels)
86 )
87 return result
88
89 def calc_is_empty(self, recur: RecurT) -> bool:
90 return any(recur(e) for e in self.element_strategies)
91
92
93@overload
94def tuples() -> SearchStrategy[tuple[()]]: # pragma: no cover
95 ...
96
97
98@overload
99def tuples(__a1: SearchStrategy[Ex]) -> SearchStrategy[tuple[Ex]]: # pragma: no cover
100 ...
101
102
103@overload
104def tuples(
105 __a1: SearchStrategy[Ex], __a2: SearchStrategy[T]
106) -> SearchStrategy[tuple[Ex, T]]: # pragma: no cover
107 ...
108
109
110@overload
111def tuples(
112 __a1: SearchStrategy[Ex], __a2: SearchStrategy[T], __a3: SearchStrategy[T3]
113) -> SearchStrategy[tuple[Ex, T, T3]]: # pragma: no cover
114 ...
115
116
117@overload
118def tuples(
119 __a1: SearchStrategy[Ex],
120 __a2: SearchStrategy[T],
121 __a3: SearchStrategy[T3],
122 __a4: SearchStrategy[T4],
123) -> SearchStrategy[tuple[Ex, T, T3, T4]]: # pragma: no cover
124 ...
125
126
127@overload
128def tuples(
129 __a1: SearchStrategy[Ex],
130 __a2: SearchStrategy[T],
131 __a3: SearchStrategy[T3],
132 __a4: SearchStrategy[T4],
133 __a5: SearchStrategy[T5],
134) -> SearchStrategy[tuple[Ex, T, T3, T4, T5]]: # pragma: no cover
135 ...
136
137
138@overload
139def tuples(
140 *args: SearchStrategy[Any],
141) -> SearchStrategy[tuple[Any, ...]]: # pragma: no cover
142 ...
143
144
145@cacheable
146@defines_strategy()
147def tuples(*args: SearchStrategy[Any]) -> SearchStrategy[tuple[Any, ...]]:
148 """Return a strategy which generates a tuple of the same length as args by
149 generating the value at index i from args[i].
150
151 e.g. tuples(integers(), integers()) would generate a tuple of length
152 two with both values an integer.
153
154 Examples from this strategy shrink by shrinking their component parts.
155 """
156 for arg in args:
157 check_strategy(arg)
158
159 return TupleStrategy(args)
160
161
162class ListStrategy(SearchStrategy[list[Ex]]):
163 """A strategy for lists which takes a strategy for its elements and the
164 allowed lengths, and generates lists with the correct size and contents."""
165
166 _nonempty_filters: tuple[Callable[[Any], Any], ...] = (bool, len, tuple, list)
167
168 def __init__(
169 self,
170 elements: SearchStrategy[Ex],
171 min_size: int = 0,
172 max_size: float | int | None = math.inf,
173 ):
174 super().__init__()
175 self.min_size = min_size or 0
176 self.max_size = max_size if max_size is not None else math.inf
177 assert 0 <= self.min_size <= self.max_size
178 self.average_size = min(
179 max(self.min_size * 2, self.min_size + 5),
180 0.5 * (self.min_size + self.max_size),
181 )
182 self.element_strategy = elements
183 if min_size > BUFFER_SIZE:
184 raise InvalidArgument(
185 f"{self!r} can never generate an example, because min_size is larger "
186 "than Hypothesis supports. Including it is at best slowing down your "
187 "tests for no benefit; at worst making them fail (maybe flakily) with "
188 "a HealthCheck error."
189 )
190
191 def calc_label(self) -> int:
192 return combine_labels(self.class_label, self.element_strategy.label)
193
194 def do_validate(self) -> None:
195 self.element_strategy.validate()
196 if self.is_empty:
197 raise InvalidArgument(
198 "Cannot create non-empty lists with elements drawn from "
199 f"strategy {self.element_strategy!r} because it has no values."
200 )
201 if self.element_strategy.is_empty and 0 < self.max_size < float("inf"):
202 raise InvalidArgument(
203 f"Cannot create a collection of max_size={self.max_size!r}, "
204 "because no elements can be drawn from the element strategy "
205 f"{self.element_strategy!r}"
206 )
207
208 def calc_is_empty(self, recur: RecurT) -> bool:
209 if self.min_size == 0:
210 return False
211 return recur(self.element_strategy)
212
213 def do_draw(self, data: ConjectureData) -> list[Ex]:
214 if self.element_strategy.is_empty:
215 assert self.min_size == 0
216 return []
217
218 elements = cu.many(
219 data,
220 min_size=self.min_size,
221 max_size=self.max_size,
222 average_size=self.average_size,
223 )
224 result = []
225 while elements.more():
226 result.append(data.draw(self.element_strategy))
227 return result
228
229 def __repr__(self) -> str:
230 return (
231 f"{self.__class__.__name__}({self.element_strategy!r}, "
232 f"min_size={self.min_size:_}, max_size={self.max_size:_})"
233 )
234
235 @overload
236 def filter(
237 self, condition: Callable[[list[Ex]], TypeGuard[T]]
238 ) -> "SearchStrategy[T]": ...
239 @overload
240 def filter(
241 self, condition: Callable[[list[Ex]], Any]
242 ) -> "SearchStrategy[list[Ex]]": ...
243 def filter(self, condition):
244 if condition in self._nonempty_filters or is_identity_function(condition):
245 assert self.max_size >= 1, "Always-empty is special cased in st.lists()"
246 if self.min_size >= 1:
247 return self
248 new = copy.copy(self)
249 new.min_size = 1
250 return new
251
252 constraints, pred = get_integer_predicate_bounds(condition)
253 if constraints.get("len") and (
254 "min_value" in constraints or "max_value" in constraints
255 ):
256 new = copy.copy(self)
257 new.min_size = max(
258 self.min_size, constraints.get("min_value", self.min_size)
259 )
260 new.max_size = min(
261 self.max_size, constraints.get("max_value", self.max_size)
262 )
263 # Unsatisfiable filters are easiest to understand without rewriting.
264 if new.min_size > new.max_size:
265 return SearchStrategy.filter(self, condition)
266 # Recompute average size; this is cheaper than making it into a property.
267 new.average_size = min(
268 max(new.min_size * 2, new.min_size + 5),
269 0.5 * (new.min_size + new.max_size),
270 )
271 if pred is None:
272 return new
273 return SearchStrategy.filter(new, condition)
274
275 return SearchStrategy.filter(self, condition)
276
277
278class UniqueListStrategy(ListStrategy[Ex]):
279 def __init__(
280 self,
281 elements: SearchStrategy[Ex],
282 min_size: int,
283 max_size: float | int | None,
284 # TODO: keys are guaranteed to be Hashable, not just Any, but this makes
285 # other things harder to type
286 keys: tuple[Callable[[Ex], Any], ...],
287 tuple_suffixes: SearchStrategy[tuple[Ex, ...]] | None,
288 ):
289 super().__init__(elements, min_size, max_size)
290 self.keys = keys
291 self.tuple_suffixes = tuple_suffixes
292
293 def do_draw(self, data: ConjectureData) -> list[Ex]:
294 if self.element_strategy.is_empty:
295 assert self.min_size == 0
296 return []
297
298 elements = cu.many(
299 data,
300 min_size=self.min_size,
301 max_size=self.max_size,
302 average_size=self.average_size,
303 )
304 seen_sets: tuple[set[Ex], ...] = tuple(set() for _ in self.keys)
305 # actually list[Ex], but if self.tuple_suffixes is present then Ex is a
306 # tuple[T, ...] because self.element_strategy is a TuplesStrategy, and
307 # appending a concrete tuple to `result: list[Ex]` makes mypy unhappy
308 # without knowing that Ex = tuple.
309 result: list[Any] = []
310
311 # We construct a filtered strategy here rather than using a check-and-reject
312 # approach because some strategies have special logic for generation under a
313 # filter, and FilteredStrategy can consolidate multiple filters.
314 def not_yet_in_unique_list(val: Ex) -> bool: # type: ignore # covariant type param
315 return all(
316 key(val) not in seen
317 for key, seen in zip(self.keys, seen_sets, strict=True)
318 )
319
320 filtered = FilteredStrategy(
321 self.element_strategy, conditions=(not_yet_in_unique_list,)
322 )
323 while elements.more():
324 value = filtered.do_filtered_draw(data)
325 if value is filter_not_satisfied:
326 elements.reject(f"Aborted test because unable to satisfy {filtered!r}")
327 else:
328 assert not isinstance(value, UniqueIdentifier)
329 for key, seen in zip(self.keys, seen_sets, strict=True):
330 seen.add(key(value))
331 if self.tuple_suffixes is not None:
332 value = (value, *data.draw(self.tuple_suffixes))
333 result.append(value)
334 assert self.max_size >= len(result) >= self.min_size
335 return result
336
337
338class UniqueSampledListStrategy(UniqueListStrategy):
339 def do_draw(self, data: ConjectureData) -> list[Ex]:
340 assert isinstance(self.element_strategy, SampledFromStrategy)
341
342 should_draw = cu.many(
343 data,
344 min_size=self.min_size,
345 max_size=self.max_size,
346 average_size=self.average_size,
347 )
348 seen_sets: tuple[set[Ex], ...] = tuple(set() for _ in self.keys)
349 result: list[Any] = []
350
351 remaining = LazySequenceCopy(self.element_strategy.elements)
352
353 while remaining and should_draw.more():
354 j = data.draw_integer(0, len(remaining) - 1)
355 value = self.element_strategy._transform(remaining.pop(j))
356 if value is not filter_not_satisfied and all(
357 key(value) not in seen
358 for key, seen in zip(self.keys, seen_sets, strict=True)
359 ):
360 for key, seen in zip(self.keys, seen_sets, strict=True):
361 seen.add(key(value))
362 if self.tuple_suffixes is not None:
363 value = (value, *data.draw(self.tuple_suffixes))
364 result.append(value)
365 else:
366 should_draw.reject(
367 "UniqueSampledListStrategy filter not satisfied or value already seen"
368 )
369 assert self.max_size >= len(result) >= self.min_size
370 return result
371
372
373class FixedDictStrategy(SearchStrategy[dict[Any, Any]]):
374 """A strategy which produces dicts with a fixed set of keys, given a
375 strategy for each of their equivalent values.
376
377 e.g. {'foo' : some_int_strategy} would generate dicts with the single
378 key 'foo' mapping to some integer.
379 """
380
381 def __init__(
382 self,
383 mapping: dict[Any, SearchStrategy[Any]],
384 *,
385 optional: dict[Any, SearchStrategy[Any]] | None,
386 ):
387 super().__init__()
388 dict_type = type(mapping)
389 self.mapping = mapping
390 keys = tuple(mapping.keys())
391 self.fixed = st.tuples(*[mapping[k] for k in keys]).map(
392 lambda value: dict_type(zip(keys, value, strict=True))
393 )
394 self.optional = optional
395
396 def do_draw(self, data: ConjectureData) -> dict[Any, Any]:
397 context = current_build_context()
398 arg_labels: ArgLabelsT = {}
399 value = type(self.mapping)()
400
401 for key, strategy in self.mapping.items():
402 with context.track_arg_label(str(key)) as arg_label:
403 value[key] = data.draw(strategy)
404 arg_labels |= arg_label
405
406 if self.optional is not None:
407 remaining = [k for k, v in self.optional.items() if not v.is_empty]
408 should_draw = cu.many(
409 data,
410 min_size=0,
411 max_size=len(remaining),
412 average_size=len(remaining) / 2,
413 )
414 while should_draw.more():
415 j = data.draw_integer(0, len(remaining) - 1)
416 remaining[-1], remaining[j] = remaining[j], remaining[-1]
417 key = remaining.pop()
418 with context.track_arg_label(str(key)) as arg_label:
419 value[key] = data.draw(self.optional[key])
420 arg_labels |= arg_label
421
422 if arg_labels:
423 context.known_object_printers[IDKey(value)].append(
424 _fixeddict_pprinter(arg_labels, self.mapping)
425 )
426 return value
427
428 def calc_is_empty(self, recur: RecurT) -> bool:
429 return recur(self.fixed)
430
431 def __repr__(self) -> str:
432 if self.optional is not None:
433 return f"fixed_dictionaries({self.mapping!r}, optional={self.optional!r})"
434 return f"fixed_dictionaries({self.mapping!r})"