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 enum
12import hashlib
13import heapq
14import math
15import sys
16from collections import OrderedDict, abc
17from collections.abc import Callable, Iterator, Sequence
18from functools import lru_cache
19from types import FunctionType
20from typing import TYPE_CHECKING, TypeVar
21
22from hypothesis.errors import InvalidArgument
23from hypothesis.internal.compat import int_from_bytes
24from hypothesis.internal.floats import next_up
25from hypothesis.internal.lambda_sources import _function_key
26
27if TYPE_CHECKING:
28 from hypothesis.internal.conjecture.choice import ChoiceT
29 from hypothesis.internal.conjecture.data import ConjectureData
30
31
32LABEL_MASK = 2**64 - 1
33
34
35def calc_label_from_name(name: str) -> int:
36 hashed = hashlib.sha384(name.encode()).digest()
37 return int_from_bytes(hashed[:8])
38
39
40def calc_label_from_callable(f: Callable) -> int:
41 if isinstance(f, FunctionType):
42 return calc_label_from_hash(_function_key(f, ignore_name=True))
43 elif isinstance(f, type):
44 return calc_label_from_cls(f)
45 else:
46 # probably an instance defining __call__
47 try:
48 return calc_label_from_hash(f)
49 except Exception:
50 # not hashable
51 return calc_label_from_cls(type(f))
52
53
54def calc_label_from_cls(cls: type) -> int:
55 return calc_label_from_name(cls.__qualname__)
56
57
58def calc_label_from_hash(obj: object) -> int:
59 return calc_label_from_name(str(hash(obj)))
60
61
62def combine_labels(*labels: int) -> int:
63 label = 0
64 for l in labels:
65 label = (label << 1) & LABEL_MASK
66 label ^= l
67 return label
68
69
70SAMPLE_IN_SAMPLER_LABEL = calc_label_from_name("a sample() in Sampler")
71ONE_FROM_MANY_LABEL = calc_label_from_name("one more from many()")
72
73
74T = TypeVar("T")
75
76
77def identity(v: T) -> T:
78 return v
79
80
81def fisher_yates_shuffle(data: "ConjectureData", ls: list[T]) -> None:
82 """Shuffle ``ls`` in place, drawing from ``data``.
83
84 Reversed Fisher-Yates shuffle: swap each element with itself or with a
85 later element. This shrinks i==j for each element, i.e. towards no change,
86 so a shuffled sequence shrinks back to its original order. We don't
87 consider the last element as it's always a no-op.
88 """
89 for i in range(len(ls) - 1):
90 j = data.draw_integer(i, len(ls) - 1)
91 ls[i], ls[j] = ls[j], ls[i]
92
93
94def check_sample(
95 values: type[enum.Enum] | Sequence[T], strategy_name: str
96) -> Sequence[T]:
97 if "numpy" in sys.modules and isinstance(values, sys.modules["numpy"].ndarray):
98 if values.ndim != 1:
99 raise InvalidArgument(
100 "Only one-dimensional arrays are supported for sampling, "
101 f"and the given value has {values.ndim} dimensions (shape "
102 f"{values.shape}). This array would give samples of array slices "
103 "instead of elements! Use np.ravel(values) to convert "
104 "to a one-dimensional array, or tuple(values) if you "
105 "want to sample slices."
106 )
107 elif not isinstance(values, (OrderedDict, abc.Sequence, enum.EnumMeta)):
108 raise InvalidArgument(
109 f"Cannot sample from {values!r} because it is not an ordered collection. "
110 f"Hypothesis goes to some length to ensure that the {strategy_name} "
111 "strategy has stable results between runs. To replay a saved "
112 "test case, the sampled values must have the same iteration order "
113 "on every run - ruling out sets, dicts, etc due to hash "
114 "randomization. Most cases can simply use `sorted(values)`, but "
115 "mixed types or special values such as math.nan require careful "
116 "handling - and note that when shrinking a test case, "
117 "Hypothesis treats earlier values as simpler."
118 )
119 if isinstance(values, range):
120 # Pyright is unhappy with every way I've tried to type-annotate this
121 # function, so fine, we'll just ignore the analysis error.
122 return values # type: ignore
123 return tuple(values)
124
125
126@lru_cache(64)
127def compute_sampler_table(weights: tuple[float, ...]) -> list[tuple[int, int, float]]:
128 n = len(weights)
129 table: list[list[int | float | None]] = [[i, None, None] for i in range(n)]
130 total = sum(weights)
131 num_type = type(total)
132
133 zero = num_type(0) # type: ignore
134 one = num_type(1) # type: ignore
135
136 small: list[int] = []
137 large: list[int] = []
138
139 probabilities = [w / total for w in weights]
140 scaled_probabilities: list[float] = []
141
142 for i, alternate_chance in enumerate(probabilities):
143 scaled = alternate_chance * n
144 scaled_probabilities.append(scaled)
145 if scaled == 1:
146 table[i][2] = zero
147 elif scaled < 1:
148 small.append(i)
149 else:
150 large.append(i)
151 heapq.heapify(small)
152 heapq.heapify(large)
153
154 while small and large:
155 lo = heapq.heappop(small)
156 hi = heapq.heappop(large)
157
158 assert lo != hi
159 assert scaled_probabilities[hi] > one
160 assert table[lo][1] is None
161 table[lo][1] = hi
162 table[lo][2] = one - scaled_probabilities[lo]
163 scaled_probabilities[hi] = (
164 scaled_probabilities[hi] + scaled_probabilities[lo]
165 ) - one
166
167 if scaled_probabilities[hi] < 1:
168 heapq.heappush(small, hi)
169 elif scaled_probabilities[hi] == 1:
170 table[hi][2] = zero
171 else:
172 heapq.heappush(large, hi)
173 while large:
174 table[large.pop()][2] = zero
175 while small:
176 table[small.pop()][2] = zero
177
178 new_table: list[tuple[int, int, float]] = []
179 for base, alternate, alternate_chance in table:
180 assert isinstance(base, int)
181 assert isinstance(alternate, int) or alternate is None
182 assert alternate_chance is not None
183 if alternate is None:
184 new_table.append((base, base, alternate_chance))
185 elif alternate < base:
186 new_table.append((alternate, base, one - alternate_chance))
187 else:
188 new_table.append((base, alternate, alternate_chance))
189 new_table.sort()
190 return new_table
191
192
193class Sampler:
194 """Sampler based on Vose's algorithm for the alias method. See
195 http://www.keithschwarz.com/darts-dice-coins/ for a good explanation.
196
197 The general idea is that we store a table of triples (base, alternate, p).
198 base. We then pick a triple uniformly at random, and choose its alternate
199 value with probability p and else choose its base value. The triples are
200 chosen so that the resulting mixture has the right distribution.
201
202 We maintain the following invariants to try to produce good shrinks:
203
204 1. The table is in lexicographic (base, alternate) order, so that choosing
205 an earlier value in the list always lowers (or at least leaves
206 unchanged) the value.
207 2. base[i] < alternate[i], so that shrinking the draw always results in
208 shrinking the chosen element.
209 """
210
211 table: list[tuple[int, int, float]] # (base_idx, alt_idx, alt_chance)
212
213 def __init__(self, weights: Sequence[float], *, observe: bool = True):
214 self.observe = observe
215 self.table = compute_sampler_table(tuple(weights))
216
217 def sample(
218 self,
219 data: "ConjectureData",
220 *,
221 forced: int | None = None,
222 ) -> int:
223 if self.observe:
224 data.start_span(SAMPLE_IN_SAMPLER_LABEL)
225 forced_choice = ( # pragma: no branch # https://github.com/nedbat/coveragepy/issues/1617
226 None
227 if forced is None
228 else next(
229 (base, alternate, alternate_chance)
230 for (base, alternate, alternate_chance) in self.table
231 if forced == base or (forced == alternate and alternate_chance > 0)
232 )
233 )
234 base, alternate, alternate_chance = data.choice(
235 self.table,
236 forced=forced_choice,
237 observe=self.observe,
238 )
239 forced_use_alternate = None
240 if forced is not None:
241 # we maintain this invariant when picking forced_choice above.
242 # This song and dance about alternate_chance > 0 is to avoid forcing
243 # e.g. draw_boolean(p=0, forced=True), which is an error.
244 forced_use_alternate = forced == alternate and alternate_chance > 0
245 assert forced == base or forced_use_alternate
246
247 use_alternate = data.draw_boolean(
248 alternate_chance,
249 forced=forced_use_alternate,
250 observe=self.observe,
251 )
252 if self.observe:
253 data.stop_span()
254 if use_alternate:
255 assert forced is None or alternate == forced, (forced, alternate)
256 return alternate
257 else:
258 assert forced is None or base == forced, (forced, base)
259 return base
260
261
262_INTERLEAVED_MANY_MSG = (
263 "Iteration over many() is interleaved with other draws. Are you iterating "
264 "over several many() objects at once, e.g. with zip()?"
265)
266
267
268class many:
269 """Utility class for collections. Bundles up the logic we use for "should I
270 keep drawing more values?" and handles starting and stopping spans in
271 the right place.
272
273 Intended usage is something like:
274
275 with many(data, ...) as elements:
276 for _ in elements:
277 if not add_stuff_to_result():
278 elements.reject()
279
280 Leaving the loop early with ``break``, ``return``, or an exception is
281 supported, and keeps both the span tree and the choice sequence well-formed.
282 """
283
284 def __init__(
285 self,
286 data: "ConjectureData",
287 min_size: int,
288 max_size: int | float,
289 average_size: int | float,
290 *,
291 forced: int | None = None,
292 observe: bool = True,
293 ) -> None:
294 assert 0 <= min_size <= average_size <= max_size
295 assert forced is None or min_size <= forced <= max_size
296 self.min_size = min_size
297 self.max_size = max_size
298 self.data = data
299 self.forced_size = forced
300 self.p_continue = _calc_p_continue(average_size - min_size, max_size - min_size)
301 self.count = 0
302 self.rejections = 0
303 self.force_stop = False
304 self.observe = observe
305 self._depth: int | None = None
306 self._iterated = False
307 self._closed = False
308 self._done = False
309 self._in_element = False
310 self._rejected = False
311
312 def __enter__(self) -> "many":
313 if self._depth is not None:
314 raise InvalidArgument("Cannot reuse a many() object.")
315 self._depth = self.data.depth
316 return self
317
318 def __exit__(self, exc_type: type | None, *_: object) -> None:
319 self._closed = True
320 if self.data.frozen:
321 return
322 if self.data.depth != self._expected_depth():
323 if exc_type is None:
324 raise InvalidArgument(_INTERLEAVED_MANY_MSG)
325 return
326 self._finish_element()
327 if exc_type is None and not self._done:
328 # We exited the loop early, so draw a (forced) stop decision. This
329 # makes our choice sequence look like a natural stop.
330 self.force_stop = True
331 self._decide()
332
333 def __iter__(self) -> Iterator[None]:
334 if self._depth is None or self._closed:
335 raise InvalidArgument(
336 "many() must be used as `with many(...) as elements:` "
337 "followed by `for _ in elements:`, inside the with-block."
338 )
339 if self._iterated:
340 raise InvalidArgument("Cannot iterate over a many() more than once.")
341 self._iterated = True
342 return self._iterate()
343
344 def _iterate(self) -> Iterator[None]:
345 data = self.data
346 while True:
347 if self._closed:
348 raise InvalidArgument("Cannot iterate over many() after exiting it.")
349 if data.depth != self._expected_depth():
350 raise InvalidArgument(_INTERLEAVED_MANY_MSG)
351 if self._in_element:
352 self._finish_element()
353 if not self._decide():
354 return
355 self.count += 1
356 self._in_element = True
357 self._rejected = False
358 yield
359
360 def reject(self, why: str | None = None) -> None:
361 """Reject the current element (i.e. don't count it towards our budget
362 of elements because it's not going to go in the final collection)."""
363 if not self._in_element or self._rejected:
364 raise InvalidArgument(
365 "reject() can only be called once per element, inside the loop."
366 )
367 self.count -= 1
368 self.rejections += 1
369 self._rejected = True
370 # We set a minimum number of rejections before we give up to avoid
371 # failing too fast when we reject the first draw.
372 if self.rejections > max(3, 2 * self.count):
373 if self.count < self.min_size:
374 self.data.mark_invalid(why)
375 else:
376 self.force_stop = True
377
378 def _expected_depth(self) -> int:
379 assert self._depth is not None
380 return self._depth + (self._in_element and self.observe)
381
382 def _finish_element(self) -> None:
383 if self._in_element:
384 self._in_element = False
385 if self.observe:
386 # A rejected element does not contribute to the collection, so
387 # discard its span - the shrinker can then delete it wholesale.
388 self.data.stop_span(discard=self._rejected)
389
390 def _decide(self) -> bool:
391 if self.force_stop and self.count < self.min_size:
392 self.data.mark_invalid(
393 f"stopped with {self.count} elements, below min_size={self.min_size}"
394 )
395 if self.observe:
396 self.data.start_span(ONE_FROM_MANY_LABEL)
397 if self.min_size == self.max_size:
398 # if we have to hit an exact size, draw unconditionally until that
399 # point, and no further.
400 should_continue = self.count < self.min_size
401 else:
402 forced_result = None
403 if self.force_stop:
404 # if our size is forced, we can't stop in a way that would
405 # cause us to differ from the forced size.
406 assert self.forced_size is None or self.count == self.forced_size
407 forced_result = False
408 elif self.count < self.min_size:
409 forced_result = True
410 elif self.count >= self.max_size:
411 forced_result = False
412 elif self.forced_size is not None:
413 forced_result = self.count < self.forced_size
414 # Don't enumerate "simplest" choices for this draw, so that probes
415 # of collections inside unique collections stop at min_size.
416 depth, self.data.unique_element_depth = self.data.unique_element_depth, 0
417 should_continue = self.data.draw_boolean(
418 self.p_continue,
419 forced=forced_result,
420 observe=self.observe,
421 )
422 self.data.unique_element_depth = depth
423 if not should_continue:
424 self._done = True
425 if self.observe:
426 self.data.stop_span()
427 return should_continue
428
429
430class invert_many:
431 """The inversion counterpart of ``many``: the boolean choices that
432 ``many`` would draw around each element. Fixed-size collections draw no
433 booleans at all; variable-size ones draw a continuation boolean before
434 each element (forced while below min_size, but forced draws still consume
435 a choice) and a final False to stop (forced at max_size)."""
436
437 def __init__(self, min_size: int, max_size: int | float) -> None:
438 self._variable_size = min_size != max_size
439
440 def more(self) -> tuple["ChoiceT", ...]:
441 return (True,) if self._variable_size else ()
442
443 def done(self) -> tuple["ChoiceT", ...]:
444 return (False,) if self._variable_size else ()
445
446
447SMALLEST_POSITIVE_FLOAT: float = next_up(0.0) or sys.float_info.min
448
449
450@lru_cache
451def _calc_p_continue(desired_avg: float, max_size: int | float) -> float:
452 """Return the p_continue which will generate the desired average size."""
453 assert desired_avg <= max_size, (desired_avg, max_size)
454 if desired_avg == max_size:
455 return 1.0
456 p_continue = 1 - 1.0 / (1 + desired_avg)
457 if p_continue == 0 or max_size == math.inf:
458 assert 0 <= p_continue < 1, p_continue
459 return p_continue
460 assert 0 < p_continue < 1, p_continue
461 # For small max_size, the infinite-series p_continue is a poor approximation,
462 # and while we can't solve the polynomial a few rounds of iteration quickly
463 # gets us a good approximate solution in almost all cases (sometimes exact!).
464 while _p_continue_to_avg(p_continue, max_size) > desired_avg:
465 # This is impossible over the reals, but *can* happen with floats.
466 p_continue -= 0.0001
467 # If we've reached zero or gone negative, we want to break out of this loop,
468 # and do so even if we're on a system with the unsafe denormals-are-zero flag.
469 # We make that an explicit error in st.floats(), but here we'd prefer to
470 # just get somewhat worse precision on collection lengths.
471 if p_continue < SMALLEST_POSITIVE_FLOAT:
472 p_continue = SMALLEST_POSITIVE_FLOAT
473 break
474 # Let's binary-search our way to a better estimate! We tried fancier options
475 # like gradient descent, but this is numerically stable and works better.
476 hi = 1.0
477 while desired_avg - _p_continue_to_avg(p_continue, max_size) > 0.01:
478 assert 0 < p_continue < hi, (p_continue, hi)
479 mid = (p_continue + hi) / 2
480 if _p_continue_to_avg(mid, max_size) <= desired_avg:
481 p_continue = mid
482 else:
483 hi = mid
484 assert 0 < p_continue < 1, p_continue
485 assert _p_continue_to_avg(p_continue, max_size) <= desired_avg
486 return p_continue
487
488
489def _p_continue_to_avg(p_continue: float, max_size: int | float) -> float:
490 """Return the average_size generated by this p_continue and max_size."""
491 if p_continue >= 1:
492 return max_size
493 return (1.0 / (1 - p_continue) - 1) * (1 - p_continue**max_size)