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