1 /*
2 tests/test_sequences_and_iterators.cpp -- supporting Pythons' sequence protocol, iterators,
3 etc.
4
5 Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
6
7 All rights reserved. Use of this source code is governed by a
8 BSD-style license that can be found in the LICENSE file.
9 */
10
11 #include "pybind11_tests.h"
12 #include "constructor_stats.h"
13 #include <pybind11/operators.h>
14 #include <pybind11/stl.h>
15
16 #include <algorithm>
17
18 template<typename T>
19 class NonZeroIterator {
20 const T* ptr_;
21 public:
NonZeroIterator(const T * ptr)22 NonZeroIterator(const T* ptr) : ptr_(ptr) {}
operator *() const23 const T& operator*() const { return *ptr_; }
operator ++()24 NonZeroIterator& operator++() { ++ptr_; return *this; }
25 };
26
27 class NonZeroSentinel {};
28
29 template<typename A, typename B>
operator ==(const NonZeroIterator<std::pair<A,B>> & it,const NonZeroSentinel &)30 bool operator==(const NonZeroIterator<std::pair<A, B>>& it, const NonZeroSentinel&) {
31 return !(*it).first || !(*it).second;
32 }
33
34 template <typename PythonType>
test_random_access_iterator(PythonType x)35 py::list test_random_access_iterator(PythonType x) {
36 if (x.size() < 5)
37 throw py::value_error("Please provide at least 5 elements for testing.");
38
39 auto checks = py::list();
40 auto assert_equal = [&checks](py::handle a, py::handle b) {
41 auto result = PyObject_RichCompareBool(a.ptr(), b.ptr(), Py_EQ);
42 if (result == -1) { throw py::error_already_set(); }
43 checks.append(result != 0);
44 };
45
46 auto it = x.begin();
47 assert_equal(x[0], *it);
48 assert_equal(x[0], it[0]);
49 assert_equal(x[1], it[1]);
50
51 assert_equal(x[1], *(++it));
52 assert_equal(x[1], *(it++));
53 assert_equal(x[2], *it);
54 assert_equal(x[3], *(it += 1));
55 assert_equal(x[2], *(--it));
56 assert_equal(x[2], *(it--));
57 assert_equal(x[1], *it);
58 assert_equal(x[0], *(it -= 1));
59
60 assert_equal(it->attr("real"), x[0].attr("real"));
61 assert_equal((it + 1)->attr("real"), x[1].attr("real"));
62
63 assert_equal(x[1], *(it + 1));
64 assert_equal(x[1], *(1 + it));
65 it += 3;
66 assert_equal(x[1], *(it - 2));
67
68 checks.append(static_cast<std::size_t>(x.end() - x.begin()) == x.size());
69 checks.append((x.begin() + static_cast<std::ptrdiff_t>(x.size())) == x.end());
70 checks.append(x.begin() < x.end());
71
72 return checks;
73 }
74
TEST_SUBMODULE(sequences_and_iterators,m)75 TEST_SUBMODULE(sequences_and_iterators, m) {
76 // test_sliceable
77 class Sliceable{
78 public:
79 Sliceable(int n): size(n) {}
80 int start,stop,step;
81 int size;
82 };
83 py::class_<Sliceable>(m,"Sliceable")
84 .def(py::init<int>())
85 .def("__getitem__",[](const Sliceable &s, py::slice slice) {
86 py::ssize_t start, stop, step, slicelength;
87 if (!slice.compute(s.size, &start, &stop, &step, &slicelength))
88 throw py::error_already_set();
89 int istart = static_cast<int>(start);
90 int istop = static_cast<int>(stop);
91 int istep = static_cast<int>(step);
92 return std::make_tuple(istart,istop,istep);
93 })
94 ;
95
96 // test_sequence
97 class Sequence {
98 public:
99 Sequence(size_t size) : m_size(size) {
100 print_created(this, "of size", m_size);
101 m_data = new float[size];
102 memset(m_data, 0, sizeof(float) * size);
103 }
104 Sequence(const std::vector<float> &value) : m_size(value.size()) {
105 print_created(this, "of size", m_size, "from std::vector");
106 m_data = new float[m_size];
107 memcpy(m_data, &value[0], sizeof(float) * m_size);
108 }
109 Sequence(const Sequence &s) : m_size(s.m_size) {
110 print_copy_created(this);
111 m_data = new float[m_size];
112 memcpy(m_data, s.m_data, sizeof(float)*m_size);
113 }
114 Sequence(Sequence &&s) : m_size(s.m_size), m_data(s.m_data) {
115 print_move_created(this);
116 s.m_size = 0;
117 s.m_data = nullptr;
118 }
119
120 ~Sequence() { print_destroyed(this); delete[] m_data; }
121
122 Sequence &operator=(const Sequence &s) {
123 if (&s != this) {
124 delete[] m_data;
125 m_size = s.m_size;
126 m_data = new float[m_size];
127 memcpy(m_data, s.m_data, sizeof(float)*m_size);
128 }
129 print_copy_assigned(this);
130 return *this;
131 }
132
133 Sequence &operator=(Sequence &&s) {
134 if (&s != this) {
135 delete[] m_data;
136 m_size = s.m_size;
137 m_data = s.m_data;
138 s.m_size = 0;
139 s.m_data = nullptr;
140 }
141 print_move_assigned(this);
142 return *this;
143 }
144
145 bool operator==(const Sequence &s) const {
146 if (m_size != s.size()) return false;
147 for (size_t i = 0; i < m_size; ++i)
148 if (m_data[i] != s[i])
149 return false;
150 return true;
151 }
152 bool operator!=(const Sequence &s) const { return !operator==(s); }
153
154 float operator[](size_t index) const { return m_data[index]; }
155 float &operator[](size_t index) { return m_data[index]; }
156
157 bool contains(float v) const {
158 for (size_t i = 0; i < m_size; ++i)
159 if (v == m_data[i])
160 return true;
161 return false;
162 }
163
164 Sequence reversed() const {
165 Sequence result(m_size);
166 for (size_t i = 0; i < m_size; ++i)
167 result[m_size - i - 1] = m_data[i];
168 return result;
169 }
170
171 size_t size() const { return m_size; }
172
173 const float *begin() const { return m_data; }
174 const float *end() const { return m_data+m_size; }
175
176 private:
177 size_t m_size;
178 float *m_data;
179 };
180 py::class_<Sequence>(m, "Sequence")
181 .def(py::init<size_t>())
182 .def(py::init<const std::vector<float>&>())
183 /// Bare bones interface
184 .def("__getitem__", [](const Sequence &s, size_t i) {
185 if (i >= s.size()) throw py::index_error();
186 return s[i];
187 })
188 .def("__setitem__", [](Sequence &s, size_t i, float v) {
189 if (i >= s.size()) throw py::index_error();
190 s[i] = v;
191 })
192 .def("__len__", &Sequence::size)
193 /// Optional sequence protocol operations
194 .def("__iter__", [](const Sequence &s) { return py::make_iterator(s.begin(), s.end()); },
195 py::keep_alive<0, 1>() /* Essential: keep object alive while iterator exists */)
196 .def("__contains__", [](const Sequence &s, float v) { return s.contains(v); })
197 .def("__reversed__", [](const Sequence &s) -> Sequence { return s.reversed(); })
198 /// Slicing protocol (optional)
199 .def("__getitem__", [](const Sequence &s, py::slice slice) -> Sequence* {
200 size_t start, stop, step, slicelength;
201 if (!slice.compute(s.size(), &start, &stop, &step, &slicelength))
202 throw py::error_already_set();
203 auto *seq = new Sequence(slicelength);
204 for (size_t i = 0; i < slicelength; ++i) {
205 (*seq)[i] = s[start]; start += step;
206 }
207 return seq;
208 })
209 .def("__setitem__", [](Sequence &s, py::slice slice, const Sequence &value) {
210 size_t start, stop, step, slicelength;
211 if (!slice.compute(s.size(), &start, &stop, &step, &slicelength))
212 throw py::error_already_set();
213 if (slicelength != value.size())
214 throw std::runtime_error("Left and right hand size of slice assignment have different sizes!");
215 for (size_t i = 0; i < slicelength; ++i) {
216 s[start] = value[i]; start += step;
217 }
218 })
219 /// Comparisons
220 .def(py::self == py::self)
221 .def(py::self != py::self)
222 // Could also define py::self + py::self for concatenation, etc.
223 ;
224
225 // test_map_iterator
226 // Interface of a map-like object that isn't (directly) an unordered_map, but provides some basic
227 // map-like functionality.
228 class StringMap {
229 public:
230 StringMap() = default;
231 StringMap(std::unordered_map<std::string, std::string> init)
232 : map(std::move(init)) {}
233
234 void set(std::string key, std::string val) { map[key] = val; }
235 std::string get(std::string key) const { return map.at(key); }
236 size_t size() const { return map.size(); }
237 private:
238 std::unordered_map<std::string, std::string> map;
239 public:
240 decltype(map.cbegin()) begin() const { return map.cbegin(); }
241 decltype(map.cend()) end() const { return map.cend(); }
242 };
243 py::class_<StringMap>(m, "StringMap")
244 .def(py::init<>())
245 .def(py::init<std::unordered_map<std::string, std::string>>())
246 .def("__getitem__", [](const StringMap &map, std::string key) {
247 try { return map.get(key); }
248 catch (const std::out_of_range&) {
249 throw py::key_error("key '" + key + "' does not exist");
250 }
251 })
252 .def("__setitem__", &StringMap::set)
253 .def("__len__", &StringMap::size)
254 .def("__iter__", [](const StringMap &map) { return py::make_key_iterator(map.begin(), map.end()); },
255 py::keep_alive<0, 1>())
256 .def("items", [](const StringMap &map) { return py::make_iterator(map.begin(), map.end()); },
257 py::keep_alive<0, 1>())
258 ;
259
260 // test_generalized_iterators
261 class IntPairs {
262 public:
263 IntPairs(std::vector<std::pair<int, int>> data) : data_(std::move(data)) {}
264 const std::pair<int, int>* begin() const { return data_.data(); }
265 private:
266 std::vector<std::pair<int, int>> data_;
267 };
268 py::class_<IntPairs>(m, "IntPairs")
269 .def(py::init<std::vector<std::pair<int, int>>>())
270 .def("nonzero", [](const IntPairs& s) {
271 return py::make_iterator(NonZeroIterator<std::pair<int, int>>(s.begin()), NonZeroSentinel());
272 }, py::keep_alive<0, 1>())
273 .def("nonzero_keys", [](const IntPairs& s) {
274 return py::make_key_iterator(NonZeroIterator<std::pair<int, int>>(s.begin()), NonZeroSentinel());
275 }, py::keep_alive<0, 1>())
276 ;
277
278
279 #if 0
280 // Obsolete: special data structure for exposing custom iterator types to python
281 // kept here for illustrative purposes because there might be some use cases which
282 // are not covered by the much simpler py::make_iterator
283
284 struct PySequenceIterator {
285 PySequenceIterator(const Sequence &seq, py::object ref) : seq(seq), ref(ref) { }
286
287 float next() {
288 if (index == seq.size())
289 throw py::stop_iteration();
290 return seq[index++];
291 }
292
293 const Sequence &seq;
294 py::object ref; // keep a reference
295 size_t index = 0;
296 };
297
298 py::class_<PySequenceIterator>(seq, "Iterator")
299 .def("__iter__", [](PySequenceIterator &it) -> PySequenceIterator& { return it; })
300 .def("__next__", &PySequenceIterator::next);
301
302 On the actual Sequence object, the iterator would be constructed as follows:
303 .def("__iter__", [](py::object s) { return PySequenceIterator(s.cast<const Sequence &>(), s); })
304 #endif
305
306 // test_python_iterator_in_cpp
307 m.def("object_to_list", [](py::object o) {
308 auto l = py::list();
309 for (auto item : o) {
310 l.append(item);
311 }
312 return l;
313 });
314
315 m.def("iterator_to_list", [](py::iterator it) {
316 auto l = py::list();
317 while (it != py::iterator::sentinel()) {
318 l.append(*it);
319 ++it;
320 }
321 return l;
322 });
323
324 // test_sequence_length: check that Python sequences can be converted to py::sequence.
325 m.def("sequence_length", [](py::sequence seq) { return seq.size(); });
326
327 // Make sure that py::iterator works with std algorithms
328 m.def("count_none", [](py::object o) {
329 return std::count_if(o.begin(), o.end(), [](py::handle h) { return h.is_none(); });
330 });
331
332 m.def("find_none", [](py::object o) {
333 auto it = std::find_if(o.begin(), o.end(), [](py::handle h) { return h.is_none(); });
334 return it->is_none();
335 });
336
337 m.def("count_nonzeros", [](py::dict d) {
338 return std::count_if(d.begin(), d.end(), [](std::pair<py::handle, py::handle> p) {
339 return p.second.cast<int>() != 0;
340 });
341 });
342
343 m.def("tuple_iterator", &test_random_access_iterator<py::tuple>);
344 m.def("list_iterator", &test_random_access_iterator<py::list>);
345 m.def("sequence_iterator", &test_random_access_iterator<py::sequence>);
346
347 // test_iterator_passthrough
348 // #181: iterator passthrough did not compile
349 m.def("iterator_passthrough", [](py::iterator s) -> py::iterator {
350 return py::make_iterator(std::begin(s), std::end(s));
351 });
352
353 // test_iterator_rvp
354 // #388: Can't make iterators via make_iterator() with different r/v policies
355 static std::vector<int> list = { 1, 2, 3 };
356 m.def("make_iterator_1", []() { return py::make_iterator<py::return_value_policy::copy>(list); });
357 m.def("make_iterator_2", []() { return py::make_iterator<py::return_value_policy::automatic>(list); });
358 }
359