test_numpy_dtypes.cpp revision 11986:c12e4625ab56
1/* 2 tests/test_numpy_dtypes.cpp -- Structured and compound NumPy dtypes 3 4 Copyright (c) 2016 Ivan Smirnov 5 6 All rights reserved. Use of this source code is governed by a 7 BSD-style license that can be found in the LICENSE file. 8*/ 9 10#include "pybind11_tests.h" 11#include <pybind11/numpy.h> 12 13#ifdef __GNUC__ 14#define PYBIND11_PACKED(cls) cls __attribute__((__packed__)) 15#else 16#define PYBIND11_PACKED(cls) __pragma(pack(push, 1)) cls __pragma(pack(pop)) 17#endif 18 19namespace py = pybind11; 20 21struct SimpleStruct { 22 bool x; 23 uint32_t y; 24 float z; 25}; 26 27std::ostream& operator<<(std::ostream& os, const SimpleStruct& v) { 28 return os << "s:" << v.x << "," << v.y << "," << v.z; 29} 30 31PYBIND11_PACKED(struct PackedStruct { 32 bool x; 33 uint32_t y; 34 float z; 35}); 36 37std::ostream& operator<<(std::ostream& os, const PackedStruct& v) { 38 return os << "p:" << v.x << "," << v.y << "," << v.z; 39} 40 41PYBIND11_PACKED(struct NestedStruct { 42 SimpleStruct a; 43 PackedStruct b; 44}); 45 46std::ostream& operator<<(std::ostream& os, const NestedStruct& v) { 47 return os << "n:a=" << v.a << ";b=" << v.b; 48} 49 50struct PartialStruct { 51 bool x; 52 uint32_t y; 53 float z; 54 uint64_t dummy2; 55}; 56 57struct PartialNestedStruct { 58 uint64_t dummy1; 59 PartialStruct a; 60 uint64_t dummy2; 61}; 62 63struct UnboundStruct { }; 64 65struct StringStruct { 66 char a[3]; 67 std::array<char, 3> b; 68}; 69 70PYBIND11_PACKED(struct StructWithUglyNames { 71 int8_t __x__; 72 uint64_t __y__; 73}); 74 75enum class E1 : int64_t { A = -1, B = 1 }; 76enum E2 : uint8_t { X = 1, Y = 2 }; 77 78PYBIND11_PACKED(struct EnumStruct { 79 E1 e1; 80 E2 e2; 81}); 82 83std::ostream& operator<<(std::ostream& os, const StringStruct& v) { 84 os << "a='"; 85 for (size_t i = 0; i < 3 && v.a[i]; i++) os << v.a[i]; 86 os << "',b='"; 87 for (size_t i = 0; i < 3 && v.b[i]; i++) os << v.b[i]; 88 return os << "'"; 89} 90 91std::ostream& operator<<(std::ostream& os, const EnumStruct& v) { 92 return os << "e1=" << (v.e1 == E1::A ? "A" : "B") << ",e2=" << (v.e2 == E2::X ? "X" : "Y"); 93} 94 95template <typename T> 96py::array mkarray_via_buffer(size_t n) { 97 return py::array(py::buffer_info(nullptr, sizeof(T), 98 py::format_descriptor<T>::format(), 99 1, { n }, { sizeof(T) })); 100} 101 102template <typename S> 103py::array_t<S, 0> create_recarray(size_t n) { 104 auto arr = mkarray_via_buffer<S>(n); 105 auto req = arr.request(); 106 auto ptr = static_cast<S*>(req.ptr); 107 for (size_t i = 0; i < n; i++) { 108 ptr[i].x = i % 2 != 0; ptr[i].y = (uint32_t) i; ptr[i].z = (float) i * 1.5f; 109 } 110 return arr; 111} 112 113std::string get_format_unbound() { 114 return py::format_descriptor<UnboundStruct>::format(); 115} 116 117py::array_t<NestedStruct, 0> create_nested(size_t n) { 118 auto arr = mkarray_via_buffer<NestedStruct>(n); 119 auto req = arr.request(); 120 auto ptr = static_cast<NestedStruct*>(req.ptr); 121 for (size_t i = 0; i < n; i++) { 122 ptr[i].a.x = i % 2 != 0; ptr[i].a.y = (uint32_t) i; ptr[i].a.z = (float) i * 1.5f; 123 ptr[i].b.x = (i + 1) % 2 != 0; ptr[i].b.y = (uint32_t) (i + 1); ptr[i].b.z = (float) (i + 1) * 1.5f; 124 } 125 return arr; 126} 127 128py::array_t<PartialNestedStruct, 0> create_partial_nested(size_t n) { 129 auto arr = mkarray_via_buffer<PartialNestedStruct>(n); 130 auto req = arr.request(); 131 auto ptr = static_cast<PartialNestedStruct*>(req.ptr); 132 for (size_t i = 0; i < n; i++) { 133 ptr[i].a.x = i % 2 != 0; ptr[i].a.y = (uint32_t) i; ptr[i].a.z = (float) i * 1.5f; 134 } 135 return arr; 136} 137 138py::array_t<StringStruct, 0> create_string_array(bool non_empty) { 139 auto arr = mkarray_via_buffer<StringStruct>(non_empty ? 4 : 0); 140 if (non_empty) { 141 auto req = arr.request(); 142 auto ptr = static_cast<StringStruct*>(req.ptr); 143 for (size_t i = 0; i < req.size * req.itemsize; i++) 144 static_cast<char*>(req.ptr)[i] = 0; 145 ptr[1].a[0] = 'a'; ptr[1].b[0] = 'a'; 146 ptr[2].a[0] = 'a'; ptr[2].b[0] = 'a'; 147 ptr[3].a[0] = 'a'; ptr[3].b[0] = 'a'; 148 149 ptr[2].a[1] = 'b'; ptr[2].b[1] = 'b'; 150 ptr[3].a[1] = 'b'; ptr[3].b[1] = 'b'; 151 152 ptr[3].a[2] = 'c'; ptr[3].b[2] = 'c'; 153 } 154 return arr; 155} 156 157py::array_t<EnumStruct, 0> create_enum_array(size_t n) { 158 auto arr = mkarray_via_buffer<EnumStruct>(n); 159 auto ptr = (EnumStruct *) arr.mutable_data(); 160 for (size_t i = 0; i < n; i++) { 161 ptr[i].e1 = static_cast<E1>(-1 + ((int) i % 2) * 2); 162 ptr[i].e2 = static_cast<E2>(1 + (i % 2)); 163 } 164 return arr; 165} 166 167template <typename S> 168py::list print_recarray(py::array_t<S, 0> arr) { 169 const auto req = arr.request(); 170 const auto ptr = static_cast<S*>(req.ptr); 171 auto l = py::list(); 172 for (size_t i = 0; i < req.size; i++) { 173 std::stringstream ss; 174 ss << ptr[i]; 175 l.append(py::str(ss.str())); 176 } 177 return l; 178} 179 180py::list print_format_descriptors() { 181 const auto fmts = { 182 py::format_descriptor<SimpleStruct>::format(), 183 py::format_descriptor<PackedStruct>::format(), 184 py::format_descriptor<NestedStruct>::format(), 185 py::format_descriptor<PartialStruct>::format(), 186 py::format_descriptor<PartialNestedStruct>::format(), 187 py::format_descriptor<StringStruct>::format(), 188 py::format_descriptor<EnumStruct>::format() 189 }; 190 auto l = py::list(); 191 for (const auto &fmt : fmts) { 192 l.append(py::cast(fmt)); 193 } 194 return l; 195} 196 197py::list print_dtypes() { 198 const auto dtypes = { 199 py::str(py::dtype::of<SimpleStruct>()), 200 py::str(py::dtype::of<PackedStruct>()), 201 py::str(py::dtype::of<NestedStruct>()), 202 py::str(py::dtype::of<PartialStruct>()), 203 py::str(py::dtype::of<PartialNestedStruct>()), 204 py::str(py::dtype::of<StringStruct>()), 205 py::str(py::dtype::of<EnumStruct>()), 206 py::str(py::dtype::of<StructWithUglyNames>()) 207 }; 208 auto l = py::list(); 209 for (const auto &s : dtypes) { 210 l.append(s); 211 } 212 return l; 213} 214 215py::array_t<int32_t, 0> test_array_ctors(int i) { 216 using arr_t = py::array_t<int32_t, 0>; 217 218 std::vector<int32_t> data { 1, 2, 3, 4, 5, 6 }; 219 std::vector<size_t> shape { 3, 2 }; 220 std::vector<size_t> strides { 8, 4 }; 221 222 auto ptr = data.data(); 223 auto vptr = (void *) ptr; 224 auto dtype = py::dtype("int32"); 225 226 py::buffer_info buf_ndim1(vptr, 4, "i", 6); 227 py::buffer_info buf_ndim1_null(nullptr, 4, "i", 6); 228 py::buffer_info buf_ndim2(vptr, 4, "i", 2, shape, strides); 229 py::buffer_info buf_ndim2_null(nullptr, 4, "i", 2, shape, strides); 230 231 auto fill = [](py::array arr) { 232 auto req = arr.request(); 233 for (int i = 0; i < 6; i++) ((int32_t *) req.ptr)[i] = i + 1; 234 return arr; 235 }; 236 237 switch (i) { 238 // shape: (3, 2) 239 case 10: return arr_t(shape, strides, ptr); 240 case 11: return py::array(shape, strides, ptr); 241 case 12: return py::array(dtype, shape, strides, vptr); 242 case 13: return arr_t(shape, ptr); 243 case 14: return py::array(shape, ptr); 244 case 15: return py::array(dtype, shape, vptr); 245 case 16: return arr_t(buf_ndim2); 246 case 17: return py::array(buf_ndim2); 247 // shape: (3, 2) - post-fill 248 case 20: return fill(arr_t(shape, strides)); 249 case 21: return py::array(shape, strides, ptr); // can't have nullptr due to templated ctor 250 case 22: return fill(py::array(dtype, shape, strides)); 251 case 23: return fill(arr_t(shape)); 252 case 24: return py::array(shape, ptr); // can't have nullptr due to templated ctor 253 case 25: return fill(py::array(dtype, shape)); 254 case 26: return fill(arr_t(buf_ndim2_null)); 255 case 27: return fill(py::array(buf_ndim2_null)); 256 // shape: (6, ) 257 case 30: return arr_t(6, ptr); 258 case 31: return py::array(6, ptr); 259 case 32: return py::array(dtype, 6, vptr); 260 case 33: return arr_t(buf_ndim1); 261 case 34: return py::array(buf_ndim1); 262 // shape: (6, ) 263 case 40: return fill(arr_t(6)); 264 case 41: return py::array(6, ptr); // can't have nullptr due to templated ctor 265 case 42: return fill(py::array(dtype, 6)); 266 case 43: return fill(arr_t(buf_ndim1_null)); 267 case 44: return fill(py::array(buf_ndim1_null)); 268 } 269 return arr_t(); 270} 271 272py::list test_dtype_ctors() { 273 py::list list; 274 list.append(py::dtype("int32")); 275 list.append(py::dtype(std::string("float64"))); 276 list.append(py::dtype::from_args(py::str("bool"))); 277 py::list names, offsets, formats; 278 py::dict dict; 279 names.append(py::str("a")); names.append(py::str("b")); dict["names"] = names; 280 offsets.append(py::int_(1)); offsets.append(py::int_(10)); dict["offsets"] = offsets; 281 formats.append(py::dtype("int32")); formats.append(py::dtype("float64")); dict["formats"] = formats; 282 dict["itemsize"] = py::int_(20); 283 list.append(py::dtype::from_args(dict)); 284 list.append(py::dtype(names, formats, offsets, 20)); 285 list.append(py::dtype(py::buffer_info((void *) 0, sizeof(unsigned int), "I", 1))); 286 list.append(py::dtype(py::buffer_info((void *) 0, 0, "T{i:a:f:b:}", 1))); 287 return list; 288} 289 290struct TrailingPaddingStruct { 291 int32_t a; 292 char b; 293}; 294 295py::dtype trailing_padding_dtype() { 296 return py::dtype::of<TrailingPaddingStruct>(); 297} 298 299py::dtype buffer_to_dtype(py::buffer& buf) { 300 return py::dtype(buf.request()); 301} 302 303py::list test_dtype_methods() { 304 py::list list; 305 auto dt1 = py::dtype::of<int32_t>(); 306 auto dt2 = py::dtype::of<SimpleStruct>(); 307 list.append(dt1); list.append(dt2); 308 list.append(py::bool_(dt1.has_fields())); list.append(py::bool_(dt2.has_fields())); 309 list.append(py::int_(dt1.itemsize())); list.append(py::int_(dt2.itemsize())); 310 return list; 311} 312 313test_initializer numpy_dtypes([](py::module &m) { 314 try { 315 py::module::import("numpy"); 316 } catch (...) { 317 return; 318 } 319 320 // typeinfo may be registered before the dtype descriptor for scalar casts to work... 321 py::class_<SimpleStruct>(m, "SimpleStruct"); 322 323 PYBIND11_NUMPY_DTYPE(SimpleStruct, x, y, z); 324 PYBIND11_NUMPY_DTYPE(PackedStruct, x, y, z); 325 PYBIND11_NUMPY_DTYPE(NestedStruct, a, b); 326 PYBIND11_NUMPY_DTYPE(PartialStruct, x, y, z); 327 PYBIND11_NUMPY_DTYPE(PartialNestedStruct, a); 328 PYBIND11_NUMPY_DTYPE(StringStruct, a, b); 329 PYBIND11_NUMPY_DTYPE(EnumStruct, e1, e2); 330 PYBIND11_NUMPY_DTYPE(TrailingPaddingStruct, a, b); 331 332 // ... or after 333 py::class_<PackedStruct>(m, "PackedStruct"); 334 335 PYBIND11_NUMPY_DTYPE_EX(StructWithUglyNames, __x__, "x", __y__, "y"); 336 337 m.def("create_rec_simple", &create_recarray<SimpleStruct>); 338 m.def("create_rec_packed", &create_recarray<PackedStruct>); 339 m.def("create_rec_nested", &create_nested); 340 m.def("create_rec_partial", &create_recarray<PartialStruct>); 341 m.def("create_rec_partial_nested", &create_partial_nested); 342 m.def("print_format_descriptors", &print_format_descriptors); 343 m.def("print_rec_simple", &print_recarray<SimpleStruct>); 344 m.def("print_rec_packed", &print_recarray<PackedStruct>); 345 m.def("print_rec_nested", &print_recarray<NestedStruct>); 346 m.def("print_dtypes", &print_dtypes); 347 m.def("get_format_unbound", &get_format_unbound); 348 m.def("create_string_array", &create_string_array); 349 m.def("print_string_array", &print_recarray<StringStruct>); 350 m.def("create_enum_array", &create_enum_array); 351 m.def("print_enum_array", &print_recarray<EnumStruct>); 352 m.def("test_array_ctors", &test_array_ctors); 353 m.def("test_dtype_ctors", &test_dtype_ctors); 354 m.def("test_dtype_methods", &test_dtype_methods); 355 m.def("trailing_padding_dtype", &trailing_padding_dtype); 356 m.def("buffer_to_dtype", &buffer_to_dtype); 357 m.def("f_simple", [](SimpleStruct s) { return s.y * 10; }); 358 m.def("f_packed", [](PackedStruct s) { return s.y * 10; }); 359 m.def("f_nested", [](NestedStruct s) { return s.a.y * 10; }); 360 m.def("register_dtype", []() { PYBIND11_NUMPY_DTYPE(SimpleStruct, x, y, z); }); 361}); 362 363#undef PYBIND11_PACKED 364