Merge pull request #4611 from Zac-HD/claude/fix-hypothesis-4607-018dHECpz3mwL9NUfXjptxqC

Improve type annotations for array strategies and `assume()`

Author: Zac-HD

.claude/CLAUDE.md

@@ -40,3 +40,4 @@ When creating a PR:
    - **Idiomatic** - follows Python and Hypothesis conventions
    - **Minimally commented** - code should be self-documenting; only add comments where truly needed
 2. **Run `./build.sh format; ./build.sh lint`** immediately before committing to auto-format and lint code
+3. **Do not reference issues or PRs in commit messages** (e.g., avoid `Fixes #1234` or `See #5678`) - this clutters the issue timeline with unnecessary links

build.sh

@@ -19,9 +19,9 @@ SCRIPTS="$ROOT/tooling/scripts"
 # shellcheck source=tooling/scripts/common.sh
 source "$SCRIPTS/common.sh"
 
-if [ -n "${GITHUB_ACTIONS-}" ] || [ -n "${CODESPACES-}" ] ; then
-    # We're on GitHub Actions or Codespaces and already set up a suitable Python
-    PYTHON=$(command -v python)
+if [ -n "${GITHUB_ACTIONS-}" ] || [ -n "${CODESPACES-}" ] || [ -n "${CLAUDECODE-}" ] ; then
+    # We're on GitHub Actions, Codespaces, or Claude Code and already have a suitable Python
+    PYTHON=$(command -v python3 || command -v python)
 else
     # Otherwise, we install it from scratch
     # NOTE: tooling keeps this version in sync with ci_version in tooling
@@ -40,9 +40,14 @@ export PYTHONPATH="$ROOT/tooling/src"
 
 if ! "$TOOL_PYTHON" -m hypothesistooling check-installed ; then
   rm -rf "$TOOL_VIRTUALENV"
-  "$PYTHON" -m pip install --upgrade pip
-  "$PYTHON" -m pip install --upgrade virtualenv
-  "$PYTHON" -m virtualenv "$TOOL_VIRTUALENV"
+  if [ -n "${CLAUDECODE-}" ] ; then
+    # Claude Code: use venv (available) and skip pip upgrades (debian-managed)
+    "$PYTHON" -m venv "$TOOL_VIRTUALENV"
+  else
+    "$PYTHON" -m pip install --upgrade pip
+    "$PYTHON" -m pip install --upgrade virtualenv
+    "$PYTHON" -m virtualenv "$TOOL_VIRTUALENV"
+  fi
   "$TOOL_PYTHON" -m pip install --no-warn-script-location -r requirements/tools.txt
 fi
 

hypothesis-python/RELEASE.rst

@@ -0,0 +1,10 @@
+RELEASE_TYPE: patch
+
+This patch improves the type annotations for :func:`~hypothesis.extra.numpy.basic_indices`.
+The return type now accurately reflects the ``allow_ellipsis`` and ``allow_newaxis``
+parameters, excluding ``EllipsisType`` or ``None`` from the union when those index
+types are disabled (:issue:`4607`).
+
+Additionally, :func:`~hypothesis.assume` now has overloaded type annotations:
+``assume(True)`` returns ``Literal[True]``, while ``assume(False)`` and
+``assume(None)`` return ``NoReturn``.

hypothesis-python/src/hypothesis/control.py

@@ -14,7 +14,7 @@
 from collections import defaultdict
 from collections.abc import Callable, Sequence
 from contextlib import contextmanager
-from typing import Any, NoReturn, Optional
+from typing import Any, Literal, NoReturn, Optional, overload
 from weakref import WeakKeyDictionary
 
 from hypothesis import Verbosity, settings
@@ -49,7 +49,13 @@ def reject() -> NoReturn:
     raise UnsatisfiedAssumption(where)
 
 
