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 Generator, Hashable
17from itertools import count
18from typing import TYPE_CHECKING, Any, Callable, Optional
19from weakref import WeakValueDictionary
20
21import hypothesis.core
22from hypothesis.errors import HypothesisWarning, InvalidArgument
23from hypothesis.internal.compat import FREE_THREADED_CPYTHON, GRAALPY, PYPY
24
25if TYPE_CHECKING:
26 from typing import Protocol
27
28 # we can't use this at runtime until from_type supports
29 # protocols -- breaks ghostwriter tests
30 class RandomLike(Protocol):
31 def seed(self, *args: Any, **kwargs: Any) -> Any: ...
32 def getstate(self, *args: Any, **kwargs: Any) -> Any: ...
33 def setstate(self, *args: Any, **kwargs: Any) -> Any: ...
34
35else: # pragma: no cover
36 RandomLike = random.Random
37
38# This is effectively a WeakSet, which allows us to associate the saved states
39# with their respective Random instances even as new ones are registered and old
40# ones go out of scope and get garbage collected. Keys are ascending integers.
41_RKEY = count()
42RANDOMS_TO_MANAGE: WeakValueDictionary[int, RandomLike] = WeakValueDictionary(
43 {next(_RKEY): random}
44)
45
46
47class NumpyRandomWrapper:
48 def __init__(self) -> None:
49 assert "numpy" in sys.modules
50 # This class provides a shim that matches the numpy to stdlib random,
51 # and lets us avoid importing Numpy until it's already in use.
52 import numpy.random
53
54 self.seed = numpy.random.seed
55 self.getstate = numpy.random.get_state
56 self.setstate = numpy.random.set_state
57
58
59NP_RANDOM: Optional[RandomLike] = None
60
61
62if not (PYPY or GRAALPY):
63
64 def _get_platform_base_refcount(r: Any) -> int:
65 return sys.getrefcount(r)
66
67 # Determine the number of refcounts created by function scope for
68 # the given platform / version of Python.
69 _PLATFORM_REF_COUNT = _get_platform_base_refcount(object())
70else: # pragma: no cover
71 # PYPY and GRAALPY don't have `sys.getrefcount`
72 _PLATFORM_REF_COUNT = -1
73
74
75def register_random(r: RandomLike) -> None:
76 """Register (a weakref to) the given Random-like instance for management by
77 Hypothesis.
78
79 You can pass instances of structural subtypes of ``random.Random``
80 (i.e., objects with seed, getstate, and setstate methods) to
81 ``register_random(r)`` to have their states seeded and restored in the same
82 way as the global PRNGs from the ``random`` and ``numpy.random`` modules.
83
84 All global PRNGs, from e.g. simulation or scheduling frameworks, should
85 be registered to prevent flaky tests. Hypothesis will ensure that the
86 PRNG state is consistent for all test runs, always seeding them to zero and
87 restoring the previous state after the test, or, reproducibly varied if you
88 choose to use the :func:`~hypothesis.strategies.random_module` strategy.
89
90 ``register_random`` only makes `weakrefs
91 <https://docs.python.org/3/library/weakref.html#module-weakref>`_ to ``r``,
92 thus ``r`` will only be managed by Hypothesis as long as it has active
93 references elsewhere at runtime. The pattern ``register_random(MyRandom())``
94 will raise a ``ReferenceError`` to help protect users from this issue.
95 This check does not occur for the PyPy interpreter. See the following example for
96 an illustration of this issue
97
98 .. code-block:: python
99
100
101 def my_BROKEN_hook():
102 r = MyRandomLike()
103
104 # `r` will be garbage collected after the hook resolved
105 # and Hypothesis will 'forget' that it was registered
106 register_random(r) # Hypothesis will emit a warning
107
108
109 rng = MyRandomLike()
110
111
112 def my_WORKING_hook():
113 register_random(rng)
114 """
115 if not (hasattr(r, "seed") and hasattr(r, "getstate") and hasattr(r, "setstate")):
116 raise InvalidArgument(f"{r=} does not have all the required methods")
117
118 if r in RANDOMS_TO_MANAGE.values():
119 return
120
121 if not (PYPY or GRAALPY): # pragma: no branch
122 # PYPY and GRAALPY do not have `sys.getrefcount`.
123 gc.collect()
124 if not gc.get_referrers(r):
125 if sys.getrefcount(r) <= _PLATFORM_REF_COUNT:
126 raise ReferenceError(
127 f"`register_random` was passed `r={r}` which will be "
128 "garbage collected immediately after `register_random` creates a "
129 "weakref to it. This will prevent Hypothesis from managing this "
130 "PRNG. See the docs for `register_random` for more "
131 "details."
132 )
133 elif not FREE_THREADED_CPYTHON: # pragma: no branch
134 # On CPython, check for the free-threaded build because
135 # gc.get_referrers() ignores objects with immortal refcounts
136 # and objects are immortalized in the Python 3.13
137 # free-threading implementation at runtime.
138
139 warnings.warn(
140 "It looks like `register_random` was passed an object that could "
141 "be garbage collected immediately after `register_random` creates "
142 "a weakref to it. This will prevent Hypothesis from managing this "
143 "PRNG. See the docs for `register_random` for more details.",
144 HypothesisWarning,
145 stacklevel=2,
146 )
147
148 RANDOMS_TO_MANAGE[next(_RKEY)] = r
149
150
151def get_seeder_and_restorer(
152 seed: Hashable = 0,
153) -> tuple[Callable[[], None], Callable[[], None]]:
154 """Return a pair of functions which respectively seed all and restore
155 the state of all registered PRNGs.
156
157 This is used by the core engine via `deterministic_PRNG`, and by users
158 via `register_random`. We support registration of additional random.Random
159 instances (or other objects with seed, getstate, and setstate methods)
160 to force determinism on simulation or scheduling frameworks which avoid
161 using the global random state. See e.g. #1709.
162 """
163 assert isinstance(seed, int)
164 assert 0 <= seed < 2**32
165 states: dict[int, object] = {}
166
167 if "numpy" in sys.modules:
168 global NP_RANDOM
169 if NP_RANDOM is None:
170 # Protect this from garbage-collection by adding it to global scope
171 NP_RANDOM = RANDOMS_TO_MANAGE[next(_RKEY)] = NumpyRandomWrapper()
172
173 def seed_all() -> None:
174 assert not states
175 for k, r in RANDOMS_TO_MANAGE.items():
176 states[k] = r.getstate()
177 r.seed(seed)
178
179 def restore_all() -> None:
180 for k, state in states.items():
181 r = RANDOMS_TO_MANAGE.get(k)
182 if r is not None: # i.e., hasn't been garbage-collected
183 r.setstate(state)
184 states.clear()
185
186 return seed_all, restore_all
187
188
189@contextlib.contextmanager
190def deterministic_PRNG(seed: int = 0) -> Generator[None, None, None]:
191 """Context manager that handles random.seed without polluting global state.
192
193 See issue #1255 and PR #1295 for details and motivation - in short,
194 leaving the global pseudo-random number generator (PRNG) seeded is a very
195 bad idea in principle, and breaks all kinds of independence assumptions
196 in practice.
197 """
198 if hypothesis.core._hypothesis_global_random is None: # pragma: no cover
199 hypothesis.core._hypothesis_global_random = random.Random()
200 register_random(hypothesis.core._hypothesis_global_random)
201
202 seed_all, restore_all = get_seeder_and_restorer(seed)
203 seed_all()
204 try:
205 yield
206 finally:
207 restore_all()