diff --git a/tests/test_numpy_array.cpp b/tests/test_numpy_array.cpp index ed118a0ae..0614f5717 100644 --- a/tests/test_numpy_array.cpp +++ b/tests/test_numpy_array.cpp @@ -8,38 +8,87 @@ */ #include "pybind11_tests.h" + #include #include +#include +#include + +using arr = py::array; +using arr_t = py::array_t; + +template arr data(const arr& a, Ix&&... index) { + return arr(a.nbytes() - a.offset_at(index...), (const uint8_t *) a.data(index...)); +} + +template 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 = (uint8_t *) 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 arr& mutate_data(arr& a, Ix&&... index) { + auto ptr = (uint8_t *) a.mutable_data(index...); + for (size_t i = 0; i < a.nbytes() - a.offset_at(index...); i++) + ptr[i] = (uint8_t) (ptr[i] * 2); + return a; +} + +template 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; +} + +template size_t index_at(const arr& a, Ix&&... idx) { return a.index_at(idx...); } +template size_t index_at_t(const arr_t& a, Ix&&... idx) { return a.index_at(idx...); } +template size_t offset_at(const arr& a, Ix&&... idx) { return a.offset_at(idx...); } +template size_t offset_at_t(const arr_t& a, Ix&&... idx) { return a.offset_at(idx...); } +template size_t at_t(const arr_t& a, Ix&&... idx) { return a.at(idx...); } +template arr_t& mutate_at_t(arr_t& a, Ix&&... idx) { a.mutable_at(idx...)++; return a; } + +#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); }); + 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(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(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(); }); + + def_index_fn(data, const arr&); + def_index_fn(data_t, const arr_t&); + def_index_fn(index_at, const arr&); + def_index_fn(index_at_t, const arr_t&); + def_index_fn(offset_at, const arr&); + def_index_fn(offset_at_t, const arr_t&); + def_index_fn(mutate_data, arr&); + def_index_fn(mutate_data_t, arr_t&); + def_index_fn(at_t, const arr_t&); + def_index_fn(mutate_at_t, arr_t&); }); diff --git a/tests/test_numpy_array.py b/tests/test_numpy_array.py index 929cdc4d9..4a6af5ef6 100644 --- a/tests/test_numpy_array.py +++ b/tests/test_numpy_array.py @@ -4,40 +4,147 @@ with pytest.suppress(ImportError): import numpy as np +@pytest.fixture(scope='function') +def arr(): + return np.array([[1, 2, 3], [4, 5, 6]], '