Add reusable type for @composite draw argument (#3069)

* Add reusable type for `@composite` draw argument

* Make DrawFn a protocol on >= 3.8

* Change DrawFn `label` parameter type to object

* Add test case for DrawFn type hints

* Prevent `DrawFn` instantiation on < 3.8

Co-authored-by: Zac Hatfield-Dodds <[email protected]>

* Condense `DrawFn` docs

Co-authored-by: Zac Hatfield-Dodds <[email protected]>

* Fix formatting

* Add test for `DrawFn`

* Update stateful docs

Fixes #3071.

Co-authored-by: Zac Hatfield-Dodds <[email protected]>

Author: ROpdebee

hypothesis-python/RELEASE.rst

@@ -0,0 +1,7 @@
+RELEASE_TYPE: minor
+
+This release adds the :class:`~hypothesis.strategies.DrawFn` type as a reusable
+type hint for the ``draw`` argument of
+:func:`@composite <hypothesis.strategies.composite>` functions.
+
+Thanks to Ruben Opdebeeck for this contribution!

hypothesis-python/src/hypothesis/stateful.py

@@ -518,8 +518,8 @@ def _convert_targets(targets, target):
 
 
 def rule(*, targets=(), target=None, **kwargs):
-    """Decorator for RuleBasedStateMachine. Any name present in target or
-    targets will define where the end result of this function should go. If
+    """Decorator for RuleBasedStateMachine. Any Bundle present in ``target`` or
+    ``targets`` will define where the end result of this function should go. If
     both are empty then the end result will be discarded.
 
     ``target`` must be a Bundle, or if the result should go to multiple

hypothesis-python/src/hypothesis/strategies/__init__.py

@@ -17,6 +17,7 @@
 from hypothesis.strategies._internal.collections import tuples
 from hypothesis.strategies._internal.core import (
     DataObject,
+    DrawFn,
     binary,
     booleans,
     builds,
@@ -82,6 +83,7 @@
     "decimals",
     "deferred",
     "dictionaries",
+    "DrawFn",
     "emails",
     "fixed_dictionaries",
     "floats",
@@ -126,6 +128,7 @@ def _check_exports(_public):
     # @declares_strategy.
     exported_strategies = set(__all__) - {
         "DataObject",
+        "DrawFn",
         "SearchStrategy",
         "composite",
         "register_type_strategy",

hypothesis-python/src/hypothesis/strategies/_internal/core.py

@@ -24,7 +24,7 @@
 from decimal import Context, Decimal, localcontext
 from fractions import Fraction
 from functools import reduce
-from inspect import Parameter, getfullargspec, isabstract, isclass, signature
+from inspect import Parameter, Signature, getfullargspec, isabstract, isclass, signature
 from types import FunctionType
 from typing import (
     Any,
@@ -111,6 +111,11 @@
 from hypothesis.strategies._internal.utils import cacheable, defines_strategy
 from hypothesis.utils.conventions import InferType, infer, not_set
 
+try:
+    from typing import Protocol
+except ImportError:  # < py3.8
+    Protocol = object  # type: ignore[assignment]
+
 UniqueBy = Union[Callable[[Ex], Hashable], Tuple[Callable[[Ex], Hashable], ...]]
 
 
@@ -1404,6 +1409,33 @@ def calc_label(self):
         return calc_label_from_cls(self.definition)
 
 
+class DrawFn(Protocol):
+    """This type only exists so that you can write type hints for functions
+    decorated with :func:`@composite <hypothesis.strategies.composite>`.
+
+    .. code-block:: python
+
+        @composite
+        def list_and_index(draw: DrawFn) -> Tuple[int, str]:
+            i = draw(integers())  # type inferred as 'int'
+            s = draw(text())  # type inferred as 'str'
+
+    """
+
+    def __init__(self):
+        raise TypeError("Protocols cannot be instantiated")  # pragma: no cover
+
+    # On Python 3.8+, Protocol overrides our signature for __init__,
+    # so we override it right back to make the docs look nice.
+    __signature__: Signature = Signature(parameters=[])
+
+    # We define this as a callback protocol because a simple typing.Callable is
+    # insufficient to fully represent the interface, due to the optional `label`
+    # parameter.
+    def __call__(self, strategy: SearchStrategy[Ex], label: object = None) -> Ex:
+        raise NotImplementedError
+
+
 @cacheable
 def composite(f: Callable[..., Ex]) -> Callable[..., SearchStrategy[Ex]]:
     """Defines a strategy that is built out of potentially arbitrarily many

hypothesis-python/tests/cover/test_composite.py

@@ -176,3 +176,8 @@ def test_applying_composite_decorator_to_methods(data):
         x = data.draw(strategy)
         assert isinstance(x, int)
         assert 0 <= x <= 10
+
+
+def test_drawfn_cannot_be_instantiated():
+    with pytest.raises(TypeError):
+        st.DrawFn()

whole-repo-tests/test_type_hints.py

@@ -109,6 +109,19 @@ def test_data_object_type_tracing(tmpdir):
     assert got == "int"
 
 
+def test_drawfn_type_tracing(tmpdir):
+    f = tmpdir.join("check_mypy_on_st_drawfn.py")
+    f.write(
+        "from hypothesis.strategies import DrawFn, text\n"
+        "def comp(draw: DrawFn) -> str:\n"
+        "    s = draw(text(), 123)\n"
+        "    reveal_type(s)\n"
+        "    return s\n"
+    )
+    got = get_mypy_analysed_type(str(f.realpath()), ...)
+    assert got == "str"
+
+
 def test_settings_preserves_type(tmpdir):
     f = tmpdir.join("check_mypy_on_settings.py")
     f.write(