Add tests for array data access /index methods

Author: aldanor

tests/test_numpy_array.cpp

@@ -8,38 +8,97 @@
 */
 
 #include "pybind11_tests.h"
+
 #include <pybind11/numpy.h>
 #include <pybind11/stl.h>
 
+#include <cstdint>
+#include <vector>
+
+using arr = py::array;
+using arr_t = py::array_t<uint16_t, 0>;
+
+template<typename... Ix> arr data(const arr& a, Ix&&... index) {
+    return arr(a.nbytes() - a.index_at(index...), a.data(index...));
+}
+
+template<typename... Ix> arr data_t(const arr_t& a, Ix&&... index) {
+    return arr(a.size() - a.index_at(index...), a.data(index...));
+}
+
+arr& mutate_data(arr& a) {
+    auto ptr = a.mutable_data();
+    for (size_t i = 0; i < a.nbytes(); i++)
+        ptr[i] = (uint8_t) (ptr[i] * 2);
+    return a;
+}
+
+arr_t& mutate_data_t(arr_t& a) {
+    auto ptr = a.mutable_data();
+    for (size_t i = 0; i < a.size(); i++)
+        ptr[i]++;
+    return a;
+}
+
+template<typename... Ix> arr& mutate_data(arr& a, Ix&&... index) {
+    auto ptr = a.mutable_data(index...);
+    for (size_t i = 0; i < a.nbytes() - a.index_at(index...); i++)
+        ptr[i] = (uint8_t) (ptr[i] * 2);
+    return a;
+}
+
+template<typename... Ix> arr_t& mutate_data_t(arr_t& a, Ix&&... index) {
+    auto ptr = a.mutable_data(index...);
+    for (size_t i = 0; i < a.size() - a.index_at(index...); i++)
+        ptr[i]++;
+    return a;
+}
+
 test_initializer numpy_array([](py::module &m) {
-    m.def("get_arr_ndim", [](const py::array& arr) {
-        return arr.ndim();
-    });
-    m.def("get_arr_shape", [](const py::array& arr) {
-        return std::vector<size_t>(arr.shape(), arr.shape() + arr.ndim());
-    });
-    m.def("get_arr_shape", [](const py::array& arr, size_t dim) {
-        return arr.shape(dim);
-    });
-    m.def("get_arr_strides", [](const py::array& arr) {
-        return std::vector<size_t>(arr.strides(), arr.strides() + arr.ndim());
-    });
-    m.def("get_arr_strides", [](const py::array& arr, size_t dim) {
-        return arr.strides(dim);
-    });
-    m.def("get_arr_writeable", [](const py::array& arr) {
-        return arr.writeable();
-    });
-    m.def("get_arr_size", [](const py::array& arr) {
-        return arr.size();
-    });
-    m.def("get_arr_itemsize", [](const py::array& arr) {
-        return arr.itemsize();
-    });
-    m.def("get_arr_nbytes", [](const py::array& arr) {
-        return arr.nbytes();
-    });
-    m.def("get_arr_owndata", [](const py::array& arr) {
-        return arr.owndata();
-    });
+    auto sm = m.def_submodule("array");
+
+    sm.def("ndim", [](const arr& a) { return a.ndim(); });
+    sm.def("shape", [](const arr& a) { return arr(a.ndim(), a.shape()); });
+    sm.def("shape", [](const arr& a, size_t dim) { return a.shape(dim); });
+    sm.def("strides", [](const arr& a) { return arr(a.ndim(), a.strides()); });
+    sm.def("strides", [](const arr& a, size_t dim) { return a.strides(dim); });
+    sm.def("writeable", [](const arr& a) { return a.writeable(); });
+    sm.def("size", [](const arr& a) { return a.size(); });
+    sm.def("itemsize", [](const arr& a) { return a.itemsize(); });
+    sm.def("nbytes", [](const arr& a) { return a.nbytes(); });
+    sm.def("owndata", [](const arr& a) { return a.owndata(); });
+    sm.def("index_at", [](const arr& a) { return a.index_at(); });
+    sm.def("index_at", [](const arr& a, int i) { return a.index_at(i); });
+    sm.def("index_at", [](const arr& a, int i, int j) { return a.index_at(i, j); });
+    sm.def("index_at", [](const arr& a, int i, int j, int k) { return a.index_at(i, j, k); });
+    sm.def("data", [](const arr& a) { return data(a); });
+    sm.def("data", [](const arr& a, int i) { return data(a, i); });
+    sm.def("data", [](const arr& a, int i, int j) { return data(a, i, j); });
+    sm.def("data", [](const arr& a, int i, int j, int k) { return data(a, i, j, k); });
+    sm.def("mutate_data", [](arr& a) { return mutate_data(a); });
+    sm.def("mutate_data", [](arr& a) { return mutate_data(a); });
+    sm.def("mutate_data", [](arr& a, int i) { return mutate_data(a, i); });
+    sm.def("mutate_data", [](arr& a, int i, int j) { return mutate_data(a, i, j); });
+    sm.def("mutate_data", [](arr& a, int i, int j, int k) { return mutate_data(a, i, j, k); });
+    sm.def("index_at_t", [](const arr_t& a) { return a.index_at(); });
+    sm.def("index_at_t", [](const arr_t& a, int i) { return a.index_at(i); });
+    sm.def("index_at_t", [](const arr_t& a, int i, int j) { return a.index_at(i, j); });
+    sm.def("index_at_t", [](const arr_t& a, int i, int j, int k) { return a.index_at(i, j, k); });
+    sm.def("data_t", [](const arr_t& a) { return data_t(a); });
+    sm.def("data_t", [](const arr_t& a, int i) { return data_t(a, i); });
+    sm.def("data_t", [](const arr_t& a, int i, int j) { return data_t(a, i, j); });
+    sm.def("data_t", [](const arr_t& a, int i, int j, int k) { return data_t(a, i, j, k); });
+    sm.def("mutate_data_t", [](arr_t& a) { return mutate_data_t(a); });
+    sm.def("mutate_data_t", [](arr_t& a) { return mutate_data_t(a); });
+    sm.def("mutate_data_t", [](arr_t& a, int i) { return mutate_data_t(a, i); });
+    sm.def("mutate_data_t", [](arr_t& a, int i, int j) { return mutate_data_t(a, i, j); });
+    sm.def("mutate_data_t", [](arr_t& a, int i, int j, int k) { return mutate_data_t(a, i, j, k); });
+    sm.def("at_t", [](const arr_t& a) { return a.at(); });
+    sm.def("at_t", [](const arr_t& a, int i) { return a.at(i); });
+    sm.def("at_t", [](const arr_t& a, int i, int j) { return a.at(i, j); });
+    sm.def("at_t", [](const arr_t& a, int i, int j, int k) { return a.at(i, j, k); });
+    sm.def("mutate_at_t", [](arr_t& a) { a.mutable_at()++; return a; });
+    sm.def("mutate_at_t", [](arr_t& a, int i) { a.mutable_at(i)++; return a; });
+    sm.def("mutate_at_t", [](arr_t& a, int i, int j) { a.mutable_at(i, j)++; return a; });
+    sm.def("mutate_at_t", [](arr_t& a, int i, int j, int k) { a.mutable_at(i, j, k)++; return a; });
 });

