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, 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
262class many:
263 """Utility class for collections. Bundles up the logic we use for "should I
264 keep drawing more values?" and handles starting and stopping spans in
265 the right place.
266
267 Intended usage is something like:
268
269 elements = many(data, ...)
270 while elements.more():
271 add_stuff_to_result()
272 """
273
274 def __init__(
275 self,
276 data: "ConjectureData",
277 min_size: int,
278 max_size: int | float,
279 average_size: int | float,
280 *,
281 forced: int | None = None,
282 observe: bool = True,
283 ) -> None:
284 assert 0 <= min_size <= average_size <= max_size
285 assert forced is None or min_size <= forced <= max_size
286 self.min_size = min_size
287 self.max_size = max_size
288 self.data = data
289 self.forced_size = forced
290 self.p_continue = _calc_p_continue(average_size - min_size, max_size - min_size)
291 self.count = 0
292 self.rejections = 0
293 self.drawn = False
294 self.force_stop = False
295 self.rejected = False
296 self.observe = observe
297
298 def stop_span(self, *, discard: bool = False) -> None:
299 if self.observe:
300 self.data.stop_span(discard=discard)
301
302 def start_span(self, label):
303 if self.observe:
304 self.data.start_span(label)
305
306 def more(self) -> bool:
307 """Should I draw another element to add to the collection?"""
308 if self.drawn:
309 # A rejected element does not contribute to the collection, so
310 # discard its span - the shrinker can then delete it wholesale.
311 self.stop_span(discard=self.rejected)
312
313 self.drawn = True
314 self.rejected = False
315
316 self.start_span(ONE_FROM_MANY_LABEL)
317 if self.min_size == self.max_size:
318 # if we have to hit an exact size, draw unconditionally until that
319 # point, and no further.
320 should_continue = self.count < self.min_size
321 else:
322 forced_result = None
323 if self.force_stop:
324 # if our size is forced, we can't reject in a way that would
325 # cause us to differ from the forced size.
326 assert self.forced_size is None or self.count == self.forced_size
327 forced_result = False
328 elif self.count < self.min_size:
329 forced_result = True
330 elif self.count >= self.max_size:
331 forced_result = False
332 elif self.forced_size is not None:
333 forced_result = self.count < self.forced_size
334 should_continue = self.data.draw_boolean(
335 self.p_continue,
336 forced=forced_result,
337 observe=self.observe,
338 )
339
340 if should_continue:
341 self.count += 1
342 return True
343 else:
344 self.stop_span()
345 return False
346
347 def reject(self, why: str | None = None) -> None:
348 """Reject the last element (i.e. don't count it towards our budget of
349 elements because it's not going to go in the final collection)."""
350 assert self.count > 0
351 self.count -= 1
352 self.rejections += 1
353 self.rejected = True
354 # We set a minimum number of rejections before we give up to avoid
355 # failing too fast when we reject the first draw.
356 if self.rejections > max(3, 2 * self.count):
357 if self.count < self.min_size:
358 self.data.mark_invalid(why)
359 else:
360 self.force_stop = True
361
362
363class invert_many:
364 """The inversion counterpart of ``many``: the boolean choices that
365 ``many`` would draw around each element. Fixed-size collections draw no
366 booleans at all; variable-size ones draw a continuation boolean before
367 each element (forced while below min_size, but forced draws still consume
368 a choice) and a final False to stop (forced at max_size)."""
369
370 def __init__(self, min_size: int, max_size: int | float) -> None:
371 self._variable_size = min_size != max_size
372
373 def more(self) -> tuple["ChoiceT", ...]:
374 return (True,) if self._variable_size else ()
375
376 def done(self) -> tuple["ChoiceT", ...]:
377 return (False,) if self._variable_size else ()
378
379
380SMALLEST_POSITIVE_FLOAT: float = next_up(0.0) or sys.float_info.min
381
382
383@lru_cache
384def _calc_p_continue(desired_avg: float, max_size: int | float) -> float:
385 """Return the p_continue which will generate the desired average size."""
386 assert desired_avg <= max_size, (desired_avg, max_size)
387 if desired_avg == max_size:
388 return 1.0
389 p_continue = 1 - 1.0 / (1 + desired_avg)
390 if p_continue == 0 or max_size == math.inf:
391 assert 0 <= p_continue < 1, p_continue
392 return p_continue
393 assert 0 < p_continue < 1, p_continue
394 # For small max_size, the infinite-series p_continue is a poor approximation,
395 # and while we can't solve the polynomial a few rounds of iteration quickly
396 # gets us a good approximate solution in almost all cases (sometimes exact!).
397 while _p_continue_to_avg(p_continue, max_size) > desired_avg:
398 # This is impossible over the reals, but *can* happen with floats.
399 p_continue -= 0.0001
400 # If we've reached zero or gone negative, we want to break out of this loop,
401 # and do so even if we're on a system with the unsafe denormals-are-zero flag.
402 # We make that an explicit error in st.floats(), but here we'd prefer to
403 # just get somewhat worse precision on collection lengths.
404 if p_continue < SMALLEST_POSITIVE_FLOAT:
405 p_continue = SMALLEST_POSITIVE_FLOAT
406 break
407 # Let's binary-search our way to a better estimate! We tried fancier options
408 # like gradient descent, but this is numerically stable and works better.
409 hi = 1.0
410 while desired_avg - _p_continue_to_avg(p_continue, max_size) > 0.01:
411 assert 0 < p_continue < hi, (p_continue, hi)
412 mid = (p_continue + hi) / 2
413 if _p_continue_to_avg(mid, max_size) <= desired_avg:
414 p_continue = mid
415 else:
416 hi = mid
417 assert 0 < p_continue < 1, p_continue
418 assert _p_continue_to_avg(p_continue, max_size) <= desired_avg
419 return p_continue
420
421
422def _p_continue_to_avg(p_continue: float, max_size: int | float) -> float:
423 """Return the average_size generated by this p_continue and max_size."""
424 if p_continue >= 1:
425 return max_size
426 return (1.0 / (1 - p_continue) - 1) * (1 - p_continue**max_size)