Merge pull request #4366 from tybug/constants-speedup

Improve constants collection speed

Author: Zac-HD

hypothesis-python/RELEASE.rst

@@ -0,0 +1,4 @@
+RELEASE_TYPE: patch
+
+This patch makes the new features introduced in :ref:`version 6.131.1 <v6.131.1>` much
+faster, and fixes an internal ``RecursionError`` when working with deeply-nested code.

hypothesis-python/src/hypothesis/_settings.py

@@ -919,7 +919,7 @@ def note_deprecation(
 settings.register_profile("ci", CI)
 
 
-if is_in_ci():
+if is_in_ci():  # pragma: no cover # covered in ci, but not locally
     settings.load_profile("ci")
 
 assert settings.default is not None

hypothesis-python/src/hypothesis/internal/conjecture/engine.py

@@ -657,7 +657,7 @@ def on_pareto_evict(self, data: ConjectureResult) -> None:
         self.settings.database.delete(self.pareto_key, choices_to_bytes(data.choices))
 
     def generate_novel_prefix(self) -> tuple[ChoiceT, ...]:
-        """Uses the tree to proactively generate a starting sequence of bytes
+        """Uses the tree to proactively generate a starting choice sequence
         that we haven't explored yet for this test.
 
         When this method is called, we assume that there must be at
@@ -1047,10 +1047,8 @@ def generate_new_examples(self) -> None:
         ran_optimisations = False
 
         while self.should_generate_more():
-            # Unfortunately generate_novel_prefix still operates in terms of
-            # a buffer and uses HypothesisProvider as its backing provider,
-            # not whatever is specified by the backend. We can improve this
-            # once more things are on the ir.
+            # we don't yet integrate DataTree with backends. Instead of generating
+            # a novel prefix, ask the backend for an input.
             if not self.using_hypothesis_backend:
                 data = self.new_conjecture_data([])
                 with suppress(BackendCannotProceed):

hypothesis-python/src/hypothesis/internal/conjecture/providers.py

@@ -13,8 +13,10 @@
 import math
 import warnings
 from collections.abc import Iterable
+from functools import cached_property
 from random import Random
 from sys import float_info
+from types import ModuleType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -45,7 +47,7 @@
     Sampler,
     many,
 )
-from hypothesis.internal.constants_ast import local_constants
+from hypothesis.internal.constants_ast import constants_from_module, local_modules
 from hypothesis.internal.floats import (
     SIGNALING_NAN,
     float_to_int,
@@ -197,21 +199,40 @@
 }
 
 
-_local_constants_hash: Optional[int] = None
-
-
-def _get_local_constants():
-    global _local_constants_hash
-
-    constants = local_constants()
-    constants_hash = hash(tuple((k, tuple(v)) for k, v in constants.items()))
-    # if we've added new constants since the last time we checked, invalidate
-    # the cache.
-    if constants_hash != _local_constants_hash:
-        CONSTANTS_CACHE.cache.clear()
-        _local_constants_hash = constants_hash
+_local_constants: "ConstantsT" = {
+    "integer": SortedSet(),
+    "float": SortedSet(key=float_to_int),
+    "bytes": SortedSet(),
+    "string": SortedSet(),
+}
+# modules that we've already seen and processed for local constants.
+_local_modules: set[ModuleType] = set()
+
+
+def _get_local_constants() -> "ConstantsT":
+    new_constants: set[ConstantT] = set()
+    new_modules = list(local_modules() - _local_modules)
+    for new_module in new_modules:
+        new_constants |= constants_from_module(new_module)
+
+    for constant in new_constants:
+        choice_type = {
+            int: "integer",
+            float: "float",
+            bytes: "bytes",
+            str: "string",
+        }[type(constant)]
+        # if we add any new constant, invalidate the constant cache for permitted values.
+        # A more efficient approach would be invalidating just the keys with this
+        # choice_type.
+        if (
+            constant not in _local_constants[choice_type]  # type: ignore
+        ):  # pragma: no branch
+            CONSTANTS_CACHE.cache.clear()
+        _local_constants[choice_type].add(constant)  # type: ignore
 
-    return constants
+    _local_modules.update(new_modules)
+    return _local_constants
 
 
 class _BackendInfoMsg(TypedDict):
@@ -371,9 +392,13 @@ class HypothesisProvider(PrimitiveProvider):
 
     def __init__(self, conjecturedata: Optional["ConjectureData"], /):
         super().__init__(conjecturedata)
-        self.local_constants = _get_local_constants()
         self._random = None if self._cd is None else self._cd._random
 
