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