Merge pull request #3509 from HypothesisWorks/register-random

Raise if register_random is passed unreferenced object

Author: rsokl

hypothesis-python/RELEASE.rst

@@ -0,0 +1,9 @@
+RELEASE_TYPE: minor
+
+:func:`~hypothesis.register_random` has used :mod:`weakref` since :ref:`v6.27.1`,
+allowing the :class:`~random.Random`-compatible objects to be garbage-collected when
+there are no other references remaining in order to avoid memory leaks.
+We now raise an error or emit a warning when this seems likely to happen immediately.
+
+The type annotation of :func:`~hypothesis.register_random` was also widened so that
+structural subtypes of :class:`~random.Random` are accepted by static typecheckers.

hypothesis-python/docs/changes.rst

@@ -1181,7 +1181,7 @@ This patch makes the :command:`hypothesis codemod`
 -------------------
 
 This patch changes the backing datastructures of :func:`~hypothesis.register_random`
-and a few internal caches to use :class:`weakref.WeakKeyDictionary`.  This reduces
+and a few internal caches to use :class:`weakref.WeakValueDictionary`.  This reduces
 memory usage and may improve performance when registered :class:`~random.Random`
 instances are only used for a subset of your tests (:issue:`3131`).
 

hypothesis-python/src/hypothesis/internal/entropy.py

@@ -9,14 +9,33 @@
 # obtain one at https://mozilla.org/MPL/2.0/.
 
 import contextlib
+import gc
 import random
 import sys
+import warnings
 from itertools import count
-from typing import Callable, Hashable, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Hashable, Tuple
 from weakref import WeakValueDictionary
 
 import hypothesis.core
-from hypothesis.errors import InvalidArgument
+from hypothesis.errors import HypothesisWarning, InvalidArgument
+from hypothesis.internal.compat import PYPY
+
+if TYPE_CHECKING:
+    if sys.version_info >= (3, 8):  # pragma: no cover
+        from typing import Protocol
+    else:
+        from typing_extensions import Protocol
+
+    # we can't use this at runtime until from_type supports
+    # protocols -- breaks ghostwriter tests
+    class RandomLike(Protocol):
+        seed: Callable[..., Any]
+        getstate: Callable[[], Any]
+        setstate: Callable[..., Any]
+
+else:  # pragma: no cover
+    RandomLike = random.Random
 
 # This is effectively a WeakSet, which allows us to associate the saved states
 # with their respective Random instances even as new ones are registered and old
@@ -40,23 +59,89 @@ def __init__(self):
 NP_RANDOM = None
 
 
-def register_random(r: random.Random) -> None:
-    """Register the given Random instance for management by Hypothesis.
+if not PYPY:
+
+    def _get_platform_base_refcount(r: Any) -> int:
+        return sys.getrefcount(r)
+
+    # Determine the number of refcounts created by function scope for
+    # the given platform / version of Python.
+    _PLATFORM_REF_COUNT = _get_platform_base_refcount(object())
+else:  # pragma: no cover
+    # PYPY doesn't have `sys.getrefcount`
+    _PLATFORM_REF_COUNT = -1
 
-    You can pass ``random.Random`` instances (or other objects with seed,
-    getstate, and setstate methods) to ``register_random(r)`` to have their
-    states seeded and restored in the same way as the global PRNGs from the
-    ``random`` and ``numpy.random`` modules.
+
+def register_random(r: RandomLike) -> None:
+    """Register (a weakref to) the given Random-like instance for management by
+    Hypothesis.
+
+    You can pass instances of structural subtypes of ``random.Random``
+    (i.e., objects with seed, getstate, and setstate methods) to
+    ``register_random(r)`` to have their states seeded and restored in the same
+    way as the global PRNGs from the ``random`` and ``numpy.random`` modules.
 
     All global PRNGs, from e.g. simulation or scheduling frameworks, should
-    be registered to prevent flaky tests.  Hypothesis will ensure that the
-    PRNG state is consistent for all test runs, or reproducibly varied if you
+    be registered to prevent flaky tests. Hypothesis will ensure that the
+    PRNG state is consistent for all test runs, always seeding them to zero and
+    restoring the previous state after the test, or, reproducibly varied if you
     choose to use the :func:`~hypothesis.strategies.random_module` strategy.
