Add tests for nested recarrays

This commit is contained in:
Ivan Smirnov 2016-06-19 16:41:15 +01:00
parent 80a3785a66
commit 7f913aecab
2 changed files with 41 additions and 14 deletions

View File

@ -30,23 +30,36 @@ struct PackedStruct {
struct NestedStruct { struct NestedStruct {
Struct a; Struct a;
PackedStruct b; PackedStruct b;
}; } __attribute__((packed));
template <typename T>
py::array mkarray_via_buffer(size_t n) {
return py::array(py::buffer_info(nullptr, sizeof(T),
py::format_descriptor<T>::value(),
1, { n }, { sizeof(T) }));
}
template <typename S> template <typename S>
py::array_t<S> create_recarray(size_t n) { py::array_t<S> create_recarray(size_t n) {
auto arr = py::array(py::buffer_info(nullptr, sizeof(S), auto arr = mkarray_via_buffer<S>(n);
py::format_descriptor<S>::value(), auto ptr = static_cast<S*>(arr.request().ptr);
1, { n }, { sizeof(S) }));
auto buf = arr.request();
auto ptr = static_cast<S*>(buf.ptr);
for (size_t i = 0; i < n; i++) { for (size_t i = 0; i < n; i++) {
ptr[i].x = i % 2; ptr[i].x = i % 2; ptr[i].y = (uint32_t) i; ptr[i].z = (float) i * 1.5f;
ptr[i].y = i;
ptr[i].z = i * 1.5;
} }
return arr; return arr;
} }
py::array_t<NestedStruct> create_nested(size_t n) {
auto arr = mkarray_via_buffer<NestedStruct>(n);
auto ptr = static_cast<NestedStruct*>(arr.request().ptr);
for (size_t i = 0; i < n; i++) {
ptr[i].a.x = i % 2; ptr[i].a.y = (uint32_t) i; ptr[i].a.z = (float) i * 1.5f;
ptr[i].b.x = (i + 1) % 2; ptr[i].b.y = (uint32_t) (i + 1); ptr[i].b.z = (float) (i + 1) * 1.5f;
}
return arr;
}
void init_ex20(py::module &m) { void init_ex20(py::module &m) {
PYBIND11_DTYPE(Struct, x, y, z); PYBIND11_DTYPE(Struct, x, y, z);
PYBIND11_DTYPE(PackedStruct, x, y, z); PYBIND11_DTYPE(PackedStruct, x, y, z);
@ -54,4 +67,5 @@ void init_ex20(py::module &m) {
m.def("create_rec_simple", &create_recarray<Struct>); m.def("create_rec_simple", &create_recarray<Struct>);
m.def("create_rec_packed", &create_recarray<PackedStruct>); m.def("create_rec_packed", &create_recarray<PackedStruct>);
m.def("create_rec_nested", &create_nested);
} }

View File

@ -2,7 +2,7 @@
from __future__ import print_function from __future__ import print_function
import numpy as np import numpy as np
from example import create_rec_simple, create_rec_packed from example import create_rec_simple, create_rec_packed, create_rec_nested
def check_eq(arr, data, dtype): def check_eq(arr, data, dtype):
@ -14,12 +14,25 @@ simple_dtype = np.dtype({'names': ['x', 'y', 'z'],
packed_dtype = np.dtype([('x', '?'), ('y', 'u4'), ('z', 'f4')]) packed_dtype = np.dtype([('x', '?'), ('y', 'u4'), ('z', 'f4')])
for func, dtype in [(create_rec_simple, simple_dtype), (create_rec_packed, packed_dtype)]: for func, dtype in [(create_rec_simple, simple_dtype), (create_rec_packed, packed_dtype)]:
arr = func(0)
assert arr.dtype == dtype
check_eq(arr, [], simple_dtype)
check_eq(arr, [], packed_dtype)
arr = func(3) arr = func(3)
assert arr.dtype == dtype assert arr.dtype == dtype
check_eq(arr, [(False, 0, 0.0), (True, 1, 1.5), (False, 2, 3.0)], simple_dtype) check_eq(arr, [(False, 0, 0.0), (True, 1, 1.5), (False, 2, 3.0)], simple_dtype)
check_eq(arr, [(False, 0, 0.0), (True, 1, 1.5), (False, 2, 3.0)], packed_dtype) check_eq(arr, [(False, 0, 0.0), (True, 1, 1.5), (False, 2, 3.0)], packed_dtype)
arr = func(0)
assert arr.dtype == dtype nested_dtype = np.dtype([('a', simple_dtype), ('b', packed_dtype)])
check_eq(arr, [], simple_dtype)
check_eq(arr, [], packed_dtype) arr = create_rec_nested(0)
assert arr.dtype == nested_dtype
check_eq(arr, [], nested_dtype)
arr = create_rec_nested(3)
assert arr.dtype == nested_dtype
check_eq(arr, [((False, 0, 0.0), (True, 1, 1.5)),
((True, 1, 1.5), (False, 2, 3.0)),
((False, 2, 3.0), (True, 3, 4.5))], nested_dtype)