tests/test_numpy_array.py

@@ -4,40 +4,172 @@
     import numpy as np
 
 
+@pytest.fixture(scope='function')
+def arr():
+    return np.array([[1, 2, 3], [4, 5, 6]], '<u2')
+
+
 @pytest.requires_numpy
 def test_array_attributes():
-    from pybind11_tests import (get_arr_ndim, get_arr_shape, get_arr_strides, get_arr_writeable,
-                                get_arr_size, get_arr_itemsize, get_arr_nbytes, get_arr_owndata)
+    from pybind11_tests.array import (
+        ndim, shape, strides, writeable, size, itemsize, nbytes, owndata
+    )
 
     a = np.array(0, 'f8')
-    assert get_arr_ndim(a) == 0
-    assert get_arr_shape(a) == []
-    assert get_arr_strides(a) == []
-    with pytest.raises(RuntimeError):
-        get_arr_shape(a, 1)
-    with pytest.raises(RuntimeError):
-        get_arr_strides(a, 0)
-    assert get_arr_writeable(a)
-    assert get_arr_size(a) == 1
-    assert get_arr_itemsize(a) == 8
-    assert get_arr_nbytes(a) == 8
-    assert get_arr_owndata(a)
+    assert ndim(a) == 0
+    assert all(shape(a) == [])
+    assert all(strides(a) == [])
+    with pytest.raises(IndexError) as excinfo:
+        shape(a, 0)
+    assert str(excinfo.value) == 'invalid axis: 0 (ndim = 0)'
+    with pytest.raises(IndexError) as excinfo:
+        strides(a, 0)
+    assert str(excinfo.value) == 'invalid axis: 0 (ndim = 0)'
+    assert writeable(a)
+    assert size(a) == 1
+    assert itemsize(a) == 8
+    assert nbytes(a) == 8
+    assert owndata(a)
 
     a = np.array([[1, 2, 3], [4, 5, 6]], 'u2').view()
     a.flags.writeable = False
