Improve functions() annotations

hypothesis-python/RELEASE.rst

@@ -3,4 +3,5 @@ RELEASE_TYPE: minor
 This release uses :pep:`612` :obj:`python:typing.ParamSpec` (or the
 :pypi:`typing_extensions` backport) to express the first-argument-removing
 behaviour of :func:`@st.composite <hypothesis.strategies.composite>`
+and signature-preservation of :func:`~hypothesis.strategies.functions`
 to IDEs, editor plugins, and static type checkers such as :pypi:`mypy`.

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

@@ -1966,52 +1966,113 @@ def emails() -> SearchStrategy[str]:
     )
 
 
-@defines_strategy()
-def functions(
-    *,
-    like: Callable[..., Any] = lambda: None,
-    returns: Union[SearchStrategy[Any], InferType] = infer,
-    pure: bool = False,
-) -> SearchStrategy[Callable[..., Any]]:
-    # The proper type signature of `functions()` would have T instead of Any, but mypy
-    # disallows default args for generics: https://github.com/python/mypy/issues/3737
-    """functions(*, like=lambda: None, returns=..., pure=False)
-
-    A strategy for functions, which can be used in callbacks.
-
-    The generated functions will mimic the interface of ``like``, which must
-    be a callable (including a class, method, or function).  The return value
-    for the function is drawn from the ``returns`` argument, which must be a
-    strategy.  If ``returns`` is not passed, we attempt to infer a strategy
-    from the return-type annotation if present, falling back to :func:`~none`.
-
-    If ``pure=True``, all arguments passed to the generated function must be
-    hashable, and if passed identical arguments the original return value will
-    be returned again - *not* regenerated, so beware mutable values.
-
-    If ``pure=False``, generated functions do not validate their arguments, and
-    may return a different value if called again with the same arguments.
-
-    Generated functions can only be called within the scope of the ``@given``
-    which created them.  This strategy does not support ``.example()``.
-    """
+def _functions(*, like=lambda: None, returns=infer, pure=False):
+    # Wrapped up to use ParamSpec below
     check_type(bool, pure, "pure")
     if not callable(like):
         raise InvalidArgument(
             "The first argument to functions() must be a callable to imitate, "
             f"but got non-callable like={nicerepr(like)!r}"
         )
-
     if returns is None or returns is infer:
         # Passing `None` has never been *documented* as working, but it still
         # did from May 2020 to Jan 2022 so we'll avoid breaking it without cause.
         hints = get_type_hints(like)
         returns = from_type(hints.get("return", type(None)))
-
     check_strategy(returns, "returns")
     return FunctionStrategy(like, returns, pure)
 
 
+if typing.TYPE_CHECKING or ParamSpec is not None:
+    P = ParamSpec("P")
+
+    @overload
+    def functions(*, pure: bool = False) -> SearchStrategy[Callable[[], None]]:
+        pass
+
+    @overload
+    def functions(
+        *,
+        like: Callable[P, T],
+        pure: bool = False,
+    ) -> SearchStrategy[Callable[P, T]]:
+        pass
+
+    @overload
+    def functions(
+        *,
+        returns: SearchStrategy[T],
+        pure: bool = False,
+    ) -> SearchStrategy[Callable[[], T]]:
+        pass
+
+    @overload
+    def functions(
+        *,
+        like: Callable[P, Any],
+        returns: SearchStrategy[T],
+        pure: bool = False,
+    ) -> SearchStrategy[Callable[P, T]]:
+        pass
+
+    @defines_strategy()
+    def functions(*, like=lambda: None, returns=infer, pure=False):
+        # We shouldn't need overloads here, but mypy disallows default args for
+        # generics: https://github.com/python/mypy/issues/3737
+        """functions(*, like=lambda: None, returns=..., pure=False)
+
+        A strategy for functions, which can be used in callbacks.
+
+        The generated functions will mimic the interface of ``like``, which must
+        be a callable (including a class, method, or function).  The return value
+        for the function is drawn from the ``returns`` argument, which must be a
+        strategy.  If ``returns`` is not passed, we attempt to infer a strategy
+        from the return-type annotation if present, falling back to :func:`~none`.
+
+        If ``pure=True``, all arguments passed to the generated function must be
+        hashable, and if passed identical arguments the original return value will
+        be returned again - *not* regenerated, so beware mutable values.
+
+        If ``pure=False``, generated functions do not validate their arguments, and
+        may return a different value if called again with the same arguments.
+
+        Generated functions can only be called within the scope of the ``@given``
+        which created them.  This strategy does not support ``.example()``.
+        """
+        return _functions(like=like, returns=returns, pure=pure)
+
+else:
+
+    @defines_strategy()
+    def functions(
+        *,
+        like: Callable[..., Any] = lambda: None,
+        returns: Union[SearchStrategy[Any], InferType] = infer,
+        pure: bool = False,
+    ) -> SearchStrategy[Callable[..., Any]]:
+        """functions(*, like=lambda: None, returns=..., pure=False)
+
+        A strategy for functions, which can be used in callbacks.
+
+        The generated functions will mimic the interface of ``like``, which must
+        be a callable (including a class, method, or function).  The return value
+        for the function is drawn from the ``returns`` argument, which must be a
+        strategy.  If ``returns`` is not passed, we attempt to infer a strategy
+        from the return-type annotation if present, falling back to :func:`~none`.
+
+        If ``pure=True``, all arguments passed to the generated function must be
+        hashable, and if passed identical arguments the original return value will
+        be returned again - *not* regenerated, so beware mutable values.
+
+        If ``pure=False``, generated functions do not validate their arguments, and
+        may return a different value if called again with the same arguments.
+
+        Generated functions can only be called within the scope of the ``@given``
+        which created them.  This strategy does not support ``.example()``.
+        """
+        return _functions(like=like, returns=returns, pure=pure)
+
+
 @composite
 def slices(draw: Any, size: int) -> slice:
     """Generates slices that will select indices up to the supplied size

whole-repo-tests/test_mypy.py

@@ -172,6 +172,27 @@ def test_composite_type_tracing(tmpdir):
     assert got == "def (x: int) -> int"
 
 
+@pytest.mark.parametrize(
+    "source, expected",
+    [
+        ("", "def ()"),
+        ("like=f", "def (x: int) -> int"),
+        ("returns=booleans()", "def () -> bool"),
+        ("like=f, returns=booleans()", "def (x: int) -> bool"),
+    ],
+)
+def test_functions_type_tracing(tmpdir, source, expected):
+    f = tmpdir.join("check_mypy_on_st_composite.py")
+    f.write(
+        "from hypothesis.strategies import booleans, functions\n"
+        "def f(x: int) -> int: return x\n"
+        f"g = functions({source}).example()\n"
+        "reveal_type(g)\n"
+    )
+    got = get_mypy_analysed_type(str(f.realpath()), ...)
+    assert got == expected, (got, expected)
+
+
 def test_settings_preserves_type(tmpdir):
     f = tmpdir.join("check_mypy_on_settings.py")
     f.write(
@@ -434,7 +455,7 @@ def test_pos_only_args(tmpdir):
 
             st.tuples(a1=st.integers())
             st.tuples(a1=st.integers(), a2=st.integers())
-            
+
             st.one_of(a1=st.integers())
             st.one_of(a1=st.integers(), a2=st.integers())
             """