diff --git a/include/pybind11/numpy.h b/include/pybind11/numpy.h index a26ac2b4cd..1c3613c9df 100644 --- a/include/pybind11/numpy.h +++ b/include/pybind11/numpy.h @@ -1156,12 +1156,14 @@ class array : public buffer { } #endif - /// Dimension along a given axis + /// Dimension along a given axis. Negative axes count from the end. + /// Throws if dim is outside [-ndim(), ndim()). ssize_t shape(ssize_t dim) const { - if (dim >= ndim()) { + const ssize_t axis = dim < 0 ? dim + ndim() : dim; + if (static_cast(axis) >= static_cast(ndim())) { fail_dim_check(dim, "invalid axis"); } - return shape()[dim]; + return shape()[axis]; } /// Strides of the array @@ -1174,12 +1176,14 @@ class array : public buffer { } #endif - /// Stride along a given axis + /// Stride along a given axis. Negative axes count from the end. + /// Throws if dim is outside [-ndim(), ndim()). ssize_t strides(ssize_t dim) const { - if (dim >= ndim()) { + const ssize_t axis = dim < 0 ? dim + ndim() : dim; + if (static_cast(axis) >= static_cast(ndim())) { fail_dim_check(dim, "invalid axis"); } - return strides()[dim]; + return strides()[axis]; } /// Return the NumPy array flags @@ -1196,15 +1200,16 @@ class array : public buffer { } /// Pointer to the contained data. If index is not provided, points to the - /// beginning of the buffer. May throw if the index would lead to out of bounds access. + /// beginning of the buffer. Negative indices count from the end of each dimension. + /// Throws if an index is out of bounds. template const void *data(Ix... index) const { return static_cast(detail::array_proxy(m_ptr)->data + offset_at(index...)); } /// Mutable pointer to the contained data. If index is not provided, points to the - /// beginning of the buffer. May throw if the index would lead to out of bounds access. - /// May throw if the array is not writeable. + /// beginning of the buffer. Negative indices count from the end of each dimension. + /// Throws if an index is out of bounds or the array is not writeable. template void *mutable_data(Ix... index) { check_writeable(); @@ -1212,7 +1217,7 @@ class array : public buffer { } /// Byte offset from beginning of the array to a given index (full or partial). - /// May throw if the index would lead to out of bounds access. + /// Negative indices count from the end of each dimension. Throws if an index is out of bounds. template ssize_t offset_at(Ix... index) const { if ((ssize_t) sizeof...(index) > ndim()) { @@ -1224,7 +1229,7 @@ class array : public buffer { ssize_t offset_at() const { return 0; } /// Item count from beginning of the array to a given index (full or partial). - /// May throw if the index would lead to out of bounds access. + /// Negative indices count from the end of each dimension. Throws if an index is out of bounds. template ssize_t index_at(Ix... index) const { return offset_at(index...) / itemsize(); @@ -1337,8 +1342,7 @@ class array : public buffer { template ssize_t byte_offset(Ix... index) const { - check_dimensions(index...); - return detail::byte_offset_unsafe(strides(), ssize_t(index)...); + return byte_offset_impl(ssize_t(0), shape(), strides(), ssize_t(index)...); } void check_writeable() const { @@ -1347,21 +1351,20 @@ class array : public buffer { } } - template - void check_dimensions(Ix... index) const { - check_dimensions_impl(ssize_t(0), shape(), ssize_t(index)...); - } - - void check_dimensions_impl(ssize_t, const ssize_t *) const {} + ssize_t byte_offset_impl(ssize_t, const ssize_t *, const ssize_t *) const { return 0; } template - void check_dimensions_impl(ssize_t axis, const ssize_t *shape, ssize_t i, Ix... index) const { - if (i >= *shape) { + ssize_t byte_offset_impl( + ssize_t axis, const ssize_t *shape, const ssize_t *strides, ssize_t i, Ix... index) const { + const ssize_t normalized = i < 0 ? i + *shape : i; + // Unsigned comparison also rejects indices that remain negative after normalization. + if (static_cast(normalized) >= static_cast(*shape)) { throw index_error(std::string("index ") + std::to_string(i) + " is out of bounds for axis " + std::to_string(axis) + " with size " + std::to_string(*shape)); } - check_dimensions_impl(axis + 1, shape + 1, index...); + const ssize_t remaining = byte_offset_impl(axis + 1, shape + 1, strides + 1, index...); + return normalized * *strides + remaining; } /// Create array from any object -- always returns a new reference @@ -1449,7 +1452,8 @@ class array_t : public array { return static_cast(array::mutable_data(index...)); } - // Reference to element at a given index + // Reference to element at a given index. Negative indices count from the end of each + // dimension. Throws if out of bounds. template const T &at(Ix... index) const { if ((ssize_t) sizeof...(index) != ndim()) { @@ -1459,7 +1463,8 @@ class array_t : public array { + byte_offset(ssize_t(index)...) / itemsize()); } - // Mutable reference to element at a given index + // Mutable reference to element at a given index. Negative indices count from the end of each + // dimension. Throws if out of bounds. template T &mutable_at(Ix... index) { if ((ssize_t) sizeof...(index) != ndim()) { diff --git a/tests/test_numpy_array.cpp b/tests/test_numpy_array.cpp index 3bce988353..d3add5b562 100644 --- a/tests/test_numpy_array.cpp +++ b/tests/test_numpy_array.cpp @@ -135,9 +135,10 @@ arr_t &mutate_at_t(arr_t &a, Ix... idx) { #define def_index_fn(name, type) \ sm.def(#name, [](type a) { return name(a); }); \ - sm.def(#name, [](type a, int i) { return name(a, i); }); \ - sm.def(#name, [](type a, int i, int j) { return name(a, i, j); }); \ - sm.def(#name, [](type a, int i, int j, int k) { return name(a, i, j, k); }); + sm.def(#name, [](type a, py::ssize_t i) { return name(a, i); }); \ + sm.def(#name, [](type a, py::ssize_t i, py::ssize_t j) { return name(a, i, j); }); \ + sm.def(#name, \ + [](type a, py::ssize_t i, py::ssize_t j, py::ssize_t k) { return name(a, i, j, k); }); template py::handle auxiliaries(T &&r, T2 &&r2) { diff --git a/tests/test_numpy_array.py b/tests/test_numpy_array.py index e3f29e92a6..b57668e371 100644 --- a/tests/test_numpy_array.py +++ b/tests/test_numpy_array.py @@ -69,6 +69,10 @@ def test_array_attributes(): with pytest.raises(IndexError) as excinfo: m.strides(a, 2) assert str(excinfo.value) == "invalid axis: 2 (ndim = 2)" + assert m.shape(a, -1) == 3 + assert m.shape(a, -2) == 2 + assert m.strides(a, -1) == 2 + assert m.strides(a, -2) == 6 assert not m.writeable(a) assert m.size(a) == 6 assert m.itemsize(a) == 2 @@ -204,6 +208,8 @@ def test_bounds_check(arr): for func in ( m.index_at, m.index_at_t, + m.offset_at, + m.offset_at_t, m.data, m.data_t, m.mutate_data, @@ -217,6 +223,110 @@ def test_bounds_check(arr): with pytest.raises(IndexError) as excinfo: func(arr, 0, 4) assert str(excinfo.value) == "index 4 is out of bounds for axis 1 with size 3" + with pytest.raises(IndexError) as excinfo: + func(arr, -3, 0) + assert str(excinfo.value) == "index -3 is out of bounds for axis 0 with size 2" + with pytest.raises(IndexError) as excinfo: + func(arr, 0, -4) + assert str(excinfo.value) == "index -4 is out of bounds for axis 1 with size 3" + + +@pytest.mark.parametrize("shape", [(), (0,), (2, 0), (2, 3)]) +def test_negative_axes(shape): + a = np.empty(shape, dtype="=u2") + for func, values in ((m.shape, a.shape), (m.strides, a.strides)): + for axis in range(-a.ndim, a.ndim): + assert func(a, axis) == values[axis] + for axis in (-sys.maxsize - 1, -a.ndim - 1, a.ndim, sys.maxsize): + with pytest.raises(IndexError) as excinfo: + func(a, axis) + assert str(excinfo.value) == f"invalid axis: {axis} (ndim = {a.ndim})" + + +@pytest.mark.parametrize( + "indices", [(-1,), (-2,), (-1, 0), (0, -1), (-2, 0), (0, -3), (-2, -3), (-1, -1)] +) +def test_negative_indices(arr, indices): + normalized = tuple(i + n if i < 0 else i for i, n in zip(indices, arr.shape)) + funcs = ( + m.index_at, + m.index_at_t, + m.offset_at, + m.offset_at_t, + m.data, + m.data_t, + m.mutate_data, + m.mutate_data_t, + ) + if len(indices) == arr.ndim: + funcs += (m.at_t, m.mutate_at_t) + for func in funcs: + actual = arr.copy() + expected = arr.copy() + np.testing.assert_array_equal( + func(actual, *indices), func(expected, *normalized) + ) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("shape", [(0,), (2, 0), (2, 3)]) +def test_negative_index_boundaries(shape): + a = np.zeros(shape, dtype="=u2") + for func in ( + m.index_at, + m.index_at_t, + m.offset_at, + m.offset_at_t, + m.data, + m.data_t, + m.mutate_data, + m.mutate_data_t, + m.at_t, + m.mutate_at_t, + ): + for axis, size in enumerate(shape): + for value in (-sys.maxsize - 1, -size - 1, size, sys.maxsize): + indices = [0] * len(shape) + indices[axis] = value + bad_axis = next( + dim + for dim, (i, n) in enumerate(zip(indices, shape)) + if i < -n or i >= n + ) + with pytest.raises(IndexError) as excinfo: + func(a, *indices) + assert str(excinfo.value) == ( + f"index {indices[bad_axis]} is out of bounds for axis {bad_axis} " + f"with size {shape[bad_axis]}" + ) + + +@pytest.mark.parametrize("layout", ["C", "F", "reversed", "broadcast"]) +def test_negative_indices_strided(layout): + a = np.arange(6, dtype="=u2").reshape(2, 3) + if layout == "F": + a = np.asfortranarray(a) + elif layout == "reversed": + a = a[::-1, ::-1] + elif layout == "broadcast": + a = np.broadcast_to(a[:1, :], (2, 3)) + for i in range(-2, 2): + for j in range(-3, 3): + assert m.at_t(a, i, j) == a[i, j] + offset = (i + 2 if i < 0 else i) * a.strides[0] + offset += (j + 3 if j < 0 else j) * a.strides[1] + assert m.offset_at(a, i, j) == offset + assert m.offset_at_t(a, i, j) == offset + assert m.index_at(a, i, j) == offset // a.itemsize + assert m.index_at_t(a, i, j) == offset // a.itemsize + if a.flags.writeable: + expected = int(a[i, j]) + 1 + m.mutate_at_t(a, i, j) + assert a[i, j] == expected + a.flags.writeable = False + for func in (m.mutate_data, m.mutate_data_t, m.mutate_at_t): + with pytest.raises(ValueError, match="array is not writeable"): + func(a, -1, -1) def test_make_c_f_array():