+    @cached_property
+    def _local_constants(self):
+        # defer computation of local constants until/if we need it
+        return _get_local_constants()
+
     def _maybe_draw_constant(
         self,
         choice_type: ChoiceTypeT,
@@ -382,6 +407,7 @@ def _maybe_draw_constant(
         p: float = 0.05,
     ) -> Optional["ConstantT"]:
         assert self._random is not None
+        assert self._local_constants is not None
         assert choice_type != "boolean"
 
         # check whether we even want a constant before spending time computing
@@ -399,7 +425,7 @@ def _maybe_draw_constant(
                 ),
                 tuple(
                     choice
-                    for choice in self.local_constants[choice_type]
+                    for choice in self._local_constants[choice_type]
                     if choice_permitted(choice, constraints)
                 ),
             )

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

@@ -9,19 +9,19 @@
 # obtain one at https://mozilla.org/MPL/2.0/.
 
 import ast
+import hashlib
 import inspect
 import math
 import sys
-from ast import AST, Constant, Expr, NodeVisitor, UnaryOp, USub
+from ast import Constant, Expr, NodeVisitor, UnaryOp, USub
 from functools import lru_cache
 from pathlib import Path
 from types import ModuleType
-from typing import TYPE_CHECKING, AbstractSet, Optional, TypedDict, Union
-
-from sortedcontainers import SortedSet
+from typing import TYPE_CHECKING, AbstractSet, TypedDict, Union
 
+import hypothesis
+from hypothesis.configuration import storage_directory
 from hypothesis.internal.escalation import is_hypothesis_file
-from hypothesis.internal.floats import float_to_int
 
 if TYPE_CHECKING:
     from typing import TypeAlias
@@ -96,22 +96,55 @@ def visit_Constant(self, node):
         self.generic_visit(node)
 
 
-@lru_cache(1024)
-def constants_from_ast(tree: AST) -> AbstractSet[ConstantT]:
+def _constants_from_source(source: Union[str, bytes]) -> AbstractSet[ConstantT]:
+    tree = ast.parse(source)
     visitor = ConstantVisitor()
     visitor.visit(tree)
     return visitor.constants
 
 
-@lru_cache(1024)
-def _module_ast(module: ModuleType) -> Optional[AST]:
+@lru_cache(4096)
+def constants_from_module(module: ModuleType) -> AbstractSet[ConstantT]:
+    try:
+        module_file = inspect.getsourcefile(module)
+        # use type: ignore because we know this might error
+        source_bytes = Path(module_file).read_bytes()  # type: ignore
+    except Exception:
+        return set()
+
+    source_hash = hashlib.sha1(source_bytes).hexdigest()[:16]
+    cache_p = storage_directory("constants") / source_hash
+    try:
+        return _constants_from_source(cache_p.read_bytes())
+    except Exception:
+        # if the cached location doesn't exist, or it does exist but there was
+        # a problem reading it, fall back to standard computation of the constants
+        pass
+
     try:
-        source = inspect.getsource(module)
-        tree = ast.parse(source)
+        constants = _constants_from_source(source_bytes)
     except Exception:
-        return None
+        # A bunch of things can go wrong here.
+        # * ast.parse may fail on the source code
+        # * NodeVisitor may hit a RecursionError (see many related issues on
+        #   e.g. libcst https://github.com/Instagram/LibCST/issues?q=recursion),
+        #   or a MemoryError (`"[1, " * 200 + "]" * 200`)
+        return set()
 
-    return tree
+    try:
+        cache_p.parent.mkdir(parents=True, exist_ok=True)
+        cache_p.write_text(
+            f"# file: {module_file}\n# hypothesis_version: {hypothesis.__version__}\n\n"
+            # somewhat arbitrary sort order. The cache file doesn't *have* to be
+            # stable... but it is aesthetically pleasing, and means we could rely
+            # on it in the future!
+            + str(sorted(constants, key=lambda v: (str(type(v)), v))),
+            encoding="utf-8",
+        )
+    except Exception:  # pragma: no cover
+        pass
+
+    return constants
 
 
 @lru_cache(4096)
@@ -141,7 +174,7 @@ def _is_local_module_file(path: str) -> bool:
     )
 
 
-def local_modules() -> tuple[ModuleType, ...]:
+def local_modules() -> set[ModuleType]:
     if sys.platform == "emscripten":  # pragma: no cover
         # pyodide builds bundle the stdlib in a nonstandard location, like
         # `/lib/python312.zip/heapq.py`. To avoid identifying the entirety of
@@ -151,44 +184,15 @@ def local_modules() -> tuple[ModuleType, ...]:
         # pyodide may provide some way to distinguish stdlib/third-party/local
         # code. I haven't looked into it. If they do, we should correctly implement
         # ModuleLocation for pyodide instead of this.
