Ivan Smirnov | 91b3d68 | 2016-08-29 02:41:05 +0100 | [diff] [blame] | 1 | /* |
| 2 | tests/test_numpy_array.cpp -- test core array functionality |
| 3 | |
| 4 | Copyright (c) 2016 Ivan Smirnov <i.s.smirnov@gmail.com> |
| 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" |
Ivan Smirnov | aca6bca | 2016-09-08 23:03:35 +0100 | [diff] [blame] | 11 | |
Ivan Smirnov | 91b3d68 | 2016-08-29 02:41:05 +0100 | [diff] [blame] | 12 | #include <pybind11/numpy.h> |
| 13 | #include <pybind11/stl.h> |
| 14 | |
Ivan Smirnov | aca6bca | 2016-09-08 23:03:35 +0100 | [diff] [blame] | 15 | #include <cstdint> |
| 16 | #include <vector> |
| 17 | |
| 18 | using arr = py::array; |
| 19 | using arr_t = py::array_t<uint16_t, 0>; |
| 20 | |
Sylvain Corlay | 5027c4f | 2016-11-16 08:53:37 -0800 | [diff] [blame] | 21 | template<typename... Ix> arr data(const arr& a, Ix... index) { |
Ivan Smirnov | aca6bca | 2016-09-08 23:03:35 +0100 | [diff] [blame] | 22 | return arr(a.nbytes() - a.offset_at(index...), (const uint8_t *) a.data(index...)); |
| 23 | } |
| 24 | |
Sylvain Corlay | 5027c4f | 2016-11-16 08:53:37 -0800 | [diff] [blame] | 25 | template<typename... Ix> arr data_t(const arr_t& a, Ix... index) { |
Ivan Smirnov | aca6bca | 2016-09-08 23:03:35 +0100 | [diff] [blame] | 26 | return arr(a.size() - a.index_at(index...), a.data(index...)); |
| 27 | } |
| 28 | |
| 29 | arr& mutate_data(arr& a) { |
| 30 | auto ptr = (uint8_t *) a.mutable_data(); |
| 31 | for (size_t i = 0; i < a.nbytes(); i++) |
| 32 | ptr[i] = (uint8_t) (ptr[i] * 2); |
| 33 | return a; |
| 34 | } |
| 35 | |
| 36 | arr_t& mutate_data_t(arr_t& a) { |
| 37 | auto ptr = a.mutable_data(); |
| 38 | for (size_t i = 0; i < a.size(); i++) |
| 39 | ptr[i]++; |
| 40 | return a; |
| 41 | } |
| 42 | |
Sylvain Corlay | 5027c4f | 2016-11-16 08:53:37 -0800 | [diff] [blame] | 43 | template<typename... Ix> arr& mutate_data(arr& a, Ix... index) { |
Ivan Smirnov | aca6bca | 2016-09-08 23:03:35 +0100 | [diff] [blame] | 44 | auto ptr = (uint8_t *) a.mutable_data(index...); |
| 45 | for (size_t i = 0; i < a.nbytes() - a.offset_at(index...); i++) |
| 46 | ptr[i] = (uint8_t) (ptr[i] * 2); |
| 47 | return a; |
| 48 | } |
| 49 | |
Sylvain Corlay | 5027c4f | 2016-11-16 08:53:37 -0800 | [diff] [blame] | 50 | template<typename... Ix> arr_t& mutate_data_t(arr_t& a, Ix... index) { |
Ivan Smirnov | aca6bca | 2016-09-08 23:03:35 +0100 | [diff] [blame] | 51 | auto ptr = a.mutable_data(index...); |
| 52 | for (size_t i = 0; i < a.size() - a.index_at(index...); i++) |
| 53 | ptr[i]++; |
| 54 | return a; |
| 55 | } |
| 56 | |
Sylvain Corlay | 5027c4f | 2016-11-16 08:53:37 -0800 | [diff] [blame] | 57 | template<typename... Ix> size_t index_at(const arr& a, Ix... idx) { return a.index_at(idx...); } |
| 58 | template<typename... Ix> size_t index_at_t(const arr_t& a, Ix... idx) { return a.index_at(idx...); } |
| 59 | template<typename... Ix> size_t offset_at(const arr& a, Ix... idx) { return a.offset_at(idx...); } |
| 60 | template<typename... Ix> size_t offset_at_t(const arr_t& a, Ix... idx) { return a.offset_at(idx...); } |
| 61 | template<typename... Ix> size_t at_t(const arr_t& a, Ix... idx) { return a.at(idx...); } |
| 62 | template<typename... Ix> arr_t& mutate_at_t(arr_t& a, Ix... idx) { a.mutable_at(idx...)++; return a; } |
Ivan Smirnov | aca6bca | 2016-09-08 23:03:35 +0100 | [diff] [blame] | 63 | |
| 64 | #define def_index_fn(name, type) \ |
| 65 | sm.def(#name, [](type a) { return name(a); }); \ |
| 66 | sm.def(#name, [](type a, int i) { return name(a, i); }); \ |
| 67 | sm.def(#name, [](type a, int i, int j) { return name(a, i, j); }); \ |
| 68 | sm.def(#name, [](type a, int i, int j, int k) { return name(a, i, j, k); }); |
| 69 | |
Ivan Smirnov | 91b3d68 | 2016-08-29 02:41:05 +0100 | [diff] [blame] | 70 | test_initializer numpy_array([](py::module &m) { |
Ivan Smirnov | aca6bca | 2016-09-08 23:03:35 +0100 | [diff] [blame] | 71 | auto sm = m.def_submodule("array"); |
| 72 | |
| 73 | sm.def("ndim", [](const arr& a) { return a.ndim(); }); |
| 74 | sm.def("shape", [](const arr& a) { return arr(a.ndim(), a.shape()); }); |
| 75 | sm.def("shape", [](const arr& a, size_t dim) { return a.shape(dim); }); |
| 76 | sm.def("strides", [](const arr& a) { return arr(a.ndim(), a.strides()); }); |
| 77 | sm.def("strides", [](const arr& a, size_t dim) { return a.strides(dim); }); |
| 78 | sm.def("writeable", [](const arr& a) { return a.writeable(); }); |
| 79 | sm.def("size", [](const arr& a) { return a.size(); }); |
| 80 | sm.def("itemsize", [](const arr& a) { return a.itemsize(); }); |
| 81 | sm.def("nbytes", [](const arr& a) { return a.nbytes(); }); |
| 82 | sm.def("owndata", [](const arr& a) { return a.owndata(); }); |
| 83 | |
| 84 | def_index_fn(data, const arr&); |
| 85 | def_index_fn(data_t, const arr_t&); |
| 86 | def_index_fn(index_at, const arr&); |
| 87 | def_index_fn(index_at_t, const arr_t&); |
| 88 | def_index_fn(offset_at, const arr&); |
| 89 | def_index_fn(offset_at_t, const arr_t&); |
| 90 | def_index_fn(mutate_data, arr&); |
| 91 | def_index_fn(mutate_data_t, arr_t&); |
| 92 | def_index_fn(at_t, const arr_t&); |
| 93 | def_index_fn(mutate_at_t, arr_t&); |
Wenzel Jakob | 43f6aa6 | 2016-10-12 23:34:06 +0200 | [diff] [blame] | 94 | |
| 95 | sm.def("make_f_array", [] { |
| 96 | return py::array_t<float>({ 2, 2 }, { 4, 8 }); |
| 97 | }); |
| 98 | |
| 99 | sm.def("make_c_array", [] { |
| 100 | return py::array_t<float>({ 2, 2 }, { 8, 4 }); |
| 101 | }); |
Wenzel Jakob | 369e9b3 | 2016-10-13 00:57:42 +0200 | [diff] [blame] | 102 | |
| 103 | sm.def("wrap", [](py::array a) { |
| 104 | return py::array( |
| 105 | a.dtype(), |
| 106 | std::vector<size_t>(a.shape(), a.shape() + a.ndim()), |
| 107 | std::vector<size_t>(a.strides(), a.strides() + a.ndim()), |
| 108 | a.data(), |
| 109 | a |
| 110 | ); |
| 111 | }); |
Wenzel Jakob | fac7c09 | 2016-10-13 10:37:52 +0200 | [diff] [blame] | 112 | |
| 113 | struct ArrayClass { |
| 114 | int data[2] = { 1, 2 }; |
| 115 | ArrayClass() { py::print("ArrayClass()"); } |
| 116 | ~ArrayClass() { py::print("~ArrayClass()"); } |
| 117 | }; |
| 118 | |
| 119 | py::class_<ArrayClass>(sm, "ArrayClass") |
| 120 | .def(py::init<>()) |
| 121 | .def("numpy_view", [](py::object &obj) { |
| 122 | py::print("ArrayClass::numpy_view()"); |
| 123 | ArrayClass &a = obj.cast<ArrayClass&>(); |
| 124 | return py::array_t<int>({2}, {4}, a.data, obj); |
| 125 | } |
| 126 | ); |
Wenzel Jakob | 496feac | 2016-10-28 00:37:07 +0200 | [diff] [blame] | 127 | |
Wenzel Jakob | 030d10e | 2016-10-28 01:23:42 +0200 | [diff] [blame] | 128 | sm.def("function_taking_uint64", [](uint64_t) { }); |
Dean Moldovan | 4de2710 | 2016-11-16 01:35:22 +0100 | [diff] [blame] | 129 | |
| 130 | sm.def("isinstance_untyped", [](py::object yes, py::object no) { |
| 131 | return py::isinstance<py::array>(yes) && !py::isinstance<py::array>(no); |
| 132 | }); |
| 133 | |
| 134 | sm.def("isinstance_typed", [](py::object o) { |
| 135 | return py::isinstance<py::array_t<double>>(o) && !py::isinstance<py::array_t<int>>(o); |
| 136 | }); |
| 137 | |
| 138 | sm.def("default_constructors", []() { |
| 139 | return py::dict( |
| 140 | "array"_a=py::array(), |
| 141 | "array_t<int32>"_a=py::array_t<std::int32_t>(), |
| 142 | "array_t<double>"_a=py::array_t<double>() |
| 143 | ); |
| 144 | }); |
| 145 | |
| 146 | sm.def("converting_constructors", [](py::object o) { |
| 147 | return py::dict( |
| 148 | "array"_a=py::array(o), |
| 149 | "array_t<int32>"_a=py::array_t<std::int32_t>(o), |
| 150 | "array_t<double>"_a=py::array_t<double>(o) |
| 151 | ); |
| 152 | }); |
Jason Rhinelander | ee2e5a5 | 2017-02-24 05:33:31 -0500 | [diff] [blame] | 153 | |
Jason Rhinelander | c44fe6f | 2017-02-26 18:03:00 -0500 | [diff] [blame^] | 154 | // Overload resolution tests: |
| 155 | sm.def("overloaded", [](py::array_t<double>) { return "double"; }); |
| 156 | sm.def("overloaded", [](py::array_t<float>) { return "float"; }); |
| 157 | sm.def("overloaded", [](py::array_t<int>) { return "int"; }); |
| 158 | sm.def("overloaded", [](py::array_t<unsigned short>) { return "unsigned short"; }); |
| 159 | sm.def("overloaded", [](py::array_t<long long>) { return "long long"; }); |
| 160 | sm.def("overloaded", [](py::array_t<std::complex<double>>) { return "double complex"; }); |
| 161 | sm.def("overloaded", [](py::array_t<std::complex<float>>) { return "float complex"; }); |
| 162 | |
| 163 | sm.def("overloaded2", [](py::array_t<std::complex<double>>) { return "double complex"; }); |
| 164 | sm.def("overloaded2", [](py::array_t<double>) { return "double"; }); |
| 165 | sm.def("overloaded2", [](py::array_t<std::complex<float>>) { return "float complex"; }); |
| 166 | sm.def("overloaded2", [](py::array_t<float>) { return "float"; }); |
| 167 | |
| 168 | // Only accept the exact types: |
| 169 | sm.def("overloaded3", [](py::array_t<int>) { return "int"; }, py::arg().noconvert()); |
| 170 | sm.def("overloaded3", [](py::array_t<double>) { return "double"; }, py::arg().noconvert()); |
| 171 | |
| 172 | // Make sure we don't do unsafe coercion (e.g. float to int) when not using forcecast, but |
| 173 | // rather that float gets converted via the safe (conversion to double) overload: |
| 174 | sm.def("overloaded4", [](py::array_t<long long, 0>) { return "long long"; }); |
| 175 | sm.def("overloaded4", [](py::array_t<double, 0>) { return "double"; }); |
| 176 | |
| 177 | // But we do allow conversion to int if forcecast is enabled (but only if no overload matches |
| 178 | // without conversion) |
| 179 | sm.def("overloaded5", [](py::array_t<unsigned int>) { return "unsigned int"; }); |
| 180 | sm.def("overloaded5", [](py::array_t<double>) { return "double"; }); |
| 181 | |
Jason Rhinelander | ee2e5a5 | 2017-02-24 05:33:31 -0500 | [diff] [blame] | 182 | // Issue 685: ndarray shouldn't go to std::string overload |
| 183 | sm.def("issue685", [](std::string) { return "string"; }); |
| 184 | sm.def("issue685", [](py::array) { return "array"; }); |
| 185 | sm.def("issue685", [](py::object) { return "other"; }); |
Ivan Smirnov | 91b3d68 | 2016-08-29 02:41:05 +0100 | [diff] [blame] | 186 | }); |