diff options
Diffstat (limited to 'numpy/core')
| -rw-r--r-- | numpy/core/src/multiarray/nditer_constr.c | 7 | ||||
| -rw-r--r-- | numpy/core/tests/test_nditer.py | 9 |
2 files changed, 15 insertions, 1 deletions
diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c index 434b2a1f3..053777f6a 100644 --- a/numpy/core/src/multiarray/nditer_constr.c +++ b/numpy/core/src/multiarray/nditer_constr.c @@ -16,6 +16,7 @@ #include "nditer_impl.h" #include "arrayobject.h" +#include "scalarmathmodule.h" /* Internal helper functions private to this file */ static int @@ -1709,7 +1710,11 @@ npyiter_fill_axisdata(NpyIter *iter, npy_uint32 flags, npyiter_opitflags *op_itf /* Now fill in the ITERSIZE member */ NIT_ITERSIZE(iter) = 1; for (idim = 0; idim < ndim; ++idim) { - NIT_ITERSIZE(iter) *= broadcast_shape[idim]; + if (npy_mul_with_overflow_intp(&NIT_ITERSIZE(iter), + NIT_ITERSIZE(iter), broadcast_shape[idim])) { + PyErr_SetString(PyExc_ValueError, "iterator is too large"); + return 0; + } } /* The range defaults to everything */ NIT_ITERSTART(iter) = 0; diff --git a/numpy/core/tests/test_nditer.py b/numpy/core/tests/test_nditer.py index bc20b900e..95c127316 100644 --- a/numpy/core/tests/test_nditer.py +++ b/numpy/core/tests/test_nditer.py @@ -2579,5 +2579,14 @@ def test_0d_nested_iter(): assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]]) +def test_iter_too_large(): + # The total size of the iterator must not exceed the maximum intp due + # to broadcasting. Dividing by 1024 will keep it small enough to + # give a legal array. + size = np.iinfo(np.intp).max // 1024 + arr = np.lib.stride_tricks.as_strided(np.zeros(1), (size,), (0,)) + assert_raises(ValueError, nditer, (arr, arr[:, None])) + + if __name__ == "__main__": run_module_suite() |