-    assert get_arr_ndim(a) == 2
-    assert get_arr_shape(a) == [2, 3]
-    assert get_arr_shape(a, 0) == 2
-    assert get_arr_shape(a, 1) == 3
-    assert get_arr_strides(a) == [6, 2]
-    assert get_arr_strides(a, 0) == 6
-    assert get_arr_strides(a, 1) == 2
+    assert ndim(a) == 2
+    assert all(shape(a) == [2, 3])
+    assert shape(a, 0) == 2
+    assert shape(a, 1) == 3
+    assert all(strides(a) == [6, 2])
+    assert strides(a, 0) == 6
+    assert strides(a, 1) == 2
+    with pytest.raises(IndexError) as excinfo:
+        shape(a, 2)
+    assert str(excinfo.value) == 'invalid axis: 2 (ndim = 2)'
+    with pytest.raises(IndexError) as excinfo:
+        strides(a, 2)
+    assert str(excinfo.value) == 'invalid axis: 2 (ndim = 2)'
+    assert not writeable(a)
+    assert size(a) == 6
+    assert itemsize(a) == 2
+    assert nbytes(a) == 12
+    assert not owndata(a)
+
+
+@pytest.requires_numpy
+def test_index_at(arr):
+    from pybind11_tests.array import index_at, index_at_t
+
+    assert index_at(arr) == 0
+    assert index_at(arr, 0) == 0
+    assert index_at(arr, 1) == 6
+    assert index_at(arr, 0, 1) == 2
+    assert index_at(arr, 1, 2) == 10
+    with pytest.raises(IndexError) as excinfo:
+        index_at(arr, 1, 2, 3)
+    assert str(excinfo.value) == 'invalid axis: 2 (ndim = 2)'
+
+    assert index_at_t(arr) == 0
+    assert index_at_t(arr, 0) == 0
+    assert index_at_t(arr, 1) == 3
+    assert index_at_t(arr, 0, 1) == 1
+    assert index_at_t(arr, 1, 2) == 5
+    with pytest.raises(IndexError) as excinfo:
+        index_at_t(arr, 1, 2, 3)
+    assert str(excinfo.value) == 'invalid axis: 2 (ndim = 2)'
+
+
+@pytest.requires_numpy
+def test_data(arr):
+    from pybind11_tests.array import data, data_t
+
+    assert all(data(arr) == [1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0])
+    assert all(data(arr) == [1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0])
+    assert all(data(arr, 1) == [4, 0, 5, 0, 6, 0])
+    assert all(data(arr, 0, 1) == [2, 0, 3, 0, 4, 0, 5, 0, 6, 0])
+    assert all(data(arr, 1, 2) == [6, 0])
+    with pytest.raises(IndexError):
+        data(arr, 1, 2, 3)
+
+    assert all(data_t(arr) == [1, 2, 3, 4, 5, 6])
+    assert all(data_t(arr) == [1, 2, 3, 4, 5, 6])
+    assert all(data_t(arr, 1) == [4, 5, 6])
+    assert all(data_t(arr, 0, 1) == [2, 3, 4, 5, 6])
+    assert all(data_t(arr, 1, 2) == [6])
+    with pytest.raises(IndexError):
+        data_t(arr, 1, 2, 3)
+
+
+@pytest.requires_numpy
+def test_mutate_reaadonly(arr):
+    from pybind11_tests.array import mutate_data, mutate_data_t, mutate_at_t
+    arr.flags.writeable = False
+
     with pytest.raises(RuntimeError):
