Merge pull request #4416 from tybug/provider-on-observability

Add `PrimitiveProvider.on_observation`

Author: Liam-DeVoe

hypothesis-python/RELEASE.rst

@@ -0,0 +1,3 @@
+RELEASE_TYPE: minor
+
+Add |PrimitiveProvider.on_observation| to the internal :ref:`alternative backends <alternative-backends-internals>` interface.

hypothesis-python/docs/prolog.rst

@@ -115,6 +115,9 @@
 .. |PrimitiveProvider.draw_float| replace:: :func:`~hypothesis.internal.conjecture.providers.PrimitiveProvider.draw_float`
 .. |PrimitiveProvider.draw_string| replace:: :func:`~hypothesis.internal.conjecture.providers.PrimitiveProvider.draw_string`
 .. |PrimitiveProvider.draw_bytes| replace:: :func:`~hypothesis.internal.conjecture.providers.PrimitiveProvider.draw_bytes`
+.. |PrimitiveProvider.on_observation| replace:: :func:`~hypothesis.internal.conjecture.providers.PrimitiveProvider.on_observation`
+.. |PrimitiveProvider.per_test_case_context_manager| replace:: :func:`~hypothesis.internal.conjecture.providers.PrimitiveProvider.per_test_case_context_manager`
+.. |PrimitiveProvider.add_observability_callback| replace:: :data:`~hypothesis.internal.conjecture.providers.PrimitiveProvider.add_observability_callback`
 
 .. |AVAILABLE_PROVIDERS| replace:: :data:`~hypothesis.internal.conjecture.providers.AVAILABLE_PROVIDERS`
 .. |TESTCASE_CALLBACKS| replace:: :data:`~hypothesis.internal.observability.TESTCASE_CALLBACKS`

hypothesis-python/src/hypothesis/core.py

@@ -68,6 +68,7 @@
     Unsatisfiable,
     UnsatisfiedAssumption,
 )
+from hypothesis.internal import observability
 from hypothesis.internal.compat import (
     PYPY,
     BaseExceptionGroup,
@@ -99,7 +100,6 @@
 )
 from hypothesis.internal.healthcheck import fail_health_check
 from hypothesis.internal.observability import (
-    OBSERVABILITY_COLLECT_COVERAGE,
     TESTCASE_CALLBACKS,
     InfoObservation,
     InfoObservationType,
@@ -936,7 +936,9 @@ def test_identifier(self):
         ) or get_pretty_function_description(self.wrapped_test)
 
     def _should_trace(self):
-        _trace_obs = TESTCASE_CALLBACKS and OBSERVABILITY_COLLECT_COVERAGE
+        # NOTE: we explicitly support monkeypatching this. Keep the namespace
+        # access intact.
+        _trace_obs = TESTCASE_CALLBACKS and observability.OBSERVABILITY_COLLECT_COVERAGE
         _trace_failure = (
             self.failed_normally
             and not self.failed_due_to_deadline

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

@@ -15,7 +15,7 @@
 import time
 from collections import defaultdict
 from collections.abc import Generator, Sequence
-from contextlib import contextmanager, suppress
+from contextlib import AbstractContextManager, contextmanager, nullcontext, suppress
 from dataclasses import dataclass, field
 from datetime import timedelta
 from enum import Enum
@@ -68,6 +68,7 @@
 from hypothesis.internal.conjecture.shrinker import Shrinker, ShrinkPredicateT, sort_key
 from hypothesis.internal.escalation import InterestingOrigin
 from hypothesis.internal.healthcheck import fail_health_check
+from hypothesis.internal.observability import Observation, with_observation_callback
 from hypothesis.reporting import base_report, report
 
 #: The maximum number of times the shrinker will reduce the complexity of a failing
@@ -830,8 +831,31 @@ def debug_data(self, data: Union[ConjectureData, ConjectureResult]) -> None:
             f"{', ' + data.output if data.output else ''}"
         )
 
+    def observe_for_provider(self) -> AbstractContextManager:
+        def on_observation(observation: Observation) -> None:
+            assert observation.type == "test_case"
+            # because lifetime == "test_function"
+            assert isinstance(self.provider, PrimitiveProvider)
+            # only fire if we actually used that provider to generate this observation
+            if not self._switch_to_hypothesis_provider:
+                self.provider.on_observation(observation)
+
+        if (
+            self.settings.backend != "hypothesis"
+            # only for lifetime = "test_function" providers (guaranteed
+            # by this isinstance check)
+            and isinstance(self.provider, PrimitiveProvider)
+            # and the provider opted-in to observations
+            and self.provider.add_observability_callback
+        ):
+            return with_observation_callback(on_observation)
+        return nullcontext()
+
     def run(self) -> None:
-        with local_settings(self.settings):
+        with (
+            local_settings(self.settings),
+            self.observe_for_provider(),
+        ):
             try:
                 self._run()
             except RunIsComplete:

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

@@ -21,6 +21,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    ClassVar,
     Literal,
     Optional,
     TypedDict,
