test_sequences_and_iterators.cpp revision 11986
111986Sandreas.sandberg@arm.com/* 211986Sandreas.sandberg@arm.com tests/test_sequences_and_iterators.cpp -- supporting Pythons' sequence protocol, iterators, 311986Sandreas.sandberg@arm.com etc. 411986Sandreas.sandberg@arm.com 511986Sandreas.sandberg@arm.com Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch> 611986Sandreas.sandberg@arm.com 711986Sandreas.sandberg@arm.com All rights reserved. Use of this source code is governed by a 811986Sandreas.sandberg@arm.com BSD-style license that can be found in the LICENSE file. 911986Sandreas.sandberg@arm.com*/ 1011986Sandreas.sandberg@arm.com 1111986Sandreas.sandberg@arm.com#include "pybind11_tests.h" 1211986Sandreas.sandberg@arm.com#include "constructor_stats.h" 1311986Sandreas.sandberg@arm.com#include <pybind11/operators.h> 1411986Sandreas.sandberg@arm.com#include <pybind11/stl.h> 1511986Sandreas.sandberg@arm.com 1611986Sandreas.sandberg@arm.comclass Sequence { 1711986Sandreas.sandberg@arm.compublic: 1811986Sandreas.sandberg@arm.com Sequence(size_t size) : m_size(size) { 1911986Sandreas.sandberg@arm.com print_created(this, "of size", m_size); 2011986Sandreas.sandberg@arm.com m_data = new float[size]; 2111986Sandreas.sandberg@arm.com memset(m_data, 0, sizeof(float) * size); 2211986Sandreas.sandberg@arm.com } 2311986Sandreas.sandberg@arm.com 2411986Sandreas.sandberg@arm.com Sequence(const std::vector<float> &value) : m_size(value.size()) { 2511986Sandreas.sandberg@arm.com print_created(this, "of size", m_size, "from std::vector"); 2611986Sandreas.sandberg@arm.com m_data = new float[m_size]; 2711986Sandreas.sandberg@arm.com memcpy(m_data, &value[0], sizeof(float) * m_size); 2811986Sandreas.sandberg@arm.com } 2911986Sandreas.sandberg@arm.com 3011986Sandreas.sandberg@arm.com Sequence(const Sequence &s) : m_size(s.m_size) { 3111986Sandreas.sandberg@arm.com print_copy_created(this); 3211986Sandreas.sandberg@arm.com m_data = new float[m_size]; 3311986Sandreas.sandberg@arm.com memcpy(m_data, s.m_data, sizeof(float)*m_size); 3411986Sandreas.sandberg@arm.com } 3511986Sandreas.sandberg@arm.com 3611986Sandreas.sandberg@arm.com Sequence(Sequence &&s) : m_size(s.m_size), m_data(s.m_data) { 3711986Sandreas.sandberg@arm.com print_move_created(this); 3811986Sandreas.sandberg@arm.com s.m_size = 0; 3911986Sandreas.sandberg@arm.com s.m_data = nullptr; 4011986Sandreas.sandberg@arm.com } 4111986Sandreas.sandberg@arm.com 4211986Sandreas.sandberg@arm.com ~Sequence() { 4311986Sandreas.sandberg@arm.com print_destroyed(this); 4411986Sandreas.sandberg@arm.com delete[] m_data; 4511986Sandreas.sandberg@arm.com } 4611986Sandreas.sandberg@arm.com 4711986Sandreas.sandberg@arm.com Sequence &operator=(const Sequence &s) { 4811986Sandreas.sandberg@arm.com if (&s != this) { 4911986Sandreas.sandberg@arm.com delete[] m_data; 5011986Sandreas.sandberg@arm.com m_size = s.m_size; 5111986Sandreas.sandberg@arm.com m_data = new float[m_size]; 5211986Sandreas.sandberg@arm.com memcpy(m_data, s.m_data, sizeof(float)*m_size); 5311986Sandreas.sandberg@arm.com } 5411986Sandreas.sandberg@arm.com 5511986Sandreas.sandberg@arm.com print_copy_assigned(this); 5611986Sandreas.sandberg@arm.com 5711986Sandreas.sandberg@arm.com return *this; 5811986Sandreas.sandberg@arm.com } 5911986Sandreas.sandberg@arm.com 6011986Sandreas.sandberg@arm.com Sequence &operator=(Sequence &&s) { 6111986Sandreas.sandberg@arm.com if (&s != this) { 6211986Sandreas.sandberg@arm.com delete[] m_data; 6311986Sandreas.sandberg@arm.com m_size = s.m_size; 6411986Sandreas.sandberg@arm.com m_data = s.m_data; 6511986Sandreas.sandberg@arm.com s.m_size = 0; 6611986Sandreas.sandberg@arm.com s.m_data = nullptr; 6711986Sandreas.sandberg@arm.com } 6811986Sandreas.sandberg@arm.com 6911986Sandreas.sandberg@arm.com print_move_assigned(this); 7011986Sandreas.sandberg@arm.com 7111986Sandreas.sandberg@arm.com return *this; 7211986Sandreas.sandberg@arm.com } 7311986Sandreas.sandberg@arm.com 7411986Sandreas.sandberg@arm.com bool operator==(const Sequence &s) const { 7511986Sandreas.sandberg@arm.com if (m_size != s.size()) 7611986Sandreas.sandberg@arm.com return false; 7711986Sandreas.sandberg@arm.com for (size_t i=0; i<m_size; ++i) 7811986Sandreas.sandberg@arm.com if (m_data[i] != s[i]) 7911986Sandreas.sandberg@arm.com return false; 8011986Sandreas.sandberg@arm.com return true; 8111986Sandreas.sandberg@arm.com } 8211986Sandreas.sandberg@arm.com 8311986Sandreas.sandberg@arm.com bool operator!=(const Sequence &s) const { 8411986Sandreas.sandberg@arm.com return !operator==(s); 8511986Sandreas.sandberg@arm.com } 8611986Sandreas.sandberg@arm.com 8711986Sandreas.sandberg@arm.com float operator[](size_t index) const { 8811986Sandreas.sandberg@arm.com return m_data[index]; 8911986Sandreas.sandberg@arm.com } 9011986Sandreas.sandberg@arm.com 9111986Sandreas.sandberg@arm.com float &operator[](size_t index) { 9211986Sandreas.sandberg@arm.com return m_data[index]; 9311986Sandreas.sandberg@arm.com } 9411986Sandreas.sandberg@arm.com 9511986Sandreas.sandberg@arm.com bool contains(float v) const { 9611986Sandreas.sandberg@arm.com for (size_t i=0; i<m_size; ++i) 9711986Sandreas.sandberg@arm.com if (v == m_data[i]) 9811986Sandreas.sandberg@arm.com return true; 9911986Sandreas.sandberg@arm.com return false; 10011986Sandreas.sandberg@arm.com } 10111986Sandreas.sandberg@arm.com 10211986Sandreas.sandberg@arm.com Sequence reversed() const { 10311986Sandreas.sandberg@arm.com Sequence result(m_size); 10411986Sandreas.sandberg@arm.com for (size_t i=0; i<m_size; ++i) 10511986Sandreas.sandberg@arm.com result[m_size-i-1] = m_data[i]; 10611986Sandreas.sandberg@arm.com return result; 10711986Sandreas.sandberg@arm.com } 10811986Sandreas.sandberg@arm.com 10911986Sandreas.sandberg@arm.com size_t size() const { return m_size; } 11011986Sandreas.sandberg@arm.com 11111986Sandreas.sandberg@arm.com const float *begin() const { return m_data; } 11211986Sandreas.sandberg@arm.com const float *end() const { return m_data+m_size; } 11311986Sandreas.sandberg@arm.com 11411986Sandreas.sandberg@arm.comprivate: 11511986Sandreas.sandberg@arm.com size_t m_size; 11611986Sandreas.sandberg@arm.com float *m_data; 11711986Sandreas.sandberg@arm.com}; 11811986Sandreas.sandberg@arm.com 11911986Sandreas.sandberg@arm.comclass IntPairs { 12011986Sandreas.sandberg@arm.compublic: 12111986Sandreas.sandberg@arm.com IntPairs(std::vector<std::pair<int, int>> data) : data_(std::move(data)) {} 12211986Sandreas.sandberg@arm.com const std::pair<int, int>* begin() const { return data_.data(); } 12311986Sandreas.sandberg@arm.com 12411986Sandreas.sandberg@arm.comprivate: 12511986Sandreas.sandberg@arm.com std::vector<std::pair<int, int>> data_; 12611986Sandreas.sandberg@arm.com}; 12711986Sandreas.sandberg@arm.com 12811986Sandreas.sandberg@arm.com// Interface of a map-like object that isn't (directly) an unordered_map, but provides some basic 12911986Sandreas.sandberg@arm.com// map-like functionality. 13011986Sandreas.sandberg@arm.comclass StringMap { 13111986Sandreas.sandberg@arm.compublic: 13211986Sandreas.sandberg@arm.com StringMap() = default; 13311986Sandreas.sandberg@arm.com StringMap(std::unordered_map<std::string, std::string> init) 13411986Sandreas.sandberg@arm.com : map(std::move(init)) {} 13511986Sandreas.sandberg@arm.com 13611986Sandreas.sandberg@arm.com void set(std::string key, std::string val) { 13711986Sandreas.sandberg@arm.com map[key] = val; 13811986Sandreas.sandberg@arm.com } 13911986Sandreas.sandberg@arm.com 14011986Sandreas.sandberg@arm.com std::string get(std::string key) const { 14111986Sandreas.sandberg@arm.com return map.at(key); 14211986Sandreas.sandberg@arm.com } 14311986Sandreas.sandberg@arm.com 14411986Sandreas.sandberg@arm.com size_t size() const { 14511986Sandreas.sandberg@arm.com return map.size(); 14611986Sandreas.sandberg@arm.com } 14711986Sandreas.sandberg@arm.com 14811986Sandreas.sandberg@arm.comprivate: 14911986Sandreas.sandberg@arm.com std::unordered_map<std::string, std::string> map; 15011986Sandreas.sandberg@arm.com 15111986Sandreas.sandberg@arm.compublic: 15211986Sandreas.sandberg@arm.com decltype(map.cbegin()) begin() const { return map.cbegin(); } 15311986Sandreas.sandberg@arm.com decltype(map.cend()) end() const { return map.cend(); } 15411986Sandreas.sandberg@arm.com}; 15511986Sandreas.sandberg@arm.com 15611986Sandreas.sandberg@arm.comtemplate<typename T> 15711986Sandreas.sandberg@arm.comclass NonZeroIterator { 15811986Sandreas.sandberg@arm.com const T* ptr_; 15911986Sandreas.sandberg@arm.compublic: 16011986Sandreas.sandberg@arm.com NonZeroIterator(const T* ptr) : ptr_(ptr) {} 16111986Sandreas.sandberg@arm.com const T& operator*() const { return *ptr_; } 16211986Sandreas.sandberg@arm.com NonZeroIterator& operator++() { ++ptr_; return *this; } 16311986Sandreas.sandberg@arm.com}; 16411986Sandreas.sandberg@arm.com 16511986Sandreas.sandberg@arm.comclass NonZeroSentinel {}; 16611986Sandreas.sandberg@arm.com 16711986Sandreas.sandberg@arm.comtemplate<typename A, typename B> 16811986Sandreas.sandberg@arm.combool operator==(const NonZeroIterator<std::pair<A, B>>& it, const NonZeroSentinel&) { 16911986Sandreas.sandberg@arm.com return !(*it).first || !(*it).second; 17011986Sandreas.sandberg@arm.com} 17111986Sandreas.sandberg@arm.com 17211986Sandreas.sandberg@arm.comtest_initializer sequences_and_iterators([](py::module &m) { 17311986Sandreas.sandberg@arm.com 17411986Sandreas.sandberg@arm.com py::class_<Sequence> seq(m, "Sequence"); 17511986Sandreas.sandberg@arm.com 17611986Sandreas.sandberg@arm.com seq.def(py::init<size_t>()) 17711986Sandreas.sandberg@arm.com .def(py::init<const std::vector<float>&>()) 17811986Sandreas.sandberg@arm.com /// Bare bones interface 17911986Sandreas.sandberg@arm.com .def("__getitem__", [](const Sequence &s, size_t i) { 18011986Sandreas.sandberg@arm.com if (i >= s.size()) 18111986Sandreas.sandberg@arm.com throw py::index_error(); 18211986Sandreas.sandberg@arm.com return s[i]; 18311986Sandreas.sandberg@arm.com }) 18411986Sandreas.sandberg@arm.com .def("__setitem__", [](Sequence &s, size_t i, float v) { 18511986Sandreas.sandberg@arm.com if (i >= s.size()) 18611986Sandreas.sandberg@arm.com throw py::index_error(); 18711986Sandreas.sandberg@arm.com s[i] = v; 18811986Sandreas.sandberg@arm.com }) 18911986Sandreas.sandberg@arm.com .def("__len__", &Sequence::size) 19011986Sandreas.sandberg@arm.com /// Optional sequence protocol operations 19111986Sandreas.sandberg@arm.com .def("__iter__", [](const Sequence &s) { return py::make_iterator(s.begin(), s.end()); }, 19211986Sandreas.sandberg@arm.com py::keep_alive<0, 1>() /* Essential: keep object alive while iterator exists */) 19311986Sandreas.sandberg@arm.com .def("__contains__", [](const Sequence &s, float v) { return s.contains(v); }) 19411986Sandreas.sandberg@arm.com .def("__reversed__", [](const Sequence &s) -> Sequence { return s.reversed(); }) 19511986Sandreas.sandberg@arm.com /// Slicing protocol (optional) 19611986Sandreas.sandberg@arm.com .def("__getitem__", [](const Sequence &s, py::slice slice) -> Sequence* { 19711986Sandreas.sandberg@arm.com size_t start, stop, step, slicelength; 19811986Sandreas.sandberg@arm.com if (!slice.compute(s.size(), &start, &stop, &step, &slicelength)) 19911986Sandreas.sandberg@arm.com throw py::error_already_set(); 20011986Sandreas.sandberg@arm.com Sequence *seq = new Sequence(slicelength); 20111986Sandreas.sandberg@arm.com for (size_t i=0; i<slicelength; ++i) { 20211986Sandreas.sandberg@arm.com (*seq)[i] = s[start]; start += step; 20311986Sandreas.sandberg@arm.com } 20411986Sandreas.sandberg@arm.com return seq; 20511986Sandreas.sandberg@arm.com }) 20611986Sandreas.sandberg@arm.com .def("__setitem__", [](Sequence &s, py::slice slice, const Sequence &value) { 20711986Sandreas.sandberg@arm.com size_t start, stop, step, slicelength; 20811986Sandreas.sandberg@arm.com if (!slice.compute(s.size(), &start, &stop, &step, &slicelength)) 20911986Sandreas.sandberg@arm.com throw py::error_already_set(); 21011986Sandreas.sandberg@arm.com if (slicelength != value.size()) 21111986Sandreas.sandberg@arm.com throw std::runtime_error("Left and right hand size of slice assignment have different sizes!"); 21211986Sandreas.sandberg@arm.com for (size_t i=0; i<slicelength; ++i) { 21311986Sandreas.sandberg@arm.com s[start] = value[i]; start += step; 21411986Sandreas.sandberg@arm.com } 21511986Sandreas.sandberg@arm.com }) 21611986Sandreas.sandberg@arm.com /// Comparisons 21711986Sandreas.sandberg@arm.com .def(py::self == py::self) 21811986Sandreas.sandberg@arm.com .def(py::self != py::self); 21911986Sandreas.sandberg@arm.com // Could also define py::self + py::self for concatenation, etc. 22011986Sandreas.sandberg@arm.com 22111986Sandreas.sandberg@arm.com py::class_<StringMap> map(m, "StringMap"); 22211986Sandreas.sandberg@arm.com 22311986Sandreas.sandberg@arm.com map .def(py::init<>()) 22411986Sandreas.sandberg@arm.com .def(py::init<std::unordered_map<std::string, std::string>>()) 22511986Sandreas.sandberg@arm.com .def("__getitem__", [](const StringMap &map, std::string key) { 22611986Sandreas.sandberg@arm.com try { return map.get(key); } 22711986Sandreas.sandberg@arm.com catch (const std::out_of_range&) { 22811986Sandreas.sandberg@arm.com throw py::key_error("key '" + key + "' does not exist"); 22911986Sandreas.sandberg@arm.com } 23011986Sandreas.sandberg@arm.com }) 23111986Sandreas.sandberg@arm.com .def("__setitem__", &StringMap::set) 23211986Sandreas.sandberg@arm.com .def("__len__", &StringMap::size) 23311986Sandreas.sandberg@arm.com .def("__iter__", [](const StringMap &map) { return py::make_key_iterator(map.begin(), map.end()); }, 23411986Sandreas.sandberg@arm.com py::keep_alive<0, 1>()) 23511986Sandreas.sandberg@arm.com .def("items", [](const StringMap &map) { return py::make_iterator(map.begin(), map.end()); }, 23611986Sandreas.sandberg@arm.com py::keep_alive<0, 1>()) 23711986Sandreas.sandberg@arm.com ; 23811986Sandreas.sandberg@arm.com 23911986Sandreas.sandberg@arm.com py::class_<IntPairs>(m, "IntPairs") 24011986Sandreas.sandberg@arm.com .def(py::init<std::vector<std::pair<int, int>>>()) 24111986Sandreas.sandberg@arm.com .def("nonzero", [](const IntPairs& s) { 24211986Sandreas.sandberg@arm.com return py::make_iterator(NonZeroIterator<std::pair<int, int>>(s.begin()), NonZeroSentinel()); 24311986Sandreas.sandberg@arm.com }, py::keep_alive<0, 1>()) 24411986Sandreas.sandberg@arm.com .def("nonzero_keys", [](const IntPairs& s) { 24511986Sandreas.sandberg@arm.com return py::make_key_iterator(NonZeroIterator<std::pair<int, int>>(s.begin()), NonZeroSentinel()); 24611986Sandreas.sandberg@arm.com }, py::keep_alive<0, 1>()); 24711986Sandreas.sandberg@arm.com 24811986Sandreas.sandberg@arm.com 24911986Sandreas.sandberg@arm.com#if 0 25011986Sandreas.sandberg@arm.com // Obsolete: special data structure for exposing custom iterator types to python 25111986Sandreas.sandberg@arm.com // kept here for illustrative purposes because there might be some use cases which 25211986Sandreas.sandberg@arm.com // are not covered by the much simpler py::make_iterator 25311986Sandreas.sandberg@arm.com 25411986Sandreas.sandberg@arm.com struct PySequenceIterator { 25511986Sandreas.sandberg@arm.com PySequenceIterator(const Sequence &seq, py::object ref) : seq(seq), ref(ref) { } 25611986Sandreas.sandberg@arm.com 25711986Sandreas.sandberg@arm.com float next() { 25811986Sandreas.sandberg@arm.com if (index == seq.size()) 25911986Sandreas.sandberg@arm.com throw py::stop_iteration(); 26011986Sandreas.sandberg@arm.com return seq[index++]; 26111986Sandreas.sandberg@arm.com } 26211986Sandreas.sandberg@arm.com 26311986Sandreas.sandberg@arm.com const Sequence &seq; 26411986Sandreas.sandberg@arm.com py::object ref; // keep a reference 26511986Sandreas.sandberg@arm.com size_t index = 0; 26611986Sandreas.sandberg@arm.com }; 26711986Sandreas.sandberg@arm.com 26811986Sandreas.sandberg@arm.com py::class_<PySequenceIterator>(seq, "Iterator") 26911986Sandreas.sandberg@arm.com .def("__iter__", [](PySequenceIterator &it) -> PySequenceIterator& { return it; }) 27011986Sandreas.sandberg@arm.com .def("__next__", &PySequenceIterator::next); 27111986Sandreas.sandberg@arm.com 27211986Sandreas.sandberg@arm.com On the actual Sequence object, the iterator would be constructed as follows: 27311986Sandreas.sandberg@arm.com .def("__iter__", [](py::object s) { return PySequenceIterator(s.cast<const Sequence &>(), s); }) 27411986Sandreas.sandberg@arm.com#endif 27511986Sandreas.sandberg@arm.com}); 276