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