diff --git a/numpy/core/src/multiarray/mapping.c b/numpy/core/src/multiarray/mapping.c index d64962f87150..8b9b67387181 100644 --- a/numpy/core/src/multiarray/mapping.c +++ b/numpy/core/src/multiarray/mapping.c @@ -2328,7 +2328,7 @@ PyArray_MapIterNext(PyArrayMapIterObject *mit) * @param Number of indices * @param The array that is being iterated * - * @return 0 on success -1 on failure + * @return 0 on success -1 on failure (broadcasting or too many fancy indices) */ static int mapiter_fill_info(PyArrayMapIterObject *mit, npy_index_info *indices, @@ -2369,6 +2369,17 @@ mapiter_fill_info(PyArrayMapIterObject *mit, npy_index_info *indices, } } + /* Before contunuing, ensure that there are not too fancy indices */ + if (indices[i].type & HAS_FANCY) { + if (NPY_UNLIKELY(j >= NPY_MAXDIMS)) { + PyErr_Format(PyExc_IndexError, + "too many advanced (array) indices. This probably " + "means you are indexing with too many booleans. " + "(more than %d found)", NPY_MAXDIMS); + return -1; + } + } + /* (iterating) fancy index, store the iterator */ if (indices[i].type == HAS_FANCY) { mit->fancy_strides[j] = PyArray_STRIDE(arr, curr_dim); @@ -2655,6 +2666,7 @@ PyArray_MapIterNew(npy_index_info *indices , int index_num, int index_type, /* For shape reporting on error */ PyArrayObject *original_extra_op = extra_op; + /* NOTE: MAXARGS is the actual limit (2*NPY_MAXDIMS is index number one) */ PyArrayObject *index_arrays[NPY_MAXDIMS]; PyArray_Descr *intp_descr; PyArray_Descr *dtypes[NPY_MAXDIMS]; /* borrowed references */ diff --git a/numpy/core/tests/test_indexing.py b/numpy/core/tests/test_indexing.py index 667c49240e4a..73dbc429c89c 100644 --- a/numpy/core/tests/test_indexing.py +++ b/numpy/core/tests/test_indexing.py @@ -3,6 +3,8 @@ import functools import operator +import pytest + import numpy as np from numpy.core._multiarray_tests import array_indexing from itertools import product @@ -547,6 +549,21 @@ def test_character_assignment(self): assert_array_equal(arr[0], np.array("asdfg", dtype="c")) assert arr[0, 1] == b"s" # make sure not all were set to "a" for both + @pytest.mark.parametrize("index", + [True, False, np.array([0])]) + @pytest.mark.parametrize("num", [32, 40]) + @pytest.mark.parametrize("original_ndim", [1, 32]) + def test_too_many_advanced_indices(self, index, num, original_ndim): + # These are limitations based on the number of arguments we can process. + # For `num=32` (and all boolean cases), the result is actually define; + # but the use of NpyIter (NPY_MAXARGS) limits it for technical reasons. + arr = np.ones((1,) * original_ndim) + with pytest.raises(IndexError): + arr[(index,) * num] + with pytest.raises(IndexError): + arr[(index,) * num] = 1. + + class TestFieldIndexing: def test_scalar_return_type(self): # Field access on an array should return an array, even if it