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 contextlib
12import gc
13import random
14import sys
15import warnings
16from collections.abc import Callable, Generator, Hashable
17from functools import lru_cache
18from itertools import count
19from random import Random
20from typing import TYPE_CHECKING, Any
21from weakref import WeakValueDictionary
22
23import hypothesis.core
24from hypothesis.errors import HypothesisWarning, InvalidArgument
25from hypothesis.internal.compat import FREE_THREADED_CPYTHON, GRAALPY, PYPY
26
27if TYPE_CHECKING:
28 from typing import Protocol
29
30 # we can't use this at runtime until from_type supports
31 # protocols -- breaks ghostwriter tests
32 class RandomLike(Protocol):
33 def seed(self, *args: Any, **kwargs: Any) -> Any: ...
34 def getstate(self, *args: Any, **kwargs: Any) -> Any: ...
35 def setstate(self, *args: Any, **kwargs: Any) -> Any: ...
36
37else:
38 RandomLike = random.Random
39
40_RKEY = count()
41_global_random_rkey = next(_RKEY)
42# This is effectively a WeakSet, which allows us to associate the saved states
43# with their respective Random instances even as new ones are registered and old
44# ones go out of scope and get garbage collected. Keys are ascending integers.
45RANDOMS_TO_MANAGE: WeakValueDictionary[int, RandomLike] = WeakValueDictionary(
46 {_global_random_rkey: random}
47)
48
49
50class NumpyRandomWrapper:
51 def __init__(self) -> None:
52 assert "numpy" in sys.modules
53 # This class provides a shim that matches the numpy to stdlib random,
54 # and lets us avoid importing Numpy until it's already in use.
55 import numpy.random
56
57 self.seed = numpy.random.seed
58 self.getstate = numpy.random.get_state
59 self.setstate = numpy.random.set_state
60
61
62NP_RANDOM: RandomLike | None = None
63
64
65if not (PYPY or GRAALPY):
66
67 def _get_platform_base_refcount(r: Any) -> int:
68 return sys.getrefcount(r)
69
70 # Determine the number of refcounts created by function scope for
71 # the given platform / version of Python.
72 _PLATFORM_REF_COUNT = _get_platform_base_refcount(object())
73else: # pragma: no cover
74 # PYPY and GRAALPY don't have `sys.getrefcount`
75 _PLATFORM_REF_COUNT = -1
76
77
78def register_random(r: RandomLike) -> None:
79 """Register (a weakref to) the given Random-like instance for management by
80 Hypothesis.
81
82 You can pass instances of structural subtypes of ``random.Random``
83 (i.e., objects with seed, getstate, and setstate methods) to
84 ``register_random(r)`` to have their states seeded and restored in the same
85 way as the global PRNGs from the ``random`` and ``numpy.random`` modules.
86
87 All global PRNGs, from e.g. simulation or scheduling frameworks, should
88 be registered to prevent flaky tests. Hypothesis will ensure that the
89 PRNG state is consistent for all test runs, always seeding them to zero and
90 restoring the previous state after the test, or, reproducibly varied if you
91 choose to use the :func:`~hypothesis.strategies.random_module` strategy.
92
93 ``register_random`` only makes `weakrefs
94 <https://docs.python.org/3/library/weakref.html#module-weakref>`_ to ``r``,
95 thus ``r`` will only be managed by Hypothesis as long as it has active
96 references elsewhere at runtime. The pattern ``register_random(MyRandom())``
97 will raise a ``ReferenceError`` to help protect users from this issue.
98 This check does not occur for the PyPy interpreter. See the following example for
99 an illustration of this issue
100
101 .. code-block:: python
102
103
104 def my_BROKEN_hook():
105 r = MyRandomLike()
106
107 # `r` will be garbage collected after the hook resolved
108 # and Hypothesis will 'forget' that it was registered
109 register_random(r) # Hypothesis will emit a warning
110
111
112 rng = MyRandomLike()
113
114
115 def my_WORKING_hook():
116 register_random(rng)
117 """
118 if not (hasattr(r, "seed") and hasattr(r, "getstate") and hasattr(r, "setstate")):
119 raise InvalidArgument(f"{r=} does not have all the required methods")
120
121 if r in [
122 random
123 for ref in RANDOMS_TO_MANAGE.data.copy().values() # type: ignore
124 if (random := ref()) is not None
125 ]:
126 return
127
128 if not (PYPY or GRAALPY): # pragma: no branch
129 # PYPY and GRAALPY do not have `sys.getrefcount`.
130 gc.collect()
131 if not gc.get_referrers(r):
132 if sys.getrefcount(r) <= _PLATFORM_REF_COUNT:
133 raise ReferenceError(
134 f"`register_random` was passed `r={r}` which will be "
135 "garbage collected immediately after `register_random` creates a "
136 "weakref to it. This will prevent Hypothesis from managing this "
137 "PRNG. See the docs for `register_random` for more "
138 "details."
139 )
140 elif not FREE_THREADED_CPYTHON: # pragma: no branch
141 # On CPython, check for the free-threaded build because
142 # gc.get_referrers() ignores objects with immortal refcounts
143 # and objects are immortalized in the Python 3.13
144 # free-threading implementation at runtime.
145
146 warnings.warn(
147 "It looks like `register_random` was passed an object that could "
148 "be garbage collected immediately after `register_random` creates "
149 "a weakref to it. This will prevent Hypothesis from managing this "
150 "PRNG. See the docs for `register_random` for more details.",
151 HypothesisWarning,
152 stacklevel=2,
153 )
154
155 RANDOMS_TO_MANAGE[next(_RKEY)] = r
156
157
158# Used to make the warning issued by `deprecate_random_in_strategy` thread-safe,
159# as well as to avoid warning on uses of st.randoms().
160# Store just the hash to reduce memory consumption. This is an underapproximation
161# of membership (distinct items might have the same hash), which is fine for the
162# warning, as it results in missed alarms, not false alarms.
163_known_random_state_hashes: set[Any] = set()
164
165
166@lru_cache(maxsize=64)
167def _seeded_state(seed: int) -> tuple[Any, int]:
168 state = Random(seed).getstate()
169 return state, hash(state)
170
171
172def get_seeder_and_restorer(
173 seed: Hashable = 0,
174) -> tuple[Callable[[], None], Callable[[], None]]:
175 """Return a pair of functions which respectively seed all and restore
176 the state of all registered PRNGs.
177
178 This is used by the core engine via `deterministic_PRNG`, and by users
179 via `register_random`. We support registration of additional random.Random
180 instances (or other objects with seed, getstate, and setstate methods)
181 to force determinism on simulation or scheduling frameworks which avoid
182 using the global random state. See e.g. #1709.
183 """
184 assert isinstance(seed, int)
185 assert 0 <= seed < 2**32
186 states: dict[int, object] = {}
187 seeded_state, seeded_hash = _seeded_state(seed)
188
189 if "numpy" in sys.modules:
190 global NP_RANDOM
191 if NP_RANDOM is None:
192 # Protect this from garbage-collection by adding it to global scope
193 NP_RANDOM = RANDOMS_TO_MANAGE[next(_RKEY)] = NumpyRandomWrapper()
194
195 def seed_all() -> None:
196 assert not states
197 # access .data.copy().items() instead of .items() to avoid a "dictionary
198 # changed size during iteration" error under multithreading.
199 #
200 # I initially expected this to be fixed by
201 # https://github.com/python/cpython/commit/96d37dbcd23e65a7a57819aeced9034296ef747e,
202 # but I believe that is addressing the size change from weakrefs expiring
203 # during gc, not from the user adding new elements to the dict.
204 #
205 # Since we're accessing .data, we have to manually handle checking for
206 # expired ref instances during iteration. Normally WeakValueDictionary
207 # handles this for us.
208 #
209 # This command reproduces at time of writing:
210 # pytest hypothesis/tests/ -k test_intervals_are_equivalent_to_their_lists
211 # --parallel-threads 2
212 for k, ref in RANDOMS_TO_MANAGE.data.copy().items(): # type: ignore
213 r = ref()
214 if r is None:
215 # ie the random instance has been gc'd
216 continue # pragma: no cover
217 states[k] = r.getstate()
218 if k == _global_random_rkey:
219 # r.seed sets the random's state. We want to add that state to
220 # _known_random_states before calling r.seed, in case a thread
221 # switch occurs between the two.
222 _known_random_state_hashes.add(seeded_hash)
223
224 if type(r) is Random:
225 r.setstate(seeded_state)
226 else:
227 r.seed(seed)
228
229 def restore_all() -> None:
230 for k, state in states.items():
231 r = RANDOMS_TO_MANAGE.get(k)
232 if r is None: # i.e., has been garbage-collected
233 continue
234
235 if k == _global_random_rkey:
236 _known_random_state_hashes.add(hash(state))
237 r.setstate(state)
238
239 states.clear()
240
241 return seed_all, restore_all
242
243
244@contextlib.contextmanager
245def deterministic_PRNG(seed: int = 0) -> Generator[None, None, None]:
246 """Context manager that handles random.seed without polluting global state.
247
248 See issue #1266 and PR #1295 for details and motivation - in short,
249 leaving the global pseudo-random number generator (PRNG) seeded is a very
250 bad idea in principle, and breaks all kinds of independence assumptions
251 in practice.
252 """
253 if hypothesis.core.threadlocal._hypothesis_global_random is None:
254 hypothesis.core.threadlocal._hypothesis_global_random = Random()
255 register_random(hypothesis.core.threadlocal._hypothesis_global_random)
256
257 seed_all, restore_all = get_seeder_and_restorer(seed)
258 seed_all()
259 try:
260 yield
261 finally:
262 restore_all()
263 # TODO it would be nice to clean up _known_random_state_hashes when no
264 # active deterministic_PRNG contexts remain, to free memory (see similar
265 # logic in StackframeLimiter). But it's a bit annoying to get right, and
266 # likely not a big deal.