Repository: HypothesisWorks/hypothesis

Issue:
Add Numpy strategy for `basic_indexes(shape)`
*This issue is extracted from #1829, and depends on `slices()` from #1929.*

MyGrad defines a strategy for [`basic_indexes`](https://github.com/rsokl/MyGrad/blob/62912d8bbb1c79ca3a72578a3e68720e4c29c594/tests/custom_strategies/__init__.py#L272), and we would like to bring a simpler version of it into the upstream `hypothesis.extra.numpy`.

Given an array shape, it generates valid indices for an array of that shape which conform to [the basic (i.e. no-copy) indexing specification](https://www.pythonlikeyoumeanit.com/Module3_IntroducingNumpy/BasicIndexing.html#Basic-Indexing). Examples from this strategy shrink from indices that produce high-dimensional views of the input array, down to indices that select a single element from the array.

An initial implementation could work something like:

```python
def basic_indexes(shape):
    # Validation logic for shape here - consider extracting a check_shape helper.
    return tuples(*(slices(s) for s in shape))
```

An optional extension would add an `allow_ellipsis=True` flag, which picks 0-2 dims to replace with `...` - only a few to ensure that the resulting arrays are typically reasonably small compared to the input shape, but may still generate a slice with a length != the number of dimensions to index.

`np.newaxis` (aka `None`) is also supported to add size-1 dimensions in any position, but we would put this behind a `allow_newaxis=False` flag.  As a disabled-by-default additional feature, we would prefer to receive a follow-up pull request to support `newaxis` after the initial version is complete and released to users.


Discussion:

Resolving patch:
diff --git a/hypothesis-python/src/hypothesis/extra/numpy.py b/hypothesis-python/src/hypothesis/extra/numpy.py
--- a/hypothesis-python/src/hypothesis/extra/numpy.py
+++ b/hypothesis-python/src/hypothesis/extra/numpy.py
@@ -23,7 +23,7 @@
 
 import hypothesis._strategies as st
 import hypothesis.internal.conjecture.utils as cu
-from hypothesis import Verbosity
+from hypothesis import Verbosity, assume
 from hypothesis._settings import note_deprecation
 from hypothesis.errors import InvalidArgument
 from hypothesis.internal.compat import PY2, hrange, integer_types
@@ -32,12 +32,14 @@
 from hypothesis.internal.validation import check_type, check_valid_interval
 from hypothesis.reporting import current_verbosity
 from hypothesis.searchstrategy import SearchStrategy
+from hypothesis.utils.conventions import not_set
 
 if False:
     from typing import Any, Union, Sequence, Tuple, Optional  # noqa
     from hypothesis.searchstrategy.strategies import T  # noqa
 
     Shape = Tuple[int, ...]  # noqa
+    BasicIndex = Tuple[Union[int, slice, ellipsis, np.newaxis], ...]  # noqa
 
 TIME_RESOLUTIONS = tuple("Y  M  D  h  m  s  ms  us  ns  ps  fs  as".split())
 
@@ -871,6 +873,121 @@ def broadcastable_shapes(shape, min_dims=0, max_dims=None, min_side=1, max_side=
     )
 
 
+class BasicIndexStrategy(SearchStrategy):
+    def __init__(self, shape, min_dims, max_dims, allow_ellipsis, allow_newaxis):
+        assert 0 <= min_dims <= max_dims <= 32
+        SearchStrategy.__init__(self)
+        self.shape = shape
+        self.min_dims = min_dims
+        self.max_dims = max_dims
+        self.allow_ellipsis = allow_ellipsis
+        self.allow_newaxis = allow_newaxis
+
+    def do_draw(self, data):
+        # General plan: determine the actual selection up front with a straightforward
+        # approach that shrinks well, then complicate it by inserting other things.
+        result = []
+        for dim_size in self.shape:
+            if dim_size == 0:
+                result.append(slice(None))
+                continue
+            strategy = st.integers(-dim_size, dim_size - 1) | st.slices(dim_size)
+            result.append(data.draw(strategy))
+        # Insert some number of new size-one dimensions if allowed
+        result_dims = sum(isinstance(idx, slice) for idx in result)
+        while (
+            self.allow_newaxis
+            and result_dims < self.max_dims
+            and (result_dims < self.min_dims or data.draw(st.booleans()))
+        ):
+            result.insert(data.draw(st.integers(0, len(result))), np.newaxis)
+            result_dims += 1
+        # Check that we'll have the right number of dimensions; reject if not.
+        # It's easy to do this by construction iff you don't care about shrinking,
+        # which is really important for array shapes.  So we filter instead.
+        assume(self.min_dims <= result_dims <= self.max_dims)
+        # This is a quick-and-dirty way to insert ..., xor shorten the indexer,
+        # but it means we don't have to do any structural analysis.
+        if self.allow_ellipsis and data.draw(st.booleans()):
+            # Choose an index; then replace all adjacent whole-dimension slices.
+            i = j = data.draw(st.integers(0, len(result)))
+            while i > 0 and result[i - 1] == slice(None):
+                i -= 1
+            while j < len(result) and result[j] == slice(None):
+                j += 1
+            result[i:j] = [Ellipsis]
+        else:
+            while result[-1:] == [slice(None, None)] and data.draw(st.integers(0, 7)):
+                result.pop()
+        return tuple(result)
+
+
+@st.defines_strategy
+def basic_indices(
+    shape,
+    __reserved=not_set,
+    min_dims=0,
+    max_dims=None,
+    allow_newaxis=False,
+    allow_ellipsis=True,
+):
+    # type: (Shape, Any, int, int, bool, bool) -> st.SearchStrategy[BasicIndex]
+    """
+    The ``basic_indices`` strategy generates `basic indexes
+    <https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html>`__  for
+    arrays of the specified shape, which may include dimensions of size zero.
+
+    It generates tuples containing some mix of integers, :obj:`python:slice` objects,
+    ``...`` (Ellipsis), and :obj:`numpy:numpy.newaxis`; which when used to index a
+    ``shape``-shaped array will produce either a scalar or a shared-memory view.
+
+    * ``shape``: the array shape that will be indexed, as a tuple of integers >= 0.
+      This must be at least two-dimensional for a tuple to be a valid basic index;
+      for one-dimensional arrays use :func:`~hypothesis.strategies.slices` instead.
+    * ``min_dims``: the minimum dimensionality of the resulting view from use of
+      the generated index.  When ``min_dims == 0``, scalars and zero-dimensional
+      arrays are both allowed.
+    * ``max_dims``: the maximum dimensionality of the resulting view.
+      If not specified, it defaults to ``max(len(shape), min_dims) + 2``.
+    * ``allow_ellipsis``: whether ``...``` is allowed in the index.
+    * ``allow_newaxis``: whether :obj:`numpy:numpy.newaxis` is allowed in the index.
+
+    Note that the length of the generated tuple may be anywhere between zero
+    and ``min_dims``.  It may not match the length of ``shape``, or even the
+    dimensionality of the array view resulting from its use!
+    """
+    # Arguments to exclude scalars, zero-dim arrays, and dims of size zero were
+    # all considered and rejected.  We want users to explicitly consider those
+    # cases if they're dealing in general indexers, and while it's fiddly we can
+    # back-compatibly add them later (hence using __reserved to sim kwonlyargs).
+    check_type(tuple, shape, "shape")
+    if __reserved is not not_set:
+        raise InvalidArgument("Do not pass the __reserved argument.")
+    check_type(bool, allow_ellipsis, "allow_ellipsis")
+    check_type(bool, allow_newaxis, "allow_newaxis")
+    check_type(integer_types, min_dims, "min_dims")
+    if max_dims is None:
+        max_dims = min(max(len(shape), min_dims) + 2, 32)
+    else:
+        check_type(integer_types, max_dims, "max_dims")
+    order_check("dims", 0, min_dims, max_dims)
+    check_argument(
+        max_dims <= 32,
+        "max_dims=%r, but numpy arrays have at most 32 dimensions" % (max_dims,),
+    )
+    check_argument(
+        all(isinstance(x, integer_types) and x >= 0 for x in shape),
+        "shape=%r, but all dimensions must be of integer size >= 0" % (shape,),
+    )
+    return BasicIndexStrategy(
+        shape,
+        min_dims=min_dims,
+        max_dims=max_dims,
+        allow_ellipsis=allow_ellipsis,
+        allow_newaxis=allow_newaxis,
+    )
+
+
 @st.defines_strategy
 def integer_array_indices(shape, result_shape=array_shapes(), dtype="int"):
     # type: (Shape, SearchStrategy[Shape], np.dtype) -> st.SearchStrategy[Tuple[np.ndarray, ...]]


Test patch:
diff --git a/hypothesis-python/tests/numpy/test_argument_validation.py b/hypothesis-python/tests/numpy/test_argument_validation.py
--- a/hypothesis-python/tests/numpy/test_argument_validation.py
+++ b/hypothesis-python/tests/numpy/test_argument_validation.py
@@ -117,6 +117,17 @@ def e(a, **kwargs):
             min_side=2,
             max_side=3,
         ),
+        e(nps.basic_indices, shape=0),
+        e(nps.basic_indices, shape=("1", "2")),
+        e(nps.basic_indices, shape=(0, -1)),
+        e(nps.basic_indices, shape=(0, 0), allow_newaxis=None),
+        e(nps.basic_indices, shape=(0, 0), allow_ellipsis=None),
+        e(nps.basic_indices, shape=(0, 0), min_dims=-1),
+        e(nps.basic_indices, shape=(0, 0), min_dims=1.0),
+        e(nps.basic_indices, shape=(0, 0), max_dims=-1),
+        e(nps.basic_indices, shape=(0, 0), max_dims=1.0),
+        e(nps.basic_indices, shape=(0, 0), min_dims=2, max_dims=1),
+        e(nps.basic_indices, shape=(0, 0), max_dims=50),
         e(nps.integer_array_indices, shape=()),
         e(nps.integer_array_indices, shape=(2, 0)),
         e(nps.integer_array_indices, shape="a"),
@@ -157,3 +168,10 @@ def test_byte_string_dtype_len_0(data):
 def test_unicode_string_dtype_len_0(data):
     s = nps.unicode_string_dtypes(min_len=0, max_len=0)
     assert data.draw(s).itemsize == 4
+
+
+def test_test_basic_indices_kwonly_emulation():
+    with pytest.raises(InvalidArgument):
+        nps.basic_indices((), 0, 1).validate()
+    with pytest.raises(InvalidArgument):
+        nps.basic_indices((), __reserved=None).validate()
diff --git a/hypothesis-python/tests/numpy/test_gen_data.py b/hypothesis-python/tests/numpy/test_gen_data.py
--- a/hypothesis-python/tests/numpy/test_gen_data.py
+++ b/hypothesis-python/tests/numpy/test_gen_data.py
@@ -818,3 +818,77 @@ def test_advanced_integer_index_can_generate_any_pattern(shape, data):
         lambda index: np.all(target == x[index]),
         settings(max_examples=10 ** 6),
     )
