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, Mapping
14from typing import Any, TypeGuard, overload
15
16from hypothesis import strategies as st
17from hypothesis.control import current_build_context
18from hypothesis.errors import CannotInvert, InvalidArgument
19from hypothesis.internal.compat import add_note
20from hypothesis.internal.conjecture import utils as cu
21from hypothesis.internal.conjecture.choice import ChoiceT
22from hypothesis.internal.conjecture.data import UNIQUE_COLLECTION_LABEL, ConjectureData
23from hypothesis.internal.conjecture.engine import BUFFER_SIZE
24from hypothesis.internal.conjecture.junkdrawer import LazySequenceCopy, equal_values
25from hypothesis.internal.conjecture.utils import combine_labels
26from hypothesis.internal.filtering import get_integer_predicate_bounds
27from hypothesis.internal.reflection import is_identity_function
28from hypothesis.strategies._internal.strategies import (
29 T1,
30 T3,
31 T4,
32 T5,
33 Ex,
34 FilteredStrategy,
35 RecurT,
36 SampledFromStrategy,
37 SearchStrategy,
38 T,
39 check_strategy,
40 filter_not_satisfied,
41)
42from hypothesis.strategies._internal.utils import cacheable, defines_strategy
43from hypothesis.utils.conventions import UniqueIdentifier
44from hypothesis.vendor.pretty import (
45 ArgLabelsT,
46 IDKey,
47 _fixeddict_pprinter,
48 _tuple_pprinter,
49)
50
51
52class TupleStrategy(SearchStrategy[tuple[Ex, ...]]):
53 """A strategy responsible for fixed length tuples based on heterogeneous
54 strategies for each of their elements."""
55
56 def __init__(self, strategies: Iterable[SearchStrategy[Any]]):
57 super().__init__()
58 self.element_strategies = tuple(strategies)
59
60 def do_validate(self) -> None:
61 for s in self.element_strategies:
62 s.validate()
63
64 def calc_label(self) -> int:
65 return combine_labels(
66 self.class_label, *(s.label for s in self.element_strategies)
67 )
68
69 def __repr__(self) -> str:
70 tuple_string = ", ".join(map(repr, self.element_strategies))
71 return f"TupleStrategy(({tuple_string}))"
72
73 def calc_has_reusable_values(self, recur: RecurT) -> bool:
74 return all(recur(e) for e in self.element_strategies)
75
76 def do_draw(self, data: ConjectureData) -> tuple[Ex, ...]:
77 context = current_build_context()
78 arg_labels: ArgLabelsT = {}
79 result = []
80 for i, strategy in enumerate(self.element_strategies):
81 with data.track_arg_label(f"arg[{i}]") as arg_label:
82 result.append(data.draw(strategy))
83 arg_labels |= arg_label
84
85 result = tuple(result)
86 if arg_labels:
87 context.known_object_printers[IDKey(result)].append(
88 _tuple_pprinter(arg_labels)
89 )
90 return result
91
92 def _invert(self, value: Any) -> tuple[ChoiceT, ...]:
93 if not isinstance(value, tuple) or len(value) != len(self.element_strategies):
94 raise CannotInvert(
95 f"{value!r} is not a tuple of length {len(self.element_strategies)}"
96 )
97 choices: list[ChoiceT] = []
98 for i, (strategy, element) in enumerate(
99 zip(self.element_strategies, value, strict=True)
100 ):
101 try:
102 choices.extend(strategy._invert(element))
103 except CannotInvert as exc:
104 add_note(exc, f"at index {i} of {value!r}, strategy={self!r}")
105 raise
106 return tuple(choices)
107
108 def calc_is_empty(self, recur: RecurT) -> bool:
109 return any(recur(e) for e in self.element_strategies)
110
111
112@overload
113def tuples() -> SearchStrategy[tuple[()]]: ...
114
115
116@overload
117def tuples(__a1: SearchStrategy[Ex]) -> SearchStrategy[tuple[Ex]]: ...
118
119
120@overload
121def tuples(
122 __a1: SearchStrategy[T1], __a2: SearchStrategy[T]
123) -> SearchStrategy[tuple[T1, T]]: ...
124
125
126@overload
127def tuples(
128 __a1: SearchStrategy[T1], __a2: SearchStrategy[T], __a3: SearchStrategy[T3]
129) -> SearchStrategy[tuple[T1, T, T3]]: ...
130
131
132@overload
133def tuples(
134 __a1: SearchStrategy[T1],
135 __a2: SearchStrategy[T],
136 __a3: SearchStrategy[T3],
137 __a4: SearchStrategy[T4],
138) -> SearchStrategy[tuple[T1, T, T3, T4]]: ...
139
140
141@overload
142def tuples(
143 __a1: SearchStrategy[T1],
144 __a2: SearchStrategy[T],
145 __a3: SearchStrategy[T3],
146 __a4: SearchStrategy[T4],
147 __a5: SearchStrategy[T5],
148) -> SearchStrategy[tuple[T1, T, T3, T4, T5]]: ...
149
150
151@overload
152def tuples(
153 *args: SearchStrategy[Any],
154) -> SearchStrategy[tuple[Any, ...]]: ...
155
156
157@cacheable
158@defines_strategy()
159def tuples(*args: SearchStrategy[Any]) -> SearchStrategy[tuple[Any, ...]]:
160 """Return a strategy which generates a tuple of the same length as args by
161 generating the value at index i from args[i].
162
163 e.g. tuples(integers(), integers()) would generate a tuple of length
164 two with both values an integer.
165
166 Examples from this strategy shrink by shrinking their component parts.
167 """
168 for arg in args:
169 check_strategy(arg)
170
171 return TupleStrategy(args)
172
173
174class ListStrategy(SearchStrategy[list[Ex]]):
175 """A strategy for lists which takes a strategy for its elements and the
176 allowed lengths, and generates lists with the correct size and contents."""
177
178 _nonempty_filters: tuple[Callable[[Any], Any], ...] = (bool, len, tuple, list)
179
180 def __init__(
181 self,
182 elements: SearchStrategy[Ex],
183 min_size: int = 0,
184 max_size: float | int | None = math.inf,
185 ):
186 super().__init__()
187 self.min_size = min_size or 0
188 self.max_size = max_size if max_size is not None else math.inf
189 assert 0 <= self.min_size <= self.max_size
190 self.average_size = min(
191 max(self.min_size * 2, self.min_size + 5),
192 0.5 * (self.min_size + self.max_size),
193 )
194 self.element_strategy = elements
195 if min_size > BUFFER_SIZE:
196 raise InvalidArgument(
197 f"{self!r} can never generate a value, because min_size is larger "
198 "than Hypothesis supports. Including it is at best slowing down your "
199 "tests for no benefit; at worst making them fail (maybe flakily) with "
200 "a HealthCheck error."
201 )
202
203 def calc_label(self) -> int:
204 return combine_labels(self.class_label, self.element_strategy.label)
205
206 def do_validate(self) -> None:
207 self.element_strategy.validate()
208 if self.is_empty:
209 raise InvalidArgument(
210 "Cannot create non-empty lists with elements drawn from "
211 f"strategy {self.element_strategy!r} because it has no values."
212 )
213 if self.element_strategy.is_empty and 0 < self.max_size < float("inf"):
214 raise InvalidArgument(
215 f"Cannot create a collection of max_size={self.max_size!r}, "
216 "because no elements can be drawn from the element strategy "
217 f"{self.element_strategy!r}"
218 )
219
220 def calc_is_empty(self, recur: RecurT) -> bool:
221 if self.min_size == 0:
222 return False
223 return recur(self.element_strategy)
224
225 def do_draw(self, data: ConjectureData) -> list[Ex]:
226 if self.element_strategy.is_empty:
227 assert self.min_size == 0
228 return []
229
230 with cu.many(
231 data,
232 min_size=self.min_size,
233 max_size=self.max_size,
234 average_size=self.average_size,
235 ) as elements:
236 return [data.draw(self.element_strategy) for _ in elements]
237
238 def _invert(self, value: Any) -> tuple[ChoiceT, ...]:
239 if not isinstance(value, list):
240 raise CannotInvert(f"{value!r} is not a list")
241 if not (self.min_size <= len(value) <= self.max_size):
242 raise CannotInvert(
243 f"len={len(value)} outside "
244 f"[{self.min_size}, {self.max_size!r}] for {self!r}"
245 )
246 if self.element_strategy.is_empty:
247 # do_draw returns [] without drawing anything
248 if value:
249 raise CannotInvert(f"elements of {self!r} are empty")
250 return ()
251 elements = cu.invert_many(self.min_size, self.max_size)
252 choices: list[ChoiceT] = []
253 for i, element in enumerate(value):
254 choices.extend(elements.more())
255 try:
256 choices.extend(self.element_strategy._invert(element))
257 except CannotInvert as exc:
258 add_note(exc, f"at index {i} of {value!r}, strategy={self!r}")
259 raise
260 choices.extend(elements.done())
261 return tuple(choices)
262
263 def __repr__(self) -> str:
264 return (
265 f"{self.__class__.__name__}({self.element_strategy!r}, "
266 f"min_size={self.min_size:_}, max_size={self.max_size:_})"
267 )
268
269 @overload
270 def filter(
271 self, condition: Callable[[list[Ex]], TypeGuard[T]]
272 ) -> "SearchStrategy[T]": ...
273 @overload
274 def filter(
275 self, condition: Callable[[list[Ex]], Any]
276 ) -> "SearchStrategy[list[Ex]]": ...
277 def filter(self, condition):
278 if condition in self._nonempty_filters or is_identity_function(condition):
279 assert self.max_size >= 1, "Always-empty is special cased in st.lists()"
280 if self.min_size >= 1:
281 return self
282 new = copy.copy(self)
283 new.min_size = 1
284 return new
285
286 constraints, pred = get_integer_predicate_bounds(condition)
287 if constraints.get("len") and (
288 "min_value" in constraints or "max_value" in constraints
289 ):
290 new = copy.copy(self)
291 new.min_size = max(
292 self.min_size, constraints.get("min_value", self.min_size)
293 )
294 new.max_size = min(
295 self.max_size, constraints.get("max_value", self.max_size)
296 )
297 # Unsatisfiable filters are easiest to understand without rewriting.
298 if new.min_size > new.max_size:
299 return SearchStrategy.filter(self, condition)
300 # Recompute average size; this is cheaper than making it into a property.
301 new.average_size = min(
302 max(new.min_size * 2, new.min_size + 5),
303 0.5 * (new.min_size + new.max_size),
304 )
305 if pred is None:
306 return new
307 return SearchStrategy.filter(new, condition)
308
309 return SearchStrategy.filter(self, condition)
310
311
312class UniqueListStrategy(ListStrategy[Ex]):
313 def __init__(
314 self,
315 elements: SearchStrategy[Ex],
316 min_size: int,
317 max_size: float | int | None,
318 # TODO: keys are guaranteed to be Hashable, not just Any, but this makes
319 # other things harder to type
320 keys: tuple[Callable[[Ex], Any], ...],
321 tuple_suffixes: SearchStrategy[tuple[Ex, ...]] | None,
322 ):
323 super().__init__(elements, min_size, max_size)
324 self.keys = keys
325 self.tuple_suffixes = tuple_suffixes
326
327 def do_draw(self, data: ConjectureData) -> list[Ex]:
328 if self.element_strategy.is_empty:
329 assert self.min_size == 0
330 return []
331
332 seen_sets: tuple[set[Ex], ...] = tuple(set() for _ in self.keys)
333 # actually list[Ex], but if self.tuple_suffixes is present then Ex is a
334 # tuple[T, ...] because self.element_strategy is a TuplesStrategy, and
335 # appending a concrete tuple to `result: list[Ex]` makes mypy unhappy
336 # without knowing that Ex = tuple.
337 result: list[Any] = []
338
339 # We construct a filtered strategy here rather than using a check-and-reject
340 # approach because some strategies have special logic for generation under a
341 # filter, and FilteredStrategy can consolidate multiple filters.
342 def not_yet_in_unique_list(val: Ex) -> bool: # type: ignore # covariant type param
343 return all(
344 key(val) not in seen
345 for key, seen in zip(self.keys, seen_sets, strict=True)
346 )
347
348 filtered = FilteredStrategy(
349 self.element_strategy, conditions=(not_yet_in_unique_list,)
350 )
351 data.start_span(UNIQUE_COLLECTION_LABEL)
352 if data.unique_element_depth == 0:
353 # Reset the enumeration of "simplest" template draws for each
354 # top-level unique collection, so that probes are canonical. A
355 # collection nested inside another unique collection keeps
356 # counting, so that sibling elements remain distinct.
357 data.simplest_index = 0
358 try:
359 with cu.many(
360 data,
361 min_size=self.min_size,
362 max_size=self.max_size,
363 average_size=self.average_size,
364 ) as elements:
365 for _ in elements:
366 data.unique_element_depth += 1
367 try:
368 value = filtered.do_filtered_draw(data)
369 finally:
370 data.unique_element_depth -= 1
371 if value is filter_not_satisfied:
372 elements.reject(
373 f"Aborted test because unable to satisfy {filtered!r}"
374 )
375 else:
376 assert not isinstance(value, UniqueIdentifier)
377 for key, seen in zip(self.keys, seen_sets, strict=True):
378 seen.add(key(value))
379 if self.tuple_suffixes is not None:
380 value = (value, *data.draw(self.tuple_suffixes))
381 result.append(value)
382 finally:
383 data.stop_span()
384 assert self.max_size >= len(result) >= self.min_size
385 return result
386
387 def _check_unique_keys(self, elements: list[Any]) -> None:
388 for keyfunc in self.keys:
389 try:
390 keys = list(map(keyfunc, elements))
391 unique = len(set(keys)) == len(keys)
392 except Exception:
393 raise CannotInvert(
394 f"could not compute uniqueness keys for {self!r}"
395 ) from None
396 if not unique:
397 raise CannotInvert(f"{elements!r} has duplicate keys for {self!r}")
398
399 def _split_suffixed(self, value: Any) -> list[Any]:
400 # With tuple_suffixes, each element is (key, *suffix): the key drawn
401 # from element_strategy, then the suffix from tuple_suffixes.
402 # Uniqueness applies to the key alone.
403 if not isinstance(value, list):
404 raise CannotInvert(f"{value!r} is not a list")
405 if not all(isinstance(e, tuple) and e for e in value):
406 raise CannotInvert(f"{value!r} is not a list of nonempty tuples")
407 return [e[0] for e in value]
408
409 def _invert(self, value: Any) -> tuple[ChoiceT, ...]:
410 if self.tuple_suffixes is None:
411 # A valid unique list draws exactly like a ListStrategy, since
412 # every element passes the uniqueness filter on its first draw.
413 self._check_unique_keys(value)
414 return ListStrategy._invert(self, value)
415 self._check_unique_keys(self._split_suffixed(value))
416 if not (self.min_size <= len(value) <= self.max_size):
417 raise CannotInvert(
418 f"len={len(value)} outside "
419 f"[{self.min_size}, {self.max_size!r}] for {self!r}"
420 )
421 elements = cu.invert_many(self.min_size, self.max_size)
422 choices: list[ChoiceT] = []
423 for i, element in enumerate(value):
424 choices.extend(elements.more())
425 try:
426 choices.extend(self.element_strategy._invert(element[0]))
427 choices.extend(self.tuple_suffixes._invert(tuple(element[1:])))
428 except CannotInvert as exc:
429 add_note(exc, f"at index {i} of {value!r}, strategy={self!r}")
430 raise
431 choices.extend(elements.done())
432 return tuple(choices)
433
434
435class UniqueSampledListStrategy(UniqueListStrategy):
436 def do_draw(self, data: ConjectureData) -> list[Ex]:
437 assert isinstance(self.element_strategy, SampledFromStrategy)
438
439 seen_sets: tuple[set[Ex], ...] = tuple(set() for _ in self.keys)
440 result: list[Any] = []
441
442 remaining = LazySequenceCopy(self.element_strategy.elements)
443
444 with cu.many(
445 data,
446 min_size=self.min_size,
447 max_size=self.max_size,
448 average_size=self.average_size,
449 ) as elements:
450 for _ in elements:
451 j = data.draw_integer(0, len(remaining) - 1)
452 value = self.element_strategy._transform(remaining.pop(j), data=data)
453 if value is not filter_not_satisfied and all(
454 key(value) not in seen
455 for key, seen in zip(self.keys, seen_sets, strict=True)
456 ):
457 for key, seen in zip(self.keys, seen_sets, strict=True):
458 seen.add(key(value))
459 if self.tuple_suffixes is not None:
460 value = (value, *data.draw(self.tuple_suffixes))
461 result.append(value)
462 else:
463 elements.reject(
464 "UniqueSampledListStrategy filter not satisfied or "
465 "value already seen"
466 )
467 if not remaining:
468 break
469 assert self.max_size >= len(result) >= self.min_size
470 return result
471
472 def _invert(self, value: Any) -> tuple[ChoiceT, ...]:
473 if not isinstance(value, list):
474 raise CannotInvert(f"{value!r} is not a list")
475 if not (self.min_size <= len(value) <= self.max_size):
476 raise CannotInvert(
477 f"len={len(value)} outside "
478 f"[{self.min_size}, {self.max_size!r}] for {self!r}"
479 )
480 targets = value if self.tuple_suffixes is None else self._split_suffixed(value)
481 # do_draw indexes into the pool of not-yet-drawn elements, so each
482 # index is relative to what remains, not to the original elements.
483 assert isinstance(self.element_strategy, SampledFromStrategy)
484 remaining = list(self.element_strategy.elements)
485 elements = cu.invert_many(self.min_size, self.max_size)
486 choices: list[ChoiceT] = []
487 for i, (element, target) in enumerate(zip(value, targets, strict=True)):
488 choices.extend(elements.more())
489 for j, candidate in enumerate(remaining):
490 if equal_values(
491 self.element_strategy._transform(candidate, data=None), target
492 ):
493 choices.append(j)
494 remaining.pop(j)
495 break
496 else:
497 raise CannotInvert(
498 f"at index {i} of {value!r}: {target!r} is not among the "
499 f"remaining elements of {self!r}"
500 )
501 if self.tuple_suffixes is not None:
502 try:
503 choices.extend(self.tuple_suffixes._invert(tuple(element[1:])))
504 except CannotInvert as exc:
505 add_note(exc, f"at index {i} of {value!r}, strategy={self!r}")
506 raise
507 choices.extend(elements.done())
508 return tuple(choices)
509
510
511class FixedDictStrategy(SearchStrategy[Mapping[Any, Any]]):
512 """A strategy which produces mappings with a fixed set of keys, given a
513 strategy for each of their equivalent values.
514
515 e.g. {'foo' : some_int_strategy} would generate mappings with the single
516 key 'foo' mapping to some integer.
517 """
518
519 def __init__(
520 self,
521 mapping: Mapping[Any, SearchStrategy[Any]],
522 *,
523 optional: Mapping[Any, SearchStrategy[Any]] | None,
524 ):
525 super().__init__()
526 dict_type = type(mapping)
527 self.mapping = mapping
528 keys = tuple(mapping.keys())
529 self.fixed = st.tuples(*[mapping[k] for k in keys]).map(
530 lambda value: dict_type(zip(keys, value, strict=True)) # type: ignore
531 )
532 self.optional = optional
533
534 def do_draw(self, data: ConjectureData) -> Mapping[Any, Any]:
535 context = current_build_context()
536 arg_labels: ArgLabelsT = {}
537 pairs: list[tuple[Any, Any]] = []
538
539 for key, strategy in self.mapping.items():
540 with data.track_arg_label(str(key)) as arg_label:
541 pairs.append((key, data.draw(strategy)))
542 arg_labels |= arg_label
543
544 if self.optional is not None:
545 remaining = [k for k, v in self.optional.items() if not v.is_empty]
546 with cu.many(
547 data,
548 min_size=0,
549 max_size=len(remaining),
550 average_size=len(remaining) / 2,
551 ) as elements:
552 for _ in elements:
553 j = data.draw_integer(0, len(remaining) - 1)
554 remaining[-1], remaining[j] = remaining[j], remaining[-1]
555 key = remaining.pop()
556 with data.track_arg_label(str(key)) as arg_label:
557 pairs.append((key, data.draw(self.optional[key])))
558 arg_labels |= arg_label
559
560 # Vary the dict's iteration order (#3906). We shuffle after choosing
561 # the optional keys, so only order varies, not the set of keys.
562 cu.fisher_yates_shuffle(data, pairs)
563 value = type(self.mapping)(pairs) # type: ignore
564
565 if arg_labels:
566 context.known_object_printers[IDKey(value)].append(
567 _fixeddict_pprinter(arg_labels)
568 )
569 return value
570
571 def _invert(self, value: Any) -> tuple[ChoiceT, ...]:
572 if not isinstance(value, Mapping):
573 raise CannotInvert(f"{value!r} is not a mapping")
574 optional = self.optional or {}
575 if any(k not in value for k in self.mapping) or any(
576 k not in self.mapping and k not in optional for k in value
577 ):
578 raise CannotInvert(f"{value!r} has the wrong keys for {self!r}")
579
580 choices: list[ChoiceT] = []
581 npairs = len(self.mapping)
582 for key, strategy in self.mapping.items():
583 try:
584 choices.extend(strategy._invert(value[key]))
585 except CannotInvert as exc:
586 add_note(exc, f"at key {key!r} of {value!r}, strategy={self!r}")
587 raise
588
589 if self.optional is not None:
590 # do_draw selects each present optional key by its index in the
591 # remaining-keys list, which pops via swap-with-last.
592 remaining = [k for k, v in self.optional.items() if not v.is_empty]
593 present = [k for k in remaining if k in value and k not in self.mapping]
594 if npairs + len(present) != len(value):
595 # a present optional key whose strategy is empty
596 raise CannotInvert(f"{value!r} has the wrong keys for {self!r}")
597 selector = cu.invert_many(0, len(remaining))
598 for key in present:
599 choices.extend(selector.more())
600 j = remaining.index(key)
601 choices.append(j)
602 remaining[-1], remaining[j] = remaining[j], remaining[-1]
603 remaining.pop()
604 try:
605 choices.extend(self.optional[key]._invert(value[key]))
606 except CannotInvert as exc:
607 add_note(exc, f"at key {key!r} of {value!r}, strategy={self!r}")
608 raise
609 choices.extend(selector.done())
610 npairs += len(present)
611
612 # an identity shuffle for do_draw's final fisher_yates_shuffle of the
613 # pairs; mappings compare equal regardless of iteration order
614 choices.extend(range(npairs - 1))
615 return tuple(choices)
616
617 def calc_is_empty(self, recur: RecurT) -> bool:
618 return recur(self.fixed)
619
620 def __repr__(self) -> str:
621 if self.optional is not None:
622 return f"fixed_dictionaries({self.mapping!r}, optional={self.optional!r})"
623 return f"fixed_dictionaries({self.mapping!r})"