Accept user-defined strategies (#2061)

Author: pradeep90

CHANGELOG.md

@@ -27,6 +27,7 @@ See [0Ver](https://0ver.org/).
 ### Features
 
 - Make `hypothesis` plugin test laws from user-defined interfaces too
+- Make `hypothesis` plugin accept user-defined strategies
 
 ### Bugfixes
 

docs/pages/contrib/hypothesis_plugins.rst

@@ -140,6 +140,27 @@ like ``Future``, ``ReaderFutureResult``, etc
 that have complex ``__init__`` signatures.
 And we don't want to mess with them.
 
+You can also register a custom strategy to be used when running your
+container's laws:
+
+.. code:: python
+
+
+  from hypothesis import strategies as st
+
+  check_all_laws(Number, container_strategy=st.builds(Number, st.integers()))
+
+The ``container_strategy`` will be used only when running the tests generated
+by the ``check_all_laws`` call above. It will have no effect on any other
+property tests that involve ``Number``. You cannot use this argument together
+with ``use_init``.
+
+Warning::
+  Avoid directly registering your container's strategy with ``hypothesis``
+  using ``st.register_type_strategy``. Because of the way we emulate
+  higher-kinded types, ``hypothesis`` may mistakenly use the strategy
+  for other incompatible containers and cause spurious test failures.
+
 Warning::
   Checking laws is not compatible with ``pytest-xdist``,
   because we use a lot of global mutable state there.

returns/contrib/hypothesis/laws.py

@@ -1,7 +1,8 @@
+import dataclasses
 import inspect
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Iterator, Mapping
 from contextlib import ExitStack, contextmanager
-from typing import Any, NamedTuple, TypeGuard, TypeVar, final
+from typing import Any, TypeGuard, TypeVar, final, overload
 
 import pytest
 from hypothesis import given
@@ -10,22 +11,57 @@
 from hypothesis.strategies._internal import types  # noqa: PLC2701
 
 from returns.contrib.hypothesis.containers import strategy_from_container
+from returns.contrib.hypothesis.type_resolver import (
+    StrategyFactory,
+    strategies_for_types,
+)
 from returns.primitives.laws import LAWS_ATTRIBUTE, Law, Lawful
 
+Example_co = TypeVar('Example_co', covariant=True)
+
 
 @final
-class _Settings(NamedTuple):
+@dataclasses.dataclass(frozen=True)
+class _Settings:
     """Settings that we provide to an end user."""
 
     settings_kwargs: dict[str, Any]
     use_init: bool
+    container_strategy: StrategyFactory | None
+
+    def __post_init__(self) -> None:
+        """Check that the settings are mutually compatible."""
+        if self.use_init and self.container_strategy is not None:
+            raise AssertionError(
+                'Expected only one of `use_init` and'
+                ' `container_strategy` to be truthy'
+            )
 
 
+@overload
 def check_all_laws(
-    container_type: type[Lawful],
+    container_type: type[Lawful[Example_co]],
+    *,
+    settings_kwargs: dict[str, Any] | None = None,
+    container_strategy: StrategyFactory[Example_co] | None = None,
+) -> None: ...
+
+
+@overload
+def check_all_laws(
+    container_type: type[Lawful[Example_co]],
     *,
     settings_kwargs: dict[str, Any] | None = None,
     use_init: bool = False,
+) -> None: ...
+
+
+def check_all_laws(
+    container_type: type[Lawful[Example_co]],
+    *,
+    settings_kwargs: dict[str, Any] | None = None,
+    use_init: bool = False,
+    container_strategy: StrategyFactory[Example_co] | None = None,
 ) -> None:
     """
     Function to check all defined mathematical laws in a specified container.
@@ -56,6 +92,7 @@ def check_all_laws(
     settings = _Settings(
         settings_kwargs or {},
         use_init,
+        container_strategy,
     )
 
     for interface, laws in container_type.laws().items():
@@ -69,62 +106,39 @@ def check_all_laws(
 
 
 @contextmanager
-def container_strategies(
+def interface_strategies(
     container_type: type[Lawful],
     *,
     settings: _Settings,
 ) -> Iterator[None]:
     """
-    Registers all types inside a container to resolve to a correct strategy.
+    Make all interfaces of a container resolve to the container's strategy.
 
     For example, let's say we have ``Result`` type.
     It is a subtype of ``ContainerN``, ``MappableN``, ``BindableN``, etc.
     When we check this type, we need ``MappableN`` to resolve to ``Result``.
 
     Can be used independently from other functions.
     """
-    our_interfaces = lawful_interfaces(container_type)
-    for interface in our_interfaces:
-        st.register_type_strategy(
-            interface,
-            strategy_from_container(
-                container_type,
-                use_init=settings.use_init,
-            ),
-        )
-
-    try:
+    mapping: Mapping[type[object], StrategyFactory] = {
+        interface: _strategy_for_container(container_type, settings)
+        for interface in lawful_interfaces(container_type)
+    }
+    with strategies_for_types(mapping):
         yield
-    finally:
-        for interface in our_interfaces:
-            types._global_type_lookup.pop(interface)  # noqa: SLF001
-        _clean_caches()
 
 
 @contextmanager
 def register_container(
     container_type: type['Lawful'],
     *,
-    use_init: bool,
+    settings: _Settings,
 ) -> Iterator[None]:
     """Temporary registers a container if it is not registered yet."""
-    used = types._global_type_lookup.pop(container_type, None)  # noqa: SLF001
-    st.register_type_strategy(
-        container_type,
-        strategy_from_container(
-            container_type,
-            use_init=use_init,
-        ),
-    )
-
-    try:
+    with strategies_for_types({
+        container_type: _strategy_for_container(container_type, settings)
+    }):
         yield
-    finally:
-        types._global_type_lookup.pop(container_type)  # noqa: SLF001
-        if used:
-            st.register_type_strategy(container_type, used)
-        else:
-            _clean_caches()
 
 
 @contextmanager
@@ -240,6 +254,17 @@ def _clean_caches() -> None:
     st.from_type.__clear_cache()  # type: ignore[attr-defined]  # noqa: SLF001
 
 
+def _strategy_for_container(
+    container_type: type[Lawful],
+    settings: _Settings,
+) -> StrategyFactory:
+    return (
+        strategy_from_container(container_type, use_init=settings.use_init)
+        if settings.container_strategy is None
+        else settings.container_strategy
+    )
+
+
 def _run_law(
     container_type: type[Lawful],
     law: Law,
@@ -252,10 +277,10 @@ def factory(source: st.DataObject) -> None:
             stack.enter_context(type_vars())
             stack.enter_context(pure_functions())
             stack.enter_context(
-                container_strategies(container_type, settings=settings),
+                interface_strategies(container_type, settings=settings),
             )
             stack.enter_context(
-                register_container(container_type, use_init=settings.use_init),
+                register_container(container_type, settings=settings),
             )
             source.draw(st.builds(law.definition))
 

returns/contrib/hypothesis/type_resolver.py

@@ -0,0 +1,75 @@
+"""Make `hypothesis` resolve types to the right strategies."""
+
+from collections.abc import Callable, Iterator, Mapping
+from contextlib import contextmanager
+from typing import TypeAlias, TypeVar
+
+from hypothesis import strategies as st
+from hypothesis.strategies._internal import types  # noqa: PLC2701
+
+Example_co = TypeVar('Example_co', covariant=True)
+
+StrategyFactory: TypeAlias = (
+    st.SearchStrategy[Example_co]
+    | Callable[[type[Example_co]], st.SearchStrategy[Example_co]]
+)
+
+
+@contextmanager
+def strategies_for_types(
+    mapping: Mapping[type[object], StrategyFactory],
+) -> Iterator[None]:
+    """
+    Temporarily register strategies with `hypothesis`.
+
+    Within this context, `hypothesis` will generate data for `MyType`
+    using `mapping[MyType]`, if available. Otherwise, it will continue to
+    use the globally registered strategy for `MyType`.
+
+    NOTE: This manually adds and removes strategies from an internal data
+    structure of `hypothesis`: `types._global_type_lookup`. This is a global
+    variable used for practically every example generated by `hypothesis`, so
+    we can easily have unintentional side-effects. We have to be very careful
+    when modifying it.
+    """
+    previous_strategies: dict[type[object], StrategyFactory | None] = {}
+    for type_, strategy in mapping.items():
+        previous_strategies[type_] = look_up_strategy(type_)
+        st.register_type_strategy(type_, strategy)
+
+    try:
+        yield
+    finally:
+        for type_, previous_strategy in previous_strategies.items():
+            if previous_strategy is None:
+                _remove_strategy(type_)
+            else:
+                st.register_type_strategy(type_, previous_strategy)
+
+
+def look_up_strategy(
+    type_: type[Example_co],
+) -> StrategyFactory[Example_co] | None:
+    """Return the strategy used by `hypothesis`."""
+    return types._global_type_lookup.get(type_)  # noqa: SLF001
+
+
+def _remove_strategy(
+    type_: type[object],
+) -> None:
+    """Remove the strategy registered for `type_`."""
+    types._global_type_lookup.pop(type_)  # noqa: SLF001
+    _clean_caches()
+
+
+def apply_strategy(
+    strategy: StrategyFactory[Example_co], type_: type[Example_co]
+) -> StrategyFactory[Example_co]:
+    """Apply `strategy` to `type_`."""
+    if isinstance(strategy, st.SearchStrategy):
+        return strategy
+    return strategy(type_)
+
+
+def _clean_caches() -> None:
+    st.from_type.__clear_cache()  # type: ignore[attr-defined]  # noqa: SLF001

tests/test_contrib/test_hypothesis/test_laws/test_user_specified_strategy.py

@@ -0,0 +1,11 @@
+from hypothesis import strategies as st
+from test_hypothesis.test_laws import test_custom_type_applicative
+
+from returns.contrib.hypothesis.laws import check_all_laws
+
+container_type = test_custom_type_applicative._Wrapper  # noqa: SLF001
+
+check_all_laws(
+    container_type,
+    container_strategy=st.builds(container_type, st.integers()),
+)