Merge pull request #4398 from tybug/next

Skip files over 512kb and over 1024 constants for constants

Author: Zac-HD

hypothesis-python/RELEASE.rst

@@ -0,0 +1,3 @@
+RELEASE_TYPE: patch
+
+Further improve the performance of the constants-collection feature introduced in :ref:`version 6.131.1 <v6.131.1>`, by ignoring large files and files with many constants.

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

@@ -568,7 +568,7 @@ def on_observation(self, observation: TestCaseObservation) -> None:  # noqa: B02
         .. important::
 
             For |PrimitiveProvider.on_observation| to be called by Hypothesis,
-            |PrimitiveProvider.add_observability_callback| must be set to ``True``,
+            |PrimitiveProvider.add_observability_callback| must be set to ``True``.
 
             |PrimitiveProvider.on_observation| is explicitly opt-in, as enabling
             observability might increase runtime or memory usage.

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

@@ -98,12 +98,24 @@ def __eq__(self, other: object) -> bool:
         )
 
 
+class TooManyConstants(Exception):
+    # a control flow exception which we raise in ConstantsVisitor when the
+    # number of constants in a module gets too large.
+    pass
+
+
 class ConstantVisitor(NodeVisitor):
-    def __init__(self):
+    CONSTANTS_LIMIT: int = 1024
+
+    def __init__(self, *, limit: bool):
         super().__init__()
         self.constants = Constants()
+        self.limit = limit
 
     def _add_constant(self, value: object) -> None:
+        if self.limit and len(self.constants) >= self.CONSTANTS_LIMIT:
+            raise TooManyConstants
+
         if isinstance(value, str) and (
             value.isspace()
             or value == ""
@@ -166,33 +178,49 @@ def visit_Constant(self, node):
         self.generic_visit(node)
 
 
-def _constants_from_source(source: Union[str, bytes]) -> Constants:
+def _constants_from_source(source: Union[str, bytes], *, limit: bool) -> Constants:
     tree = ast.parse(source)
-    visitor = ConstantVisitor()
-    visitor.visit(tree)
+    visitor = ConstantVisitor(limit=limit)
+
+    try:
+        visitor.visit(tree)
+    except TooManyConstants:
+        # in the case of an incomplete collection, return nothing, to avoid
+        # muddying caches etc.
+        return Constants()
+
     return visitor.constants
 
 
 @lru_cache(4096)
-def constants_from_module(module: ModuleType) -> Constants:
+def constants_from_module(module: ModuleType, *, limit: bool = True) -> Constants:
     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 Constants()
 
+    if limit and len(source_bytes) > 512 * 1024:
+        # Skip files over 512kb. For reference, the largest source file
+        # in Hypothesis is strategies/_internal/core.py at 107kb at time
+        # of writing.
+        return Constants()
+
     source_hash = hashlib.sha1(source_bytes).hexdigest()[:16]
-    cache_p = storage_directory("constants") / source_hash
+    # separate cache files for each limit param. see discussion in pull/4398
+    cache_p = storage_directory("constants") / (
+        source_hash + ("" if limit else "_nolimit")
+    )
     try:
-        return _constants_from_source(cache_p.read_bytes())
+        return _constants_from_source(cache_p.read_bytes(), limit=limit)
     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:
-        constants = _constants_from_source(source_bytes)
+        constants = _constants_from_source(source_bytes, limit=limit)
     except Exception:
         # A bunch of things can go wrong here.
         # * ast.parse may fail on the source code

hypothesis-python/tests/cover/test_constants_ast.py

@@ -22,6 +22,8 @@
 from hypothesis.internal.constants_ast import (
     Constants,
     ConstantVisitor,
+    TooManyConstants,
+    _constants_from_source,
     constants_from_module,
     is_local_module_file,
 )
@@ -44,7 +46,7 @@
 
 
 def constants_from_ast(tree):
-    visitor = ConstantVisitor()
+    visitor = ConstantVisitor(limit=True)
     visitor.visit(tree)
     return visitor.constants
 
@@ -234,6 +236,40 @@ def test_ignores_ast_parse_error(tmp_path):
 @given(constants_classes)
 def test_constant_visitor_roundtrips_string(constants):
     # our files in storage_directory("constants") rely on this roundtrip
-    visitor = ConstantVisitor()
+    visitor = ConstantVisitor(limit=True)
     visitor.visit(ast.parse(str(set(constants))))
     assert visitor.constants == constants
+
+
+def test_too_many_constants():
+    visitor = ConstantVisitor(limit=True)
+    # start at n=1000 to avoid ConstantVisitor ignoring small integers
+    s = "; ".join(
+        f"n = {i}" for i in range(1000, 1000 + ConstantVisitor.CONSTANTS_LIMIT + 1)
+    )
+    # visitor should raise on too many constants
+    with pytest.raises(TooManyConstants):
+        visitor.visit(ast.parse(s))
+
+    # and also _constants_from_source should return empty on too many constants
+    assert _constants_from_source(s, limit=True) == Constants()
+
+    # but it parses fine with limit=False
+    visitor = ConstantVisitor(limit=False)
+    visitor.visit(ast.parse(s))
+    assert _constants_from_source(s, limit=False) == Constants(
+        integers=set(range(1000, 1000 + ConstantVisitor.CONSTANTS_LIMIT + 1))
+    )
+
+
+def test_module_too_large(tmp_path):
+    constant = 11231783
+
+    p = tmp_path / "large_file.py"
+    content = f"a = {constant}\n\n" + "#" * (512 * 1024 + 1)
+    p.write_text(content, encoding="utf-8")
+
+    module = ModuleType("large_module")
+    module.__file__ = str(p)
+    assert constants_from_module(module) == Constants()
+    assert constants_from_module(module, limit=False) == Constants(integers={constant})

hypothesis-python/tests/cover/test_replay_logic.py

@@ -55,10 +55,9 @@ def test(ls):
 
 def test_does_not_shrink_on_replay_with_multiple_bugs():
     database = InMemoryExampleDatabase()
-
     call_count = 0
-
-    tombstone = 1000093
+    raised = False
+    marker = 1000093
 
     @settings(
         database=database,
@@ -67,12 +66,13 @@ def test_does_not_shrink_on_replay_with_multiple_bugs():
         max_examples=1000,
     )
     @given(st.integers())
-    def test(i):
-        nonlocal call_count
+    def test(n):
+        nonlocal call_count, raised
         call_count += 1
-        if i > tombstone:
+        if n >= marker:
+            raised = True
             raise AssertionError
-        elif i == tombstone:
+        elif n < marker and raised:
             raise AssertionError
 
     with pytest.raises(ExceptionGroup):
@@ -99,10 +99,10 @@ def test_will_always_shrink_if_previous_example_does_not_replay():
         max_examples=1000,
     )
     @given(st.integers(min_value=0))
-    def test(i):
+    def test(n):
         nonlocal last
-        if i not in good:
-            last = i
+        if n not in good:
+            last = n
             raise AssertionError
 
     for i in range(20):