feat: Add context argument to hook functions

README.rst

@@ -405,25 +405,29 @@ If you want to customize how data is generated, then you can use hooks of three
 - Schema-local, which are applied only for specific schema instance;
 - Test function specific, they are applied only for a specific test function;
 
-Each hook accepts a Hypothesis strategy and should return a Hypothesis strategy:
+Each hook accepts a Hypothesis strategy and a hook context. Hook context provides additional info that might be helpful to
+construct a new strategy, for example ``context.endpoint`` attribute is a reference to the currently tested endpoint.
+For more information look at ``schemathesis.hooks.HookContext`` class.
+
+Hooks should return a Hypothesis strategy:
 
 .. code:: python
 
     import schemathesis
 
-    def global_hook(strategy):
+    def global_hook(strategy, context):
         return strategy.filter(lambda x: x["id"].isdigit())
 
     schemathesis.hooks.register("query", hook)
 
     schema = schemathesis.from_uri("http://0.0.0.0:8080/swagger.json")
 
-    def schema_hook(strategy):
+    def schema_hook(strategy, context):
         return strategy.filter(lambda x: int(x["id"]) % 2 == 0)
 
     schema.register_hook("query", schema_hook)
 
-    def function_hook(strategy):
+    def function_hook(strategy, context):
         return strategy.filter(lambda x: len(x["id"]) > 5)
 
     @schema.with_hook("query", function_hook)

docs/changelog.rst

@@ -6,6 +6,17 @@ Changelog
 `Unreleased`_
 -------------
 
+Added
+~~~~~
+
+- ``context`` argument for hook functions to provide an additional context for hooks. A deprecation warning is emitted
+  for hook functions that do not accept this argument.
+
+Deprecated
+~~~~~~~~~~
+
+- Hook functions that do not accept ``context`` as their second argument. They will become not supported in Schemathesis 2.0.
+
 Fixed
 ~~~~~
 

src/schemathesis/_hypothesis.py

@@ -1,9 +1,10 @@
 """Provide strategies for given endpoint(s) definition."""
 import asyncio
+import inspect
 import re
 from base64 import b64encode
 from functools import partial
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Callable, Dict, Optional, Tuple, Union
 from urllib.parse import quote_plus
 
 import hypothesis
@@ -13,7 +14,7 @@
 from . import utils
 from ._compat import handle_warnings
 from .exceptions import InvalidSchema
-from .hooks import get_hook
+from .hooks import HookContext, get_hook
 from .models import Case, Endpoint
 from .types import Hook
 