@@ -62,7 +63,7 @@
     next_up,
 )
 from hypothesis.internal.intervalsets import IntervalSet
-from hypothesis.internal.observability import InfoObservationType
+from hypothesis.internal.observability import InfoObservationType, TestCaseObservation
 
 if TYPE_CHECKING:
     from typing import TypeAlias
@@ -356,6 +357,15 @@ class PrimitiveProvider(abc.ABC):
     #: Only set this to ``True`` if it is necessary for your backend.
     avoid_realization = False
 
+    #: If ``True``, |PrimitiveProvider.on_observation| will be added as a
+    #: callback to |TESTCASE_CALLBACKS|, enabling observability during the lifetime
+    #: of this provider. If ``False``, |PrimitiveProvider.on_observation| will
+    #: never be called by Hypothesis.
+    #:
+    #: The opt-in behavior of observability is because enabling observability
+    #: might increase runtime or memory usage.
+    add_observability_callback: ClassVar[bool] = False
+
     def __init__(self, conjecturedata: Optional["ConjectureData"], /) -> None:
         self._cd = conjecturedata
 
@@ -544,6 +554,41 @@ def observe_information_messages(
         assert lifetime in ("test_case", "test_function")
         yield from []
 
+    def on_observation(self, observation: TestCaseObservation) -> None:  # noqa: B027
+        """
+        Called at the end of each test case which uses this provider, with the same
+        ``observation["type"] == "test_case"`` observation that is passed to
+        other callbacks in |TESTCASE_CALLBACKS|. This method is not called with
+        ``observation["type"] in {"info", "alert", "error"}`` observations.
+
+        .. important::
+
+            For |PrimitiveProvider.on_observation| to be called by Hypothesis,
+            |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.
+
+        Calls to this method are guaranteed to alternate with calls to
+        |PrimitiveProvider.per_test_case_context_manager|. For example:
+
+        .. code-block:: python
+
+            # test function starts
+            per_test_case_context_manager()
+            on_observation()
+            per_test_case_context_manager()
+            on_observation()
+            ...
+            # test function ends
+
+        Note that |PrimitiveProvider.on_observation| will not be called for test
+        cases which did not use this provider during generation, for example
+        during |Phase.reuse| or |Phase.shrink|, or because Hypothesis switched
+        to the standard Hypothesis backend after this backend raised too many
+        |BackendCannotProceed| exceptions.
+        """
+
     def span_start(self, label: int, /) -> None:  # noqa: B027  # non-abstract noop
         """Marks the beginning of a semantically meaningful span of choices.
 

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

@@ -15,6 +15,8 @@
 import sys
 import time
 import warnings
+from collections.abc import Generator
+from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import date, timedelta
 from functools import lru_cache
@@ -98,6 +100,17 @@ class TestCaseObservation(BaseObservation):
 TESTCASE_CALLBACKS: list[Callable[[Observation], None]] = []
 
 
+@contextmanager
+def with_observation_callback(
+    callback: Callable[[Observation], None],
+) -> Generator[None, None, None]:
+    TESTCASE_CALLBACKS.append(callback)
+    try:
+        yield
+    finally:
+        TESTCASE_CALLBACKS.remove(callback)
+
+
 def deliver_observation(observation: Observation) -> None:
     for callback in TESTCASE_CALLBACKS:
         callback(observation)

hypothesis-python/tests/conjecture/test_alt_backend.py

@@ -46,6 +46,7 @@
 )
 from hypothesis.internal.floats import SIGNALING_NAN
 from hypothesis.internal.intervalsets import IntervalSet
+from hypothesis.internal.observability import TESTCASE_CALLBACKS, Observation
 
 from tests.common.debug import minimal
 from tests.common.utils import (
@@ -718,3 +719,43 @@ def test_replay_choices():
     # trivial covering test
     provider = TrivialProvider(None)
     provider.replay_choices([1])
+
+
+class ObservationProvider(TrivialProvider):
+    add_observability_callback = True
+
+    def __init__(self, conjecturedata: "ConjectureData", /) -> None:
+        super().__init__(conjecturedata)
+        # calls to per_test_case_context_manager and on_observation alternate,
+        # starting with per_test_case_context_manager
+        self.expected = "per_test_case_context_manager"
+
+    @contextmanager
+    def per_test_case_context_manager(self):
+        assert self.expected == "per_test_case_context_manager"
+        self.expected = "on_observation"
+        yield
+
+    def on_observation(self, observation: Observation) -> None:
+        assert self.expected == "on_observation"
+        self.expected = "per_test_case_context_manager"
+
+
+@temp_register_backend("observation", ObservationProvider)
+def test_on_observation_alternates():
+    @given(st.integers())
+    @settings(backend="observation")
+    def f(n):
+        pass
+
+    f()
+
+
+@temp_register_backend("observation", TrivialProvider)
+def test_on_observation_no_override():
+    @given(st.integers())
+    @settings(backend="observation")
+    def f(n):
+        assert TESTCASE_CALLBACKS == []
+
+    f()