-        return ()
-
-    # Prevents a `RuntimeError` that can occur when looping over `sys.modules`
-    # if it's simultaneously modified as a side effect of code in another thread.
-    # See: https://docs.python.org/3/library/sys.html#sys.modules
-    modules = sys.modules.copy().values()
+        return set()
 
-    return tuple(
+    return {
         module
-        for module in modules
+        # copy to avoid a RuntimeError if another thread imports a module while
+        # we're iterating.
+        for module in sys.modules.copy().values()
         if (
             getattr(module, "__file__", None) is not None
             and _is_local_module_file(module.__file__)
         )
-    )
-
-
-def local_constants() -> ConstantsT:
-    constants: set[ConstantT] = set()
-    for module in local_modules():
-        tree = _module_ast(module)
-        if tree is None:  # pragma: no cover
-            continue
-        constants |= constants_from_ast(tree)
-
-    local_constants: ConstantsT = {
-        "integer": SortedSet(),
-        "float": SortedSet(key=float_to_int),
-        "bytes": SortedSet(),
-        "string": SortedSet(),
     }
-    for value in constants:
-        choice_type = {
-            int: "integer",
-            float: "float",
-            bytes: "bytes",
-            str: "string",
-        }[type(value)]
-        local_constants[choice_type].add(value)  # type: ignore # hard to type
-
-    return local_constants

hypothesis-python/tests/conjecture/test_local_constants.py

@@ -12,9 +12,11 @@
 
 import pytest
 
-from hypothesis import settings, strategies as st
+from hypothesis import given, settings, strategies as st
+from hypothesis.internal import constants_ast
 from hypothesis.internal.conjecture import providers
 from hypothesis.internal.conjecture.choice import choice_equal
+from hypothesis.internal.conjecture.providers import CONSTANTS_CACHE
 
 from tests.common.debug import find_any
 from tests.common.utils import Why, xfail_on_crosshair
@@ -24,30 +26,51 @@
 # with CONSTANTS_CACHE when testing it inside of a hypothesis test.
 @pytest.mark.parametrize("value", [2**20 - 50, 2**10 - 10, 129387123, -19827321, 0])
 def test_can_draw_local_constants_integers(monkeypatch, value):
-    monkeypatch.setattr(providers, "local_constants", lambda: {"integer": {value}})
+    # _get_local_constants normally invalidates this cache for us, but we're
+    # monkeypatching it.
+    CONSTANTS_CACHE.cache.clear()
+    monkeypatch.setattr(providers, "_get_local_constants", lambda: {"integer": {value}})
     find_any(st.integers(), lambda v: choice_equal(v, value))
 
 
 @xfail_on_crosshair(Why.undiscovered)  # I think float_to_int is difficult for crosshair
 @pytest.mark.parametrize("value", [1.2938, -1823.0239, 1e999, math.nan])
 def test_can_draw_local_constants_floats(monkeypatch, value):
-    monkeypatch.setattr(providers, "local_constants", lambda: {"float": {value}})
+    CONSTANTS_CACHE.cache.clear()
+    monkeypatch.setattr(providers, "_get_local_constants", lambda: {"float": {value}})
     find_any(st.floats(), lambda v: choice_equal(v, value))
 
 
 @pytest.mark.parametrize("value", [b"abdefgh", b"a" * 50])
 def test_can_draw_local_constants_bytes(monkeypatch, value):
-    monkeypatch.setattr(providers, "local_constants", lambda: {"bytes": {value}})
+    CONSTANTS_CACHE.cache.clear()
+    monkeypatch.setattr(providers, "_get_local_constants", lambda: {"bytes": {value}})
     find_any(st.binary(), lambda v: choice_equal(v, value))
 
 
 @pytest.mark.parametrize("value", ["abdefgh", "a" * 50])
 def test_can_draw_local_constants_string(monkeypatch, value):
-    monkeypatch.setattr(providers, "local_constants", lambda: {"string": {value}})
+    CONSTANTS_CACHE.cache.clear()
+    monkeypatch.setattr(providers, "_get_local_constants", lambda: {"string": {value}})
     # we have a bunch of strings in GLOBAL_CONSTANTS, so it might take a while
     # to generate our local constant.
     find_any(
         st.text(),
         lambda v: choice_equal(v, value),
         settings=settings(max_examples=5_000),
     )
+
+
+def test_actual_collection(monkeypatch):
+    # covering test for doing some real work collecting constants. We'll fake
+    # hypothesis as being the "local" module, just to get some real constant
+    # collection going.
+    monkeypatch.setattr(
+        constants_ast, "_is_local_module_file", lambda f: "hypothesis" in f
+    )
+
+    @given(st.integers())
+    def f(n):
+        pass
+
+    f()