-def assume(condition: object) -> bool:
+@overload
+def assume(condition: Literal[False] | None) -> NoReturn: ...
+@overload
+def assume(condition: object) -> Literal[True]: ...
+
+
+def assume(condition: object) -> Literal[True]:
     """Calling ``assume`` is like an :ref:`assert <python:assert>` that marks
     the example as bad, rather than failing the test.
 

hypothesis-python/src/hypothesis/extra/_array_helpers.py

@@ -22,11 +22,15 @@
 
 __all__ = [
     "NDIM_MAX",
+    "_BIE",
     "BasicIndex",
     "BasicIndexStrategy",
     "BroadcastableShapes",
     "MutuallyBroadcastableShapesStrategy",
     "Shape",
+    "_BIENoEllipsis",
+    "_BIENoEllipsisNoNewaxis",
+    "_BIENoNewaxis",
     "array_shapes",
     "broadcastable_shapes",
     "check_argument",
@@ -38,7 +42,15 @@
 
 
 Shape = tuple[int, ...]
-BasicIndex = tuple[int | slice | None | EllipsisType, ...]
+
+# Type aliases for basic array index elements. Variants exist to accurately
+# type the return value of basic_indices() based on allow_ellipsis/allow_newaxis.
+_BIE = int | slice | None | EllipsisType
+_BIENoEllipsis = int | slice | None
+_BIENoNewaxis = int | slice | EllipsisType
+_BIENoEllipsisNoNewaxis = int | slice
+
+BasicIndex = _BIE | tuple[_BIE, ...]
 
 
 class BroadcastableShapes(NamedTuple):

hypothesis-python/src/hypothesis/extra/numpy.py

@@ -30,11 +30,15 @@
 from hypothesis._settings import note_deprecation
 from hypothesis.errors import HypothesisException, InvalidArgument
 from hypothesis.extra._array_helpers import (
+    _BIE,
     NDIM_MAX,
     BasicIndex,
     BasicIndexStrategy,
     BroadcastableShapes,
     Shape,
+    _BIENoEllipsis,
+    _BIENoEllipsisNoNewaxis,
+    _BIENoNewaxis,
     array_shapes,
     broadcastable_shapes,
     check_argument,
@@ -1092,6 +1096,52 @@ def mutually_broadcastable_shapes(*args, **kwargs):
     """
 
 
+@overload
+def basic_indices(
+    shape: Shape,
+    *,
+    min_dims: int = 0,
+    max_dims: int | None = None,
+    allow_newaxis: Literal[False] = ...,
+    allow_ellipsis: Literal[False],
+) -> st.SearchStrategy[
+    _BIENoEllipsisNoNewaxis | tuple[_BIENoEllipsisNoNewaxis, ...]
+]: ...
+
+
+@overload
+def basic_indices(
+    shape: Shape,
+    *,
+    min_dims: int = 0,
+    max_dims: int | None = None,
+    allow_newaxis: Literal[False] = ...,
+    allow_ellipsis: Literal[True] = ...,
+) -> st.SearchStrategy[_BIENoNewaxis | tuple[_BIENoNewaxis, ...]]: ...
+
+
+@overload
+def basic_indices(
+    shape: Shape,
+    *,
+    min_dims: int = 0,
+    max_dims: int | None = None,
+    allow_newaxis: Literal[True],
+    allow_ellipsis: Literal[False],
+) -> st.SearchStrategy[_BIENoEllipsis | tuple[_BIENoEllipsis, ...]]: ...
+
+
+@overload
+def basic_indices(
+    shape: Shape,
+    *,
+    min_dims: int = 0,
+    max_dims: int | None = None,
+    allow_newaxis: Literal[True],
+    allow_ellipsis: Literal[True] = ...,
+) -> st.SearchStrategy[_BIE | tuple[_BIE, ...]]: ...
+
+
 @defines_strategy()
 def basic_indices(
     shape: Shape,

whole_repo_tests/types/revealed_types.py

@@ -22,6 +22,13 @@
 else:
     NP1 = np_version.startswith("1.")
 
+ASSUME_REVEALED_TYPES = [
+    ("assume(False)", "Never"),
+    ("assume(None)", "Never"),
+    ("assume(True)", "Literal[True]"),
+    ("assume(1)", "Literal[True]"),
+]
+
 REVEALED_TYPES = [
     ("integers()", "int"),
     ("text()", "str"),
@@ -202,4 +209,30 @@ class DifferingRevealedTypes(NamedTuple):
         'integer_array_indices(shape=(2, 3), dtype=np.dtype("uint8"))',
         "tuple[ndarray[tuple[int, ...], dtype[unsignedinteger[_8Bit]]], ...]",
     ),
+    # basic_indices with allow_ellipsis=False (no EllipsisType differences)
+    (
+        "basic_indices((3, 4), allow_ellipsis=False)",
+        "int | slice[Any, Any, Any] | tuple[int | slice[Any, Any, Any], ...]",
+    ),
+]
+
+# basic_indices tests where mypy/pyright differ in EllipsisType representation
+NUMPY_DIFF_REVEALED_TYPES = [
+    # mypy uses types.EllipsisType, pyright uses EllipsisType
+    DifferingRevealedTypes(
+        "basic_indices((3, 4))",
+        "int | slice[Any, Any, Any] | types.EllipsisType | tuple[int | slice[Any, Any, Any] | types.EllipsisType, ...]",
+        "int | slice[Any, Any, Any] | EllipsisType | tuple[int | slice[Any, Any, Any] | EllipsisType, ...]",
+    ),
+    # pyright also reorders None to the end
+    DifferingRevealedTypes(
+        "basic_indices((3, 4), allow_newaxis=True, allow_ellipsis=False)",
+        "int | slice[Any, Any, Any] | None | tuple[int | slice[Any, Any, Any] | None, ...]",
+        "int | slice[Any, Any, Any] | tuple[int | slice[Any, Any, Any] | None, ...] | None",
+    ),
+    DifferingRevealedTypes(
+        "basic_indices((3, 4), allow_newaxis=True)",
+        "int | slice[Any, Any, Any] | None | types.EllipsisType | tuple[int | slice[Any, Any, Any] | None | types.EllipsisType, ...]",
+        "int | slice[Any, Any, Any] | EllipsisType | tuple[int | slice[Any, Any, Any] | EllipsisType | None, ...] | None",
+    ),
 ]