+
+    ``register_random`` only makes `weakrefs
+    <https://docs.python.org/3/library/weakref.html#module-weakref>`_ to ``r``,
+    thus ``r`` will only be managed by Hypothesis as long as it has active
+    references elsewhere at runtime. The pattern ``register_random(MyRandom())``
+    will raise a ``ReferenceError`` to help protect users from this issue.
+    This check does not occur for the PyPy interpreter. See the following example for
+    an illustration of this issue
+
+    .. code-block:: python
+
+
+       def my_BROKEN_hook():
+           r = MyRandomLike()
+
+           # `r` will be garbage collected after the hook resolved
+           # and Hypothesis will 'forget' that it was registered
+           register_random(r)  # Hypothesis will emit a warning
+
+
+       rng = MyRandomLike()
+
+
+       def my_WORKING_hook():
+           register_random(rng)
     """
     if not (hasattr(r, "seed") and hasattr(r, "getstate") and hasattr(r, "setstate")):
         raise InvalidArgument(f"r={r!r} does not have all the required methods")
-    if r not in RANDOMS_TO_MANAGE.values():
-        RANDOMS_TO_MANAGE[next(_RKEY)] = r
+
+    if r in RANDOMS_TO_MANAGE.values():
+        return
+
+    if not PYPY:  # pragma: no branch
+        # PYPY does not have `sys.getrefcount`
+        gc.collect()
+        if not gc.get_referrers(r):
+            if sys.getrefcount(r) <= _PLATFORM_REF_COUNT:
+                raise ReferenceError(
+                    f"`register_random` was passed `r={r}` which will be "
+                    "garbage collected immediately after `register_random` creates a "
+                    "weakref to it. This will prevent Hypothesis from managing this "
+                    "source of RNG. See the docs for `register_random` for more "
+                    "details."
+                )
+            else:
+                warnings.warn(
+                    HypothesisWarning(
+                        "It looks like `register_random` was passed an object "
+                        "that could be garbage collected immediately after "
+                        "`register_random` creates a weakref to it. This will "
+                        "prevent Hypothesis from managing this source of RNG. "
+                        "See the docs for `register_random` for more details."
+                    )
+                )
+
+    RANDOMS_TO_MANAGE[next(_RKEY)] = r
 
 
 def get_seeder_and_restorer(

hypothesis-python/src/hypothesis/strategies/_internal/core.py

@@ -22,6 +22,7 @@
 from inspect import Parameter, Signature, isabstract, isclass, signature
 from types import FunctionType
 from typing import (
+    TYPE_CHECKING,
     Any,
     AnyStr,
     Callable,
@@ -124,10 +125,12 @@
     EllipsisType = type(Ellipsis)
 
 
-try:
+if sys.version_info >= (3, 8):  # pragma: no cover
     from typing import Protocol
-except ImportError:  # < py3.8
-    Protocol = object  # type: ignore[assignment]
+elif TYPE_CHECKING:
+    from typing_extensions import Protocol
+else:  # pragma: no cover
+    Protocol = object
 
 
 @cacheable

hypothesis-python/tests/cover/test_random_module.py

@@ -14,7 +14,7 @@
 import pytest
 
 from hypothesis import core, find, given, register_random, strategies as st
-from hypothesis.errors import InvalidArgument
+from hypothesis.errors import HypothesisWarning, InvalidArgument
 from hypothesis.internal import entropy
 from hypothesis.internal.compat import PYPY
 from hypothesis.internal.entropy import deterministic_PRNG
@@ -54,6 +54,9 @@ def test_cannot_register_non_Random():
         register_random("not a Random instance")
 
 
+@pytest.mark.filterwarnings(
+    "ignore:It looks like `register_random` was passed an object that could be garbage collected"
+)
 def test_registering_a_Random_is_idempotent():
     gc_on_pypy()
     n_registered = len(entropy.RANDOMS_TO_MANAGE)
@@ -144,6 +147,9 @@ def test_find_does_not_pollute_state():
         assert state_a2 != state_b2
 
 
+@pytest.mark.filterwarnings(
+    "ignore:It looks like `register_random` was passed an object that could be garbage collected"
+)
 def test_evil_prng_registration_nonsense():
     gc_on_pypy()
     n_registered = len(entropy.RANDOMS_TO_MANAGE)
@@ -172,3 +178,57 @@ def test_evil_prng_registration_nonsense():
     # Implicit check, no exception was raised in __exit__
     assert r2.getstate() == s2, "reset previously registered random state"
     assert r3.getstate() == s4, "retained state when registered within the context"
+
+
+@pytest.mark.skipif(
+    PYPY, reason="We can't guard against bad no-reference patterns in pypy."
+)
+def test_passing_unreferenced_instance_raises():
+    with pytest.raises(ReferenceError):
+        register_random(random.Random(0))
+
+
+@pytest.mark.skipif(
+    PYPY, reason="We can't guard against bad no-reference patterns in pypy."
+)
+def test_passing_unreferenced_instance_within_function_scope_raises():
+    def f():
+        register_random(random.Random(0))
+
+    with pytest.raises(ReferenceError):
+        f()
+
+
+@pytest.mark.skipif(
+    PYPY, reason="We can't guard against bad no-reference patterns in pypy."
+)
+def test_passing_referenced_instance_within_function_scope_warns():
+    def f():
+        r = random.Random(0)
+        register_random(r)
+
+    with pytest.warns(
+        HypothesisWarning,
+        match="It looks like `register_random` was passed an object that could be"
+        " garbage collected",
+    ):
+        f()
+
+
+@pytest.mark.filterwarnings(
+    "ignore:It looks like `register_random` was passed an object that could be garbage collected"
+)
+@pytest.mark.skipif(
+    PYPY, reason="We can't guard against bad no-reference patterns in pypy."
+)
+def test_register_random_within_nested_function_scope():
+    n_registered = len(entropy.RANDOMS_TO_MANAGE)
+
+    def f():
+        r = random.Random()
+        register_random(r)
+        assert len(entropy.RANDOMS_TO_MANAGE) == n_registered + 1
+
+    f()
+    gc_on_pypy()
+    assert len(entropy.RANDOMS_TO_MANAGE) == n_registered

whole-repo-tests/test_mypy.py

@@ -540,3 +540,26 @@ def test_bar(x):
     assert_mypy_errors(
         str(f.realpath()), [(5, "call-overload")], python_version=python_version
     )
+
+
+def test_register_random_interface(tmpdir):
+    f = tmpdir.join("check_mypy_on_pos_arg_only_strats.py")
+    f.write(
+        textwrap.dedent(
+            """
+            from random import Random
+            from hypothesis import register_random
+
+            class MyRandom:
+                def __init__(self) -> None:
+                    r = Random()
+                    self.seed = r.seed
+                    self.setstate = r.setstate
+                    self.getstate = r.getstate
+
+            register_random(MyRandom())
+            register_random(None)  # type: ignore[arg-type]
+            """
+        )
+    )
+    assert_mypy_errors(str(f.realpath()), [])

whole-repo-tests/test_pyright.py

@@ -195,6 +195,30 @@ def test_pyright_one_of_pos_args_only(tmp_path: Path):
     )
 
 
+def test_register_random_protocol(tmp_path: Path):
+    file = tmp_path / "test.py"
+    file.write_text(
+        textwrap.dedent(
+            """
+            from random import Random
+            from hypothesis import register_random
+
+            class MyRandom:
+                def __init__(self) -> None:
+                    r = Random()
+                    self.seed = r.seed
+                    self.setstate = r.setstate
+                    self.getstate = r.getstate
+
+            register_random(MyRandom())
+            register_random(None)  # type: ignore
+            """
+        )
+    )
+    _write_config(tmp_path, {"reportUnnecessaryTypeIgnoreComment": True})
+    assert _get_pyright_errors(file) == []
+
+
 # ---------- Helpers for running pyright ---------- #