| /* |
| tests/test_numpy_dtypes.cpp -- Structured and compound NumPy dtypes |
| |
| Copyright (c) 2016 Ivan Smirnov |
| |
| All rights reserved. Use of this source code is governed by a |
| BSD-style license that can be found in the LICENSE file. |
| */ |
| |
| #include <pybind11/numpy.h> |
| |
| #include "pybind11_tests.h" |
| |
| #ifdef __GNUC__ |
| # define PYBIND11_PACKED(cls) cls __attribute__((__packed__)) |
| #else |
| # define PYBIND11_PACKED(cls) __pragma(pack(push, 1)) cls __pragma(pack(pop)) |
| #endif |
| |
| namespace py = pybind11; |
| |
| struct SimpleStruct { |
| bool bool_; |
| uint32_t uint_; |
| float float_; |
| long double ldbl_; |
| }; |
| |
| std::ostream &operator<<(std::ostream &os, const SimpleStruct &v) { |
| return os << "s:" << v.bool_ << "," << v.uint_ << "," << v.float_ << "," << v.ldbl_; |
| } |
| |
| struct SimpleStructReordered { |
| bool bool_; |
| float float_; |
| uint32_t uint_; |
| long double ldbl_; |
| }; |
| |
| PYBIND11_PACKED(struct PackedStruct { |
| bool bool_; |
| uint32_t uint_; |
| float float_; |
| long double ldbl_; |
| }); |
| |
| std::ostream &operator<<(std::ostream &os, const PackedStruct &v) { |
| return os << "p:" << v.bool_ << "," << v.uint_ << "," << v.float_ << "," << v.ldbl_; |
| } |
| |
| PYBIND11_PACKED(struct NestedStruct { |
| SimpleStruct a; |
| PackedStruct b; |
| }); |
| |
| std::ostream &operator<<(std::ostream &os, const NestedStruct &v) { |
| return os << "n:a=" << v.a << ";b=" << v.b; |
| } |
| |
| struct PartialStruct { |
| bool bool_; |
| uint32_t uint_; |
| float float_; |
| uint64_t dummy2; |
| long double ldbl_; |
| }; |
| |
| struct PartialNestedStruct { |
| uint64_t dummy1; |
| PartialStruct a; |
| uint64_t dummy2; |
| }; |
| |
| struct UnboundStruct {}; |
| |
| struct StringStruct { |
| char a[3]; |
| std::array<char, 3> b; |
| }; |
| |
| struct ComplexStruct { |
| std::complex<float> cflt; |
| std::complex<double> cdbl; |
| }; |
| |
| std::ostream &operator<<(std::ostream &os, const ComplexStruct &v) { |
| return os << "c:" << v.cflt << "," << v.cdbl; |
| } |
| |
| struct ArrayStruct { |
| char a[3][4]; |
| int32_t b[2]; |
| std::array<uint8_t, 3> c; |
| std::array<float, 2> d[4]; |
| }; |
| |
| PYBIND11_PACKED(struct StructWithUglyNames { |
| int8_t __x__; |
| uint64_t __y__; |
| }); |
| |
| enum class E1 : int64_t { A = -1, B = 1 }; |
| enum E2 : uint8_t { X = 1, Y = 2 }; |
| |
| PYBIND11_PACKED(struct EnumStruct { |
| E1 e1; |
| E2 e2; |
| }); |
| |
| std::ostream &operator<<(std::ostream &os, const StringStruct &v) { |
| os << "a='"; |
| for (size_t i = 0; i < 3 && (v.a[i] != 0); i++) { |
| os << v.a[i]; |
| } |
| os << "',b='"; |
| for (size_t i = 0; i < 3 && (v.b[i] != 0); i++) { |
| os << v.b[i]; |
| } |
| return os << "'"; |
| } |
| |
| std::ostream &operator<<(std::ostream &os, const ArrayStruct &v) { |
| os << "a={"; |
| for (int i = 0; i < 3; i++) { |
| if (i > 0) { |
| os << ','; |
| } |
| os << '{'; |
| for (int j = 0; j < 3; j++) { |
| os << v.a[i][j] << ','; |
| } |
| os << v.a[i][3] << '}'; |
| } |
| os << "},b={" << v.b[0] << ',' << v.b[1]; |
| os << "},c={" << int(v.c[0]) << ',' << int(v.c[1]) << ',' << int(v.c[2]); |
| os << "},d={"; |
| for (int i = 0; i < 4; i++) { |
| if (i > 0) { |
| os << ','; |
| } |
| os << '{' << v.d[i][0] << ',' << v.d[i][1] << '}'; |
| } |
| return os << '}'; |
| } |
| |
| std::ostream &operator<<(std::ostream &os, const EnumStruct &v) { |
| return os << "e1=" << (v.e1 == E1::A ? "A" : "B") << ",e2=" << (v.e2 == E2::X ? "X" : "Y"); |
| } |
| |
| template <typename T> |
| py::array mkarray_via_buffer(size_t n) { |
| return py::array(py::buffer_info( |
| nullptr, sizeof(T), py::format_descriptor<T>::format(), 1, {n}, {sizeof(T)})); |
| } |
| |
| #define SET_TEST_VALS(s, i) \ |
| do { \ |
| (s).bool_ = (i) % 2 != 0; \ |
| (s).uint_ = (uint32_t) (i); \ |
| (s).float_ = (float) (i) *1.5f; \ |
| (s).ldbl_ = (long double) (i) * -2.5L; \ |
| } while (0) |
| |
| template <typename S> |
| py::array_t<S, 0> create_recarray(size_t n) { |
| auto arr = mkarray_via_buffer<S>(n); |
| auto req = arr.request(); |
| auto *ptr = static_cast<S *>(req.ptr); |
| for (size_t i = 0; i < n; i++) { |
| SET_TEST_VALS(ptr[i], i); |
| } |
| return arr; |
| } |
| |
| template <typename S> |
| py::list print_recarray(py::array_t<S, 0> arr) { |
| const auto req = arr.request(); |
| auto *const ptr = static_cast<S *>(req.ptr); |
| auto l = py::list(); |
| for (py::ssize_t i = 0; i < req.size; i++) { |
| std::stringstream ss; |
| ss << ptr[i]; |
| l.append(py::str(ss.str())); |
| } |
| return l; |
| } |
| |
| py::array_t<int32_t, 0> test_array_ctors(int i) { |
| using arr_t = py::array_t<int32_t, 0>; |
| |
| std::vector<int32_t> data{1, 2, 3, 4, 5, 6}; |
| std::vector<py::ssize_t> shape{3, 2}; |
| std::vector<py::ssize_t> strides{8, 4}; |
| |
| auto *ptr = data.data(); |
| auto *vptr = (void *) ptr; |
| auto dtype = py::dtype("int32"); |
| |
| py::buffer_info buf_ndim1(vptr, 4, "i", 6); |
| py::buffer_info buf_ndim1_null(nullptr, 4, "i", 6); |
| py::buffer_info buf_ndim2(vptr, 4, "i", 2, shape, strides); |
| py::buffer_info buf_ndim2_null(nullptr, 4, "i", 2, shape, strides); |
| |
| auto fill = [](py::array arr) { |
| auto req = arr.request(); |
| for (int i = 0; i < 6; i++) { |
| ((int32_t *) req.ptr)[i] = i + 1; |
| } |
| return arr; |
| }; |
| |
| switch (i) { |
| // shape: (3, 2) |
| case 10: |
| return arr_t(shape, strides, ptr); |
| case 11: |
| return py::array(shape, strides, ptr); |
| case 12: |
| return py::array(dtype, shape, strides, vptr); |
| case 13: |
| return arr_t(shape, ptr); |
| case 14: |
| return py::array(shape, ptr); |
| case 15: |
| return py::array(dtype, shape, vptr); |
| case 16: |
| return arr_t(buf_ndim2); |
| case 17: |
| return py::array(buf_ndim2); |
| // shape: (3, 2) - post-fill |
| case 20: |
| return fill(arr_t(shape, strides)); |
| case 21: |
| return py::array(shape, strides, ptr); // can't have nullptr due to templated ctor |
| case 22: |
| return fill(py::array(dtype, shape, strides)); |
| case 23: |
| return fill(arr_t(shape)); |
| case 24: |
| return py::array(shape, ptr); // can't have nullptr due to templated ctor |
| case 25: |
| return fill(py::array(dtype, shape)); |
| case 26: |
| return fill(arr_t(buf_ndim2_null)); |
| case 27: |
| return fill(py::array(buf_ndim2_null)); |
| // shape: (6, ) |
| case 30: |
| return arr_t(6, ptr); |
| case 31: |
| return py::array(6, ptr); |
| case 32: |
| return py::array(dtype, 6, vptr); |
| case 33: |
| return arr_t(buf_ndim1); |
| case 34: |
| return py::array(buf_ndim1); |
| // shape: (6, ) |
| case 40: |
| return fill(arr_t(6)); |
| case 41: |
| return py::array(6, ptr); // can't have nullptr due to templated ctor |
| case 42: |
| return fill(py::array(dtype, 6)); |
| case 43: |
| return fill(arr_t(buf_ndim1_null)); |
| case 44: |
| return fill(py::array(buf_ndim1_null)); |
| } |
| return arr_t(); |
| } |
| |
| py::list test_dtype_ctors() { |
| py::list list; |
| list.append(py::dtype("int32")); |
| list.append(py::dtype(std::string("float64"))); |
| list.append(py::dtype::from_args(py::str("bool"))); |
| py::list names, offsets, formats; |
| py::dict dict; |
| names.append(py::str("a")); |
| names.append(py::str("b")); |
| dict["names"] = names; |
| offsets.append(py::int_(1)); |
| offsets.append(py::int_(10)); |
| dict["offsets"] = offsets; |
| formats.append(py::dtype("int32")); |
| formats.append(py::dtype("float64")); |
| dict["formats"] = formats; |
| dict["itemsize"] = py::int_(20); |
| list.append(py::dtype::from_args(dict)); |
| list.append(py::dtype(names, formats, offsets, 20)); |
| list.append(py::dtype(py::buffer_info((void *) nullptr, sizeof(unsigned int), "I", 1))); |
| list.append(py::dtype(py::buffer_info((void *) nullptr, 0, "T{i:a:f:b:}", 1))); |
| list.append(py::dtype(py::detail::npy_api::NPY_DOUBLE_)); |
| return list; |
| } |
| |
| struct A {}; |
| struct B {}; |
| |
| TEST_SUBMODULE(numpy_dtypes, m) { |
| try { |
| py::module_::import("numpy"); |
| } catch (const py::error_already_set &) { |
| return; |
| } |
| |
| // typeinfo may be registered before the dtype descriptor for scalar casts to work... |
| py::class_<SimpleStruct>(m, "SimpleStruct") |
| // Explicit construct to ensure zero-valued initialization. |
| .def(py::init([]() { return SimpleStruct(); })) |
| .def_readwrite("bool_", &SimpleStruct::bool_) |
| .def_readwrite("uint_", &SimpleStruct::uint_) |
| .def_readwrite("float_", &SimpleStruct::float_) |
| .def_readwrite("ldbl_", &SimpleStruct::ldbl_) |
| .def("astuple", |
| [](const SimpleStruct &self) { |
| return py::make_tuple(self.bool_, self.uint_, self.float_, self.ldbl_); |
| }) |
| .def_static("fromtuple", [](const py::tuple &tup) { |
| if (py::len(tup) != 4) { |
| throw py::cast_error("Invalid size"); |
| } |
| return SimpleStruct{tup[0].cast<bool>(), |
| tup[1].cast<uint32_t>(), |
| tup[2].cast<float>(), |
| tup[3].cast<long double>()}; |
| }); |
| |
| PYBIND11_NUMPY_DTYPE(SimpleStruct, bool_, uint_, float_, ldbl_); |
| PYBIND11_NUMPY_DTYPE(SimpleStructReordered, bool_, uint_, float_, ldbl_); |
| PYBIND11_NUMPY_DTYPE(PackedStruct, bool_, uint_, float_, ldbl_); |
| PYBIND11_NUMPY_DTYPE(NestedStruct, a, b); |
| PYBIND11_NUMPY_DTYPE(PartialStruct, bool_, uint_, float_, ldbl_); |
| PYBIND11_NUMPY_DTYPE(PartialNestedStruct, a); |
| PYBIND11_NUMPY_DTYPE(StringStruct, a, b); |
| PYBIND11_NUMPY_DTYPE(ArrayStruct, a, b, c, d); |
| PYBIND11_NUMPY_DTYPE(EnumStruct, e1, e2); |
| PYBIND11_NUMPY_DTYPE(ComplexStruct, cflt, cdbl); |
| |
| // ... or after |
| py::class_<PackedStruct>(m, "PackedStruct"); |
| |
| PYBIND11_NUMPY_DTYPE_EX(StructWithUglyNames, __x__, "x", __y__, "y"); |
| |
| #ifdef PYBIND11_NEVER_DEFINED_EVER |
| // If enabled, this should produce a static_assert failure telling the user that the struct |
| // is not a POD type |
| struct NotPOD { |
| std::string v; |
| NotPOD() : v("hi"){}; |
| }; |
| PYBIND11_NUMPY_DTYPE(NotPOD, v); |
| #endif |
| |
| // Check that dtypes can be registered programmatically, both from |
| // initializer lists of field descriptors and from other containers. |
| py::detail::npy_format_descriptor<A>::register_dtype({}); |
| py::detail::npy_format_descriptor<B>::register_dtype( |
| std::vector<py::detail::field_descriptor>{}); |
| |
| // test_recarray, test_scalar_conversion |
| m.def("create_rec_simple", &create_recarray<SimpleStruct>); |
| m.def("create_rec_packed", &create_recarray<PackedStruct>); |
| m.def("create_rec_nested", [](size_t n) { // test_signature |
| py::array_t<NestedStruct, 0> arr = mkarray_via_buffer<NestedStruct>(n); |
| auto req = arr.request(); |
| auto *ptr = static_cast<NestedStruct *>(req.ptr); |
| for (size_t i = 0; i < n; i++) { |
| SET_TEST_VALS(ptr[i].a, i); |
| SET_TEST_VALS(ptr[i].b, i + 1); |
| } |
| return arr; |
| }); |
| m.def("create_rec_partial", &create_recarray<PartialStruct>); |
| m.def("create_rec_partial_nested", [](size_t n) { |
| py::array_t<PartialNestedStruct, 0> arr = mkarray_via_buffer<PartialNestedStruct>(n); |
| auto req = arr.request(); |
| auto *ptr = static_cast<PartialNestedStruct *>(req.ptr); |
| for (size_t i = 0; i < n; i++) { |
| SET_TEST_VALS(ptr[i].a, i); |
| } |
| return arr; |
| }); |
| m.def("print_rec_simple", &print_recarray<SimpleStruct>); |
| m.def("print_rec_packed", &print_recarray<PackedStruct>); |
| m.def("print_rec_nested", &print_recarray<NestedStruct>); |
| |
| // test_format_descriptors |
| m.def("get_format_unbound", []() { return py::format_descriptor<UnboundStruct>::format(); }); |
| m.def("print_format_descriptors", []() { |
| py::list l; |
| for (const auto &fmt : {py::format_descriptor<SimpleStruct>::format(), |
| py::format_descriptor<PackedStruct>::format(), |
| py::format_descriptor<NestedStruct>::format(), |
| py::format_descriptor<PartialStruct>::format(), |
| py::format_descriptor<PartialNestedStruct>::format(), |
| py::format_descriptor<StringStruct>::format(), |
| py::format_descriptor<ArrayStruct>::format(), |
| py::format_descriptor<EnumStruct>::format(), |
| py::format_descriptor<ComplexStruct>::format()}) { |
| l.append(py::cast(fmt)); |
| } |
| return l; |
| }); |
| |
| // test_dtype |
| std::vector<const char *> dtype_names{ |
| "byte", "short", "intc", "int_", "longlong", "ubyte", "ushort", |
| "uintc", "uint", "ulonglong", "half", "single", "double", "longdouble", |
| "csingle", "cdouble", "clongdouble", "bool_", "datetime64", "timedelta64", "object_"}; |
| |
| m.def("print_dtypes", []() { |
| py::list l; |
| for (const py::handle &d : {py::dtype::of<SimpleStruct>(), |
| py::dtype::of<PackedStruct>(), |
| py::dtype::of<NestedStruct>(), |
| py::dtype::of<PartialStruct>(), |
| py::dtype::of<PartialNestedStruct>(), |
| py::dtype::of<StringStruct>(), |
| py::dtype::of<ArrayStruct>(), |
| py::dtype::of<EnumStruct>(), |
| py::dtype::of<StructWithUglyNames>(), |
| py::dtype::of<ComplexStruct>()}) { |
| l.append(py::str(d)); |
| } |
| return l; |
| }); |
| m.def("test_dtype_ctors", &test_dtype_ctors); |
| m.def("test_dtype_kind", [dtype_names]() { |
| py::list list; |
| for (const auto &dt_name : dtype_names) { |
| list.append(py::dtype(dt_name).kind()); |
| } |
| return list; |
| }); |
| m.def("test_dtype_char_", [dtype_names]() { |
| py::list list; |
| for (const auto &dt_name : dtype_names) { |
| list.append(py::dtype(dt_name).char_()); |
| } |
| return list; |
| }); |
| m.def("test_dtype_num", [dtype_names]() { |
| py::list list; |
| for (const auto &dt_name : dtype_names) { |
| list.append(py::dtype(dt_name).num()); |
| } |
| return list; |
| }); |
| m.def("test_dtype_byteorder", [dtype_names]() { |
| py::list list; |
| for (const auto &dt_name : dtype_names) { |
| list.append(py::dtype(dt_name).byteorder()); |
| } |
| return list; |
| }); |
| m.def("test_dtype_alignment", [dtype_names]() { |
| py::list list; |
| for (const auto &dt_name : dtype_names) { |
| list.append(py::dtype(dt_name).alignment()); |
| } |
| return list; |
| }); |
| m.def("test_dtype_flags", [dtype_names]() { |
| py::list list; |
| for (const auto &dt_name : dtype_names) { |
| list.append(py::dtype(dt_name).flags()); |
| } |
| return list; |
| }); |
| m.def("test_dtype_methods", []() { |
| py::list list; |
| auto dt1 = py::dtype::of<int32_t>(); |
| auto dt2 = py::dtype::of<SimpleStruct>(); |
| list.append(dt1); |
| list.append(dt2); |
| list.append(py::bool_(dt1.has_fields())); |
| list.append(py::bool_(dt2.has_fields())); |
| list.append(py::int_(dt1.itemsize())); |
| list.append(py::int_(dt2.itemsize())); |
| return list; |
| }); |
| struct TrailingPaddingStruct { |
| int32_t a; |
| char b; |
| }; |
| PYBIND11_NUMPY_DTYPE(TrailingPaddingStruct, a, b); |
| m.def("trailing_padding_dtype", []() { return py::dtype::of<TrailingPaddingStruct>(); }); |
| |
| // test_string_array |
| m.def("create_string_array", [](bool non_empty) { |
| py::array_t<StringStruct, 0> arr = mkarray_via_buffer<StringStruct>(non_empty ? 4 : 0); |
| if (non_empty) { |
| auto req = arr.request(); |
| auto *ptr = static_cast<StringStruct *>(req.ptr); |
| for (py::ssize_t i = 0; i < req.size * req.itemsize; i++) { |
| static_cast<char *>(req.ptr)[i] = 0; |
| } |
| ptr[1].a[0] = 'a'; |
| ptr[1].b[0] = 'a'; |
| ptr[2].a[0] = 'a'; |
| ptr[2].b[0] = 'a'; |
| ptr[3].a[0] = 'a'; |
| ptr[3].b[0] = 'a'; |
| |
| ptr[2].a[1] = 'b'; |
| ptr[2].b[1] = 'b'; |
| ptr[3].a[1] = 'b'; |
| ptr[3].b[1] = 'b'; |
| |
| ptr[3].a[2] = 'c'; |
| ptr[3].b[2] = 'c'; |
| } |
| return arr; |
| }); |
| m.def("print_string_array", &print_recarray<StringStruct>); |
| |
| // test_array_array |
| m.def("create_array_array", [](size_t n) { |
| py::array_t<ArrayStruct, 0> arr = mkarray_via_buffer<ArrayStruct>(n); |
| auto *ptr = (ArrayStruct *) arr.mutable_data(); |
| for (size_t i = 0; i < n; i++) { |
| for (size_t j = 0; j < 3; j++) { |
| for (size_t k = 0; k < 4; k++) { |
| ptr[i].a[j][k] = char('A' + (i * 100 + j * 10 + k) % 26); |
| } |
| } |
| for (size_t j = 0; j < 2; j++) { |
| ptr[i].b[j] = int32_t(i * 1000 + j); |
| } |
| for (size_t j = 0; j < 3; j++) { |
| ptr[i].c[j] = uint8_t(i * 10 + j); |
| } |
| for (size_t j = 0; j < 4; j++) { |
| for (size_t k = 0; k < 2; k++) { |
| ptr[i].d[j][k] = float(i) * 100.0f + float(j) * 10.0f + float(k); |
| } |
| } |
| } |
| return arr; |
| }); |
| m.def("print_array_array", &print_recarray<ArrayStruct>); |
| |
| // test_enum_array |
| m.def("create_enum_array", [](size_t n) { |
| py::array_t<EnumStruct, 0> arr = mkarray_via_buffer<EnumStruct>(n); |
| auto *ptr = (EnumStruct *) arr.mutable_data(); |
| for (size_t i = 0; i < n; i++) { |
| ptr[i].e1 = static_cast<E1>(-1 + ((int) i % 2) * 2); |
| ptr[i].e2 = static_cast<E2>(1 + (i % 2)); |
| } |
| return arr; |
| }); |
| m.def("print_enum_array", &print_recarray<EnumStruct>); |
| |
| // test_complex_array |
| m.def("create_complex_array", [](size_t n) { |
| py::array_t<ComplexStruct, 0> arr = mkarray_via_buffer<ComplexStruct>(n); |
| auto *ptr = (ComplexStruct *) arr.mutable_data(); |
| for (size_t i = 0; i < n; i++) { |
| ptr[i].cflt.real(float(i)); |
| ptr[i].cflt.imag(float(i) + 0.25f); |
| ptr[i].cdbl.real(double(i) + 0.5); |
| ptr[i].cdbl.imag(double(i) + 0.75); |
| } |
| return arr; |
| }); |
| m.def("print_complex_array", &print_recarray<ComplexStruct>); |
| |
| // test_array_constructors |
| m.def("test_array_ctors", &test_array_ctors); |
| |
| // test_compare_buffer_info |
| struct CompareStruct { |
| bool x; |
| uint32_t y; |
| float z; |
| }; |
| PYBIND11_NUMPY_DTYPE(CompareStruct, x, y, z); |
| m.def("compare_buffer_info", []() { |
| py::list list; |
| list.append(py::bool_(py::detail::compare_buffer_info<float>::compare( |
| py::buffer_info(nullptr, sizeof(float), "f", 1)))); |
| list.append(py::bool_(py::detail::compare_buffer_info<unsigned>::compare( |
| py::buffer_info(nullptr, sizeof(int), "I", 1)))); |
| list.append(py::bool_(py::detail::compare_buffer_info<long>::compare( |
| py::buffer_info(nullptr, sizeof(long), "l", 1)))); |
| list.append(py::bool_(py::detail::compare_buffer_info<long>::compare( |
| py::buffer_info(nullptr, sizeof(long), sizeof(long) == sizeof(int) ? "i" : "q", 1)))); |
| list.append(py::bool_(py::detail::compare_buffer_info<CompareStruct>::compare( |
| py::buffer_info(nullptr, sizeof(CompareStruct), "T{?:x:3xI:y:f:z:}", 1)))); |
| return list; |
| }); |
| m.def("buffer_to_dtype", [](py::buffer &buf) { return py::dtype(buf.request()); }); |
| |
| // test_scalar_conversion |
| auto f_simple = [](SimpleStruct s) { return s.uint_ * 10; }; |
| m.def("f_simple", f_simple); |
| m.def("f_packed", [](PackedStruct s) { return s.uint_ * 10; }); |
| m.def("f_nested", [](NestedStruct s) { return s.a.uint_ * 10; }); |
| |
| // test_vectorize |
| m.def("f_simple_vectorized", py::vectorize(f_simple)); |
| auto f_simple_pass_thru = [](SimpleStruct s) { return s; }; |
| m.def("f_simple_pass_thru_vectorized", py::vectorize(f_simple_pass_thru)); |
| |
| // test_register_dtype |
| m.def("register_dtype", |
| []() { PYBIND11_NUMPY_DTYPE(SimpleStruct, bool_, uint_, float_, ldbl_); }); |
| |
| // test_str_leak |
| m.def("dtype_wrapper", [](const py::object &d) { return py::dtype::from_args(d); }); |
| } |