whole_repo_tests/types/test_mypy.py

@@ -18,7 +18,9 @@
 from hypothesistooling.scripts import pip_tool, tool_path
 
 from .revealed_types import (
+    ASSUME_REVEALED_TYPES,
     DIFF_REVEALED_TYPES,
+    NUMPY_DIFF_REVEALED_TYPES,
     NUMPY_REVEALED_TYPES,
     PYTHON_VERSIONS,
     REVEALED_TYPES,
@@ -142,7 +144,27 @@ def test_revealed_types(tmp_path, val, expect):
     assert typ == f"SearchStrategy[{expect}]"
 
 
-@pytest.mark.parametrize("val,expect", NUMPY_REVEALED_TYPES)
+@pytest.mark.parametrize("val,expect", ASSUME_REVEALED_TYPES)
+def test_assume_revealed_types(tmp_path, val, expect):
+    """Check that Mypy infers the correct return type for assume()."""
+    f = tmp_path / "check.py"
+    f.write_text(
+        textwrap.dedent(
+            f"""
+            from hypothesis import assume
+            reveal_type({val})
+            """
+        ),
+        encoding="utf-8",
+    )
+    typ = get_mypy_analysed_type(f)
+    assert typ == expect
+
+
+@pytest.mark.parametrize(
+    "val,expect",
+    [*NUMPY_REVEALED_TYPES, *((x.value, x.mypy) for x in NUMPY_DIFF_REVEALED_TYPES)],
+)
 def test_numpy_revealed_types(tmp_path, val, expect):
     f = tmp_path / "check.py"
     f.write_text(

whole_repo_tests/types/test_pyright.py

@@ -24,7 +24,9 @@
 from hypothesistooling.scripts import pip_tool, tool_path
 
 from .revealed_types import (
+    ASSUME_REVEALED_TYPES,
     DIFF_REVEALED_TYPES,
+    NUMPY_DIFF_REVEALED_TYPES,
     NUMPY_REVEALED_TYPES,
     PYTHON_VERSIONS,
     REVEALED_TYPES,
@@ -200,9 +202,31 @@ def test_revealed_types(tmp_path, val, expect):
     assert typ == f"SearchStrategy[{expect}]"
 
 
-@pytest.mark.parametrize("val,expect", NUMPY_REVEALED_TYPES)
+@pytest.mark.parametrize("val,expect", ASSUME_REVEALED_TYPES)
+def test_assume_revealed_types(tmp_path, val, expect):
+    """Check that Pyright infers the correct return type for assume()."""
+    f = tmp_path / "check.py"
+    f.write_text(
+        textwrap.dedent(
+            f"""
+            from hypothesis import assume
+            reveal_type({val})
+            """
+        ),
+        encoding="utf-8",
+    )
+    _write_config(tmp_path)
+    typ = get_pyright_analysed_type(f)
+    # Pyright uses NoReturn, mypy uses Never
+    assert typ == (expect if expect != "Never" else "NoReturn")
+
+
+@pytest.mark.parametrize(
+    "val,expect",
+    [*NUMPY_REVEALED_TYPES, *((x.value, x.pyright) for x in NUMPY_DIFF_REVEALED_TYPES)],
+)
 def test_numpy_revealed_types(tmp_path, val, expect):
-    f = tmp_path / (expect + ".py")
+    f = tmp_path / "check.py"
     f.write_text(
         textwrap.dedent(
             f"""