@@ -187,18 +188,31 @@ def _get_case_strategy(
             raise InvalidSchema("Body parameters are defined for GET request.")
         static_parameters["body"] = None
         strategies.pop("body", None)
-    _apply_hooks(strategies, get_hook)
-    _apply_hooks(strategies, endpoint.schema.get_hook)
+    context = HookContext(endpoint)
+    _apply_hooks(strategies, get_hook, context)
+    _apply_hooks(strategies, endpoint.schema.get_hook, context)
     if hooks is not None:
-        _apply_hooks(strategies, hooks.get)
+        _apply_hooks(strategies, hooks.get, context)
     return st.builds(partial(Case, **static_parameters), **strategies)
 
 
-def _apply_hooks(strategies: Dict[str, st.SearchStrategy], getter: Callable[[str], Optional[Hook]]) -> None:
+def _apply_hooks(
+    strategies: Dict[str, st.SearchStrategy], getter: Callable[[str], Optional[Hook]], context: HookContext
+) -> None:
     for key, strategy in strategies.items():
         hook = getter(key)
         if hook is not None:
-            strategies[key] = hook(strategy)
+            args: Union[Tuple[st.SearchStrategy], Tuple[st.SearchStrategy, HookContext]]
+            if _accepts_context(hook):
+                args = (strategy, context)
+            else:
+                args = (strategy,)
+            strategies[key] = hook(*args)
+
+
+def _accepts_context(hook: Hook) -> bool:
+    # There are no restrictions on the second argument's name and we don't check its name here.
+    return len(inspect.signature(hook).parameters) == 2
 
 
 def register_string_format(name: str, strategy: st.SearchStrategy) -> None:

src/schemathesis/hooks.py

@@ -1,12 +1,35 @@
+import inspect
+import warnings
 from typing import Optional
 
+import attr
+
 from .constants import HookLocation
+from .models import Endpoint
 from .types import Hook
 
 GLOBAL_HOOKS = {}
 
 
+@attr.s(slots=True)
+class HookContext:
+    """A context that is passed to hook functions."""
+
+    endpoint: Endpoint = attr.ib()
+
+
+def warn_deprecated_hook(hook: Hook) -> None:
+    if "context" not in inspect.signature(hook).parameters:
+        warnings.warn(
+            DeprecationWarning(
+                "Hook functions that do not accept `context` argument are deprecated and "
+                "support will be removed in Schemathesis 2.0."
+            )
+        )
+
+
 def register(place: str, hook: Hook) -> None:
+    warn_deprecated_hook(hook)
     key = HookLocation[place]
     GLOBAL_HOOKS[key] = hook
 

src/schemathesis/schemas.py

@@ -26,6 +26,7 @@
 from .converter import to_json_schema, to_json_schema_recursive
 from .exceptions import InvalidSchema
 from .filters import should_skip_by_tag, should_skip_endpoint, should_skip_method
+from .hooks import warn_deprecated_hook
 from .models import Endpoint, empty_object
 from .types import Filter, Hook, NotSet
 from .utils import NOT_SET, GenericResponse, StringDatesYAMLLoader
@@ -164,11 +165,13 @@ def _get_response_schema(self, definition: Dict[str, Any]) -> Optional[Dict[str,
         raise NotImplementedError
 
     def register_hook(self, place: str, hook: Hook) -> None:
+        warn_deprecated_hook(hook)
         key = HookLocation[place]
         self.hooks[key] = hook
 
     def with_hook(self, place: str, hook: Hook) -> Callable[[GenericTest], GenericTest]:
         """Register a hook for a specific test."""
+        warn_deprecated_hook(hook)
         if place not in HookLocation.__members__:
             raise KeyError(place)
 

src/schemathesis/types.py

@@ -1,8 +1,11 @@
 from pathlib import Path
-from typing import Any, Callable, Dict, List, NewType, Set, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, NewType, Set, Tuple, Union
 
 from hypothesis.strategies import SearchStrategy
 
+if TYPE_CHECKING:
+    from .hooks import HookContext
+
 Schema = NewType("Schema", Dict[str, Any])  # pragma: no mutate
 PathLike = Union[Path, str]  # pragma: no mutate
 
@@ -22,6 +25,8 @@ class NotSet:
 # A filter for endpoint / method
 Filter = Union[str, List[str], Tuple[str], Set[str], NotSet]  # pragma: no mutate
 
-Hook = Callable[[SearchStrategy], SearchStrategy]  # pragma: no mutate
+Hook = Union[
+    Callable[[SearchStrategy], SearchStrategy], Callable[[SearchStrategy, "HookContext"], SearchStrategy]
+]  # pragma: no mutate
 
 RawAuth = Tuple[str, str]  # pragma: no mutate

test/test_hooks.py

@@ -4,7 +4,7 @@
 import schemathesis
 
 
-def hook(strategy):
+def hook(strategy, context):
     return strategy.filter(lambda x: x["id"].isdigit())
 
 
@@ -52,7 +52,8 @@ def test(case):
 @pytest.mark.usefixtures("query_hook")
 @pytest.mark.endpoints("custom_format")
 def test_hooks_combination(schema, schema_url):
-    def extra(st):
+    def extra(st, context):
+        assert context.endpoint == schema.endpoints["/api/custom_format"]["GET"]
         return st.filter(lambda x: int(x["id"]) % 2 == 0)
 
     schema.register_hook("query", extra)
@@ -89,7 +90,7 @@ def test_per_test_hooks(testdir):
         """
 from hypothesis import strategies as st
 
-def replacement(strategy):
+def replacement(strategy, context):
     return st.just({"id": "foobar"})
 
 @schema.with_hook("query", replacement)
@@ -104,10 +105,10 @@ def test_a(case):
 def test_b(case):
     assert case.query["id"] == "foobar"
 
-def another_replacement(strategy):
+def another_replacement(strategy, context):
     return st.just({"id": "foobaz"})
 
-def third_replacement(strategy):
+def third_replacement(strategy, context):
     return st.just({"value": "spam"})
 
 @schema.parametrize()
@@ -131,7 +132,7 @@ def test_d(case):
 
 
 def test_invalid_hook(schema):
-    def foo(strategy):
+    def foo(strategy, context):
         pass
 
     with pytest.raises(KeyError, match="wrong"):
@@ -144,7 +145,7 @@ def test(case):
 def test_hooks_via_parametrize(testdir):
     testdir.make_test(
         """
-def extra(st):
+def extra(st, context):
     return st.filter(lambda x: x["id"].isdigit() and int(x["id"]) % 2 == 0)
 
 schema.register_hook("query", extra)
@@ -159,3 +160,26 @@ def test(case):
     )
     result = testdir.runpytest()
     result.assert_outcomes(passed=1)
+
+
+@pytest.mark.hypothesis_nested
+@pytest.mark.endpoints("custom_format")
+def test_deprecated_hook(recwarn, schema):
+    def deprecated_hook(strategy):
+        return strategy.filter(lambda x: x["id"].isdigit())
+
+    schema.register_hook("query", deprecated_hook)
+    assert (
+        str(recwarn.list[0].message) == "Hook functions that do not accept `context` argument are deprecated and "
+        "support will be removed in Schemathesis 2.0."
+    )
+
+    strategy = schema.endpoints["/api/custom_format"]["GET"].as_strategy()
+
+    @given(case=strategy)
+    @settings(max_examples=3)
+    def test(case):
+        assert case.query["id"].isdigit()
+        assert int(case.query["id"]) % 2 == 0
+
+    test()