+
+
+@pytest.mark.parametrize(
+    "condition",
+    [
+        lambda ix: Ellipsis in ix,
+        lambda ix: Ellipsis not in ix,
+        lambda ix: np.newaxis in ix,
+        lambda ix: np.newaxis not in ix,
+    ],
+)
+def test_test_basic_indices_options(condition):
+    indexers = nps.array_shapes(min_dims=0, max_dims=32).flatmap(
+        lambda shape: nps.basic_indices(shape, allow_newaxis=True)
+    )
+    find_any(indexers, condition)
+
+
+def test_test_basic_indices_can_generate_empty_tuple():
+    find_any(nps.basic_indices(shape=(0, 0), allow_ellipsis=True), lambda ix: ix == ())
+
+
+def test_test_basic_indices_can_generate_long_ellipsis():
+    # Runs of slice(None) - such as [0,:,:,:,0] - can be replaced by e.g. [0,...,0]
+    find_any(
+        nps.basic_indices(shape=(1, 0, 0, 0, 1), allow_ellipsis=True),
+        lambda ix: len(ix) == 3 and ix[1] == Ellipsis,
+    )
+
+
+@given(
+    shape=nps.array_shapes(min_dims=0, max_side=4)
+    | nps.array_shapes(min_dims=0, min_side=0, max_side=10),
+    min_dims=st.integers(0, 5),
+    allow_ellipsis=st.booleans(),
+    allow_newaxis=st.booleans(),
+    data=st.data(),
+)
+def test_basic_indices_generate_valid_indexers(
+    shape, min_dims, allow_ellipsis, allow_newaxis, data
+):
+    max_dims = data.draw(st.none() | st.integers(min_dims, 32), label="max_dims")
+    indexer = data.draw(
+        nps.basic_indices(
+            shape,
+            min_dims=min_dims,
+            max_dims=max_dims,
+            allow_ellipsis=allow_ellipsis,
+            allow_newaxis=allow_newaxis,
+        ),
+        label="indexer",
+    )
+    # Check that disallowed things are indeed absent
+    if not allow_newaxis:
+        assert 0 <= len(indexer) <= len(shape) + int(allow_ellipsis)
+        assert np.newaxis not in shape
+    if not allow_ellipsis:
+        assert Ellipsis not in shape
+
+    if 0 in shape:
+        # If there's a zero in the shape, the array will have no elements.
+        array = np.zeros(shape)
+        assert array.size == 0
+    elif np.prod(shape) <= 10 ** 5:
+        # If it's small enough to instantiate, do so with distinct elements.
+        array = np.arange(np.prod(shape)).reshape(shape)
+    else:
+        # We can't cheat on this one, so just try another.
+        assume(False)
+    view = array[indexer]
+    if not np.isscalar(view):
+        assert min_dims <= view.ndim <= (32 if max_dims is None else max_dims)
+        if view.size:
+            assert np.shares_memory(view, array)