-        get_arr_shape(a, 2)
+        mutate_data(arr)
     with pytest.raises(RuntimeError):
-        get_arr_strides(a, 2)
-    assert not get_arr_writeable(a)
-    assert get_arr_size(a) == 6
-    assert get_arr_itemsize(a) == 2
-    assert get_arr_nbytes(a) == 12
-    assert not get_arr_owndata(a)
+        mutate_data_t(arr)
+    with pytest.raises(RuntimeError):
+        mutate_at_t(arr, 0, 0)
+
+
+@pytest.requires_numpy
+def test_at(arr):
+    from pybind11_tests.array import at_t, mutate_at_t
+
+    with pytest.raises(IndexError) as excinfo:
+        at_t(arr)
+    assert str(excinfo.value) == 'invalid index shape: 0 (ndim = 2)'
+    with pytest.raises(IndexError) as excinfo:
+        at_t(arr, 1)
+    assert str(excinfo.value) == 'invalid index shape: 1 (ndim = 2)'
+    with pytest.raises(IndexError) as excinfo:
+        at_t(arr, 1, 2, 3)
+    assert str(excinfo.value) == 'invalid index shape: 3 (ndim = 2)'
+    assert at_t(arr, 0, 2) == 3
+    assert at_t(arr, 1, 0) == 4
+
+    with pytest.raises(IndexError) as excinfo:
+        mutate_at_t(arr)
+    assert str(excinfo.value) == 'invalid index shape: 0 (ndim = 2)'
+    with pytest.raises(IndexError) as excinfo:
+        mutate_at_t(arr, 1)
+    assert str(excinfo.value) == 'invalid index shape: 1 (ndim = 2)'
+    with pytest.raises(IndexError) as excinfo:
+        mutate_at_t(arr, 1, 2, 3)
+    assert str(excinfo.value) == 'invalid index shape: 3 (ndim = 2)'
+    assert all(mutate_at_t(arr, 0, 2).ravel() == [1, 2, 4, 4, 5, 6])
+    assert all(mutate_at_t(arr, 1, 0).ravel() == [1, 2, 4, 5, 5, 6])
+
+
+@pytest.requires_numpy
+def test_mutate_data(arr):
+    from pybind11_tests.array import mutate_data, mutate_data_t
+
+    assert all(mutate_data(arr).ravel() == [2, 4, 6, 8, 10, 12])
+    assert all(mutate_data(arr).ravel() == [4, 8, 12, 16, 20, 24])
+    assert all(mutate_data(arr, 1).ravel() == [4, 8, 12, 32, 40, 48])
+    assert all(mutate_data(arr, 0, 1).ravel() == [4, 16, 24, 64, 80, 96])
+    assert all(mutate_data(arr, 1, 2).ravel() == [4, 16, 24, 64, 80, 192])
+    with pytest.raises(IndexError):
+        mutate_data(arr, 1, 2, 3)
+
+    assert all(mutate_data_t(arr).ravel() == [5, 17, 25, 65, 81, 193])
+    assert all(mutate_data_t(arr).ravel() == [6, 18, 26, 66, 82, 194])
+    assert all(mutate_data_t(arr, 1).ravel() == [6, 18, 26, 67, 83, 195])
+    assert all(mutate_data_t(arr, 0, 1).ravel() == [6, 19, 27, 68, 84, 196])
+    assert all(mutate_data_t(arr, 1, 2).ravel() == [6, 19, 27, 68, 84, 197])
+    with pytest.raises(IndexError):
+        mutate_data_t(arr, 1, 2, 3)
+
+
+@pytest.requires_numpy
+def test_bounds_check(arr):
+    from pybind11_tests.array import (index_at, index_at_t, data, data_t,
+                                      mutate_data, mutate_data_t, at_t, mutate_at_t)
+    funcs = (index_at, index_at_t, data, data_t,
+             mutate_data, mutate_data_t, at_t, mutate_at_t)
+    for func in funcs:
+        with pytest.raises(IndexError) as excinfo:
+            index_at(arr, 2, 0)
+        assert str(excinfo.value) == 'out of bounds access'
+        with pytest.raises(IndexError) as excinfo:
+            index_at(arr, 0, 3)
+        assert str(excinfo.value) == 'out of bounds access'