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