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