test_sequences_and_iterators.cpp 15.8 KB
Newer Older
Wenzel Jakob's avatar
Wenzel Jakob committed
1
/*
Dean Moldovan's avatar
Dean Moldovan committed
2
    tests/test_sequences_and_iterators.cpp -- supporting Pythons' sequence protocol, iterators,
3
    etc.
Wenzel Jakob's avatar
Wenzel Jakob committed
4

5
    Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
Wenzel Jakob's avatar
Wenzel Jakob committed
6
7
8
9
10

    All rights reserved. Use of this source code is governed by a
    BSD-style license that can be found in the LICENSE file.
*/

Dean Moldovan's avatar
Dean Moldovan committed
11
12
#include "pybind11_tests.h"
#include "constructor_stats.h"
13
14
#include <pybind11/operators.h>
#include <pybind11/stl.h>
Wenzel Jakob's avatar
Wenzel Jakob committed
15

16
#include <algorithm>
17
#include <utility>
Bruce Merry's avatar
Bruce Merry committed
18
#include <vector>
19

20
21
22
23
template<typename T>
class NonZeroIterator {
    const T* ptr_;
public:
24
    explicit NonZeroIterator(const T *ptr) : ptr_(ptr) {}
25
26
27
28
29
30
31
32
33
34
    const T& operator*() const { return *ptr_; }
    NonZeroIterator& operator++() { ++ptr_; return *this; }
};

class NonZeroSentinel {};

template<typename A, typename B>
bool operator==(const NonZeroIterator<std::pair<A, B>>& it, const NonZeroSentinel&) {
    return !(*it).first || !(*it).second;
}
35

Bruce Merry's avatar
Bruce Merry committed
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
class NonCopyableInt {
public:
    explicit NonCopyableInt(int value) : value_(value) {}
    NonCopyableInt(const NonCopyableInt &) = delete;
    NonCopyableInt(NonCopyableInt &&other) noexcept : value_(other.value_) {
        other.value_ = -1;  // detect when an unwanted move occurs
    }
    NonCopyableInt &operator=(const NonCopyableInt &) = delete;
    NonCopyableInt &operator=(NonCopyableInt &&other) noexcept {
        value_ = other.value_;
        other.value_ = -1;  // detect when an unwanted move occurs
        return *this;
    }
    int get() const { return value_; }
    void set(int value) { value_ = value; }
    ~NonCopyableInt() = default;
private:
    int value_;
};
using NonCopyableIntPair = std::pair<NonCopyableInt, NonCopyableInt>;
PYBIND11_MAKE_OPAQUE(std::vector<NonCopyableInt>);
PYBIND11_MAKE_OPAQUE(std::vector<NonCopyableIntPair>);

59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
template <typename PythonType>
py::list test_random_access_iterator(PythonType x) {
    if (x.size() < 5)
        throw py::value_error("Please provide at least 5 elements for testing.");

    auto checks = py::list();
    auto assert_equal = [&checks](py::handle a, py::handle b) {
        auto result = PyObject_RichCompareBool(a.ptr(), b.ptr(), Py_EQ);
        if (result == -1) { throw py::error_already_set(); }
        checks.append(result != 0);
    };

    auto it = x.begin();
    assert_equal(x[0], *it);
    assert_equal(x[0], it[0]);
    assert_equal(x[1], it[1]);

    assert_equal(x[1], *(++it));
    assert_equal(x[1], *(it++));
    assert_equal(x[2], *it);
    assert_equal(x[3], *(it += 1));
    assert_equal(x[2], *(--it));
    assert_equal(x[2], *(it--));
    assert_equal(x[1], *it);
    assert_equal(x[0], *(it -= 1));

    assert_equal(it->attr("real"), x[0].attr("real"));
    assert_equal((it + 1)->attr("real"), x[1].attr("real"));

    assert_equal(x[1], *(it + 1));
    assert_equal(x[1], *(1 + it));
    it += 3;
    assert_equal(x[1], *(it - 2));

    checks.append(static_cast<std::size_t>(x.end() - x.begin()) == x.size());
    checks.append((x.begin() + static_cast<std::ptrdiff_t>(x.size())) == x.end());
    checks.append(x.begin() < x.end());

    return checks;
}

100
TEST_SUBMODULE(sequences_and_iterators, m) {
101
102
103
    // test_sliceable
    class Sliceable{
    public:
104
105
106
        explicit Sliceable(int n) : size(n) {}
        int start, stop, step;
        int size;
107
    };
108
    py::class_<Sliceable>(m, "Sliceable")
109
        .def(py::init<int>())
110
        .def("__getitem__", [](const Sliceable &s, const py::slice &slice) {
111
            py::ssize_t start = 0, stop = 0, step = 0, slicelength = 0;
112
113
114
115
116
117
118
            if (!slice.compute(s.size, &start, &stop, &step, &slicelength))
                throw py::error_already_set();
            int istart = static_cast<int>(start);
            int istop  = static_cast<int>(stop);
            int istep  = static_cast<int>(step);
            return std::make_tuple(istart, istop, istep);
        });
119

120
121
122
    // test_sequence
    class Sequence {
    public:
123
        explicit Sequence(size_t size) : m_size(size) {
124
125
126
127
            print_created(this, "of size", m_size);
            m_data = new float[size];
            memset(m_data, 0, sizeof(float) * size);
        }
128
        explicit Sequence(const std::vector<float> &value) : m_size(value.size()) {
129
130
131
132
133
134
135
136
137
            print_created(this, "of size", m_size, "from std::vector");
            m_data = new float[m_size];
            memcpy(m_data, &value[0], sizeof(float) * m_size);
        }
        Sequence(const Sequence &s) : m_size(s.m_size) {
            print_copy_created(this);
            m_data = new float[m_size];
            memcpy(m_data, s.m_data, sizeof(float)*m_size);
        }
138
        Sequence(Sequence &&s) noexcept : m_size(s.m_size), m_data(s.m_data) {
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
            print_move_created(this);
            s.m_size = 0;
            s.m_data = nullptr;
        }

        ~Sequence() { print_destroyed(this); delete[] m_data; }

        Sequence &operator=(const Sequence &s) {
            if (&s != this) {
                delete[] m_data;
                m_size = s.m_size;
                m_data = new float[m_size];
                memcpy(m_data, s.m_data, sizeof(float)*m_size);
            }
            print_copy_assigned(this);
            return *this;
        }

157
        Sequence &operator=(Sequence &&s) noexcept {
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
            if (&s != this) {
                delete[] m_data;
                m_size = s.m_size;
                m_data = s.m_data;
                s.m_size = 0;
                s.m_data = nullptr;
            }
            print_move_assigned(this);
            return *this;
        }

        bool operator==(const Sequence &s) const {
            if (m_size != s.size()) return false;
            for (size_t i = 0; i < m_size; ++i)
                if (m_data[i] != s[i])
                    return false;
            return true;
        }
        bool operator!=(const Sequence &s) const { return !operator==(s); }

        float operator[](size_t index) const { return m_data[index]; }
        float &operator[](size_t index) { return m_data[index]; }
Wenzel Jakob's avatar
Wenzel Jakob committed
180

181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
        bool contains(float v) const {
            for (size_t i = 0; i < m_size; ++i)
                if (v == m_data[i])
                    return true;
            return false;
        }

        Sequence reversed() const {
            Sequence result(m_size);
            for (size_t i = 0; i < m_size; ++i)
                result[m_size - i - 1] = m_data[i];
            return result;
        }

        size_t size() const { return m_size; }

        const float *begin() const { return m_data; }
        const float *end() const { return m_data+m_size; }

    private:
        size_t m_size;
        float *m_data;
    };
    py::class_<Sequence>(m, "Sequence")
        .def(py::init<size_t>())
206
        .def(py::init<const std::vector<float> &>())
207
        /// Bare bones interface
208
209
210
211
212
213
214
215
216
217
218
219
        .def("__getitem__",
             [](const Sequence &s, size_t i) {
                 if (i >= s.size())
                     throw py::index_error();
                 return s[i];
             })
        .def("__setitem__",
             [](Sequence &s, size_t i, float v) {
                 if (i >= s.size())
                     throw py::index_error();
                 s[i] = v;
             })
220
221
        .def("__len__", &Sequence::size)
        /// Optional sequence protocol operations
222
223
224
225
        .def(
            "__iter__",
            [](const Sequence &s) { return py::make_iterator(s.begin(), s.end()); },
            py::keep_alive<0, 1>() /* Essential: keep object alive while iterator exists */)
226
227
228
        .def("__contains__", [](const Sequence &s, float v) { return s.contains(v); })
        .def("__reversed__", [](const Sequence &s) -> Sequence { return s.reversed(); })
        /// Slicing protocol (optional)
229
230
        .def("__getitem__",
             [](const Sequence &s, const py::slice &slice) -> Sequence * {
231
                 size_t start = 0, stop = 0, step = 0, slicelength = 0;
232
233
234
235
236
237
238
239
240
241
242
                 if (!slice.compute(s.size(), &start, &stop, &step, &slicelength))
                     throw py::error_already_set();
                 auto *seq = new Sequence(slicelength);
                 for (size_t i = 0; i < slicelength; ++i) {
                     (*seq)[i] = s[start];
                     start += step;
                 }
                 return seq;
             })
        .def("__setitem__",
             [](Sequence &s, const py::slice &slice, const Sequence &value) {
243
                 size_t start = 0, stop = 0, step = 0, slicelength = 0;
244
245
246
247
248
249
250
251
252
253
                 if (!slice.compute(s.size(), &start, &stop, &step, &slicelength))
                     throw py::error_already_set();
                 if (slicelength != value.size())
                     throw std::runtime_error(
                         "Left and right hand size of slice assignment have different sizes!");
                 for (size_t i = 0; i < slicelength; ++i) {
                     s[start] = value[i];
                     start += step;
                 }
             })
254
255
256
257
258
        /// Comparisons
        .def(py::self == py::self)
        .def(py::self != py::self)
        // Could also define py::self + py::self for concatenation, etc.
        ;
259

260
261
262
263
264
265
    // test_map_iterator
    // Interface of a map-like object that isn't (directly) an unordered_map, but provides some basic
    // map-like functionality.
    class StringMap {
    public:
        StringMap() = default;
266
        explicit StringMap(std::unordered_map<std::string, std::string> init)
267
268
            : map(std::move(init)) {}

269
270
        void set(const std::string &key, std::string val) { map[key] = std::move(val); }
        std::string get(const std::string &key) const { return map.at(key); }
271
272
273
274
275
276
277
278
279
        size_t size() const { return map.size(); }
    private:
        std::unordered_map<std::string, std::string> map;
    public:
        decltype(map.cbegin()) begin() const { return map.cbegin(); }
        decltype(map.cend()) end() const { return map.cend(); }
    };
    py::class_<StringMap>(m, "StringMap")
        .def(py::init<>())
280
        .def(py::init<std::unordered_map<std::string, std::string>>())
281
282
283
284
285
286
287
288
        .def("__getitem__",
             [](const StringMap &map, const std::string &key) {
                 try {
                     return map.get(key);
                 } catch (const std::out_of_range &) {
                     throw py::key_error("key '" + key + "' does not exist");
                 }
             })
289
290
        .def("__setitem__", &StringMap::set)
        .def("__len__", &StringMap::size)
291
292
293
294
295
296
297
        .def(
            "__iter__",
            [](const StringMap &map) { return py::make_key_iterator(map.begin(), map.end()); },
            py::keep_alive<0, 1>())
        .def(
            "items",
            [](const StringMap &map) { return py::make_iterator(map.begin(), map.end()); },
Bruce Merry's avatar
Bruce Merry committed
298
299
300
301
            py::keep_alive<0, 1>())
        .def(
            "values",
            [](const StringMap &map) { return py::make_value_iterator(map.begin(), map.end()); },
302
            py::keep_alive<0, 1>());
303

304
305
306
    // test_generalized_iterators
    class IntPairs {
    public:
307
        explicit IntPairs(std::vector<std::pair<int, int>> data) : data_(std::move(data)) {}
308
309
310
311
        const std::pair<int, int>* begin() const { return data_.data(); }
    private:
        std::vector<std::pair<int, int>> data_;
    };
312
313
314
315
    py::class_<IntPairs>(m, "IntPairs")
        .def(py::init<std::vector<std::pair<int, int>>>())
        .def("nonzero", [](const IntPairs& s) {
                return py::make_iterator(NonZeroIterator<std::pair<int, int>>(s.begin()), NonZeroSentinel());
316
        }, py::keep_alive<0, 1>())
317
318
        .def("nonzero_keys", [](const IntPairs& s) {
            return py::make_key_iterator(NonZeroIterator<std::pair<int, int>>(s.begin()), NonZeroSentinel());
319
        }, py::keep_alive<0, 1>())
Bruce Merry's avatar
Bruce Merry committed
320
321
322
        .def("nonzero_values", [](const IntPairs& s) {
            return py::make_value_iterator(NonZeroIterator<std::pair<int, int>>(s.begin()), NonZeroSentinel());
        }, py::keep_alive<0, 1>())
323
        ;
324

Bruce Merry's avatar
Bruce Merry committed
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
    // test_iterater_referencing
    py::class_<NonCopyableInt>(m, "NonCopyableInt")
        .def(py::init<int>())
        .def("set", &NonCopyableInt::set)
        .def("__int__", &NonCopyableInt::get)
        ;
    py::class_<std::vector<NonCopyableInt>>(m, "VectorNonCopyableInt")
        .def(py::init<>())
        .def("append", [](std::vector<NonCopyableInt> &vec, int value) {
            vec.emplace_back(value);
        })
        .def("__iter__", [](std::vector<NonCopyableInt> &vec) {
            return py::make_iterator(vec.begin(), vec.end());
        })
        ;
    py::class_<std::vector<NonCopyableIntPair>>(m, "VectorNonCopyableIntPair")
        .def(py::init<>())
        .def("append", [](std::vector<NonCopyableIntPair> &vec, const std::pair<int, int> &value) {
            vec.emplace_back(NonCopyableInt(value.first), NonCopyableInt(value.second));
        })
        .def("keys", [](std::vector<NonCopyableIntPair> &vec) {
            return py::make_key_iterator(vec.begin(), vec.end());
        })
        .def("values", [](std::vector<NonCopyableIntPair> &vec) {
            return py::make_value_iterator(vec.begin(), vec.end());
        })
        ;
352

353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
#if 0
    // Obsolete: special data structure for exposing custom iterator types to python
    // kept here for illustrative purposes because there might be some use cases which
    // are not covered by the much simpler py::make_iterator

    struct PySequenceIterator {
        PySequenceIterator(const Sequence &seq, py::object ref) : seq(seq), ref(ref) { }

        float next() {
            if (index == seq.size())
                throw py::stop_iteration();
            return seq[index++];
        }

        const Sequence &seq;
        py::object ref; // keep a reference
        size_t index = 0;
    };

Wenzel Jakob's avatar
Wenzel Jakob committed
372
373
374
    py::class_<PySequenceIterator>(seq, "Iterator")
        .def("__iter__", [](PySequenceIterator &it) -> PySequenceIterator& { return it; })
        .def("__next__", &PySequenceIterator::next);
375
376
377
378

    On the actual Sequence object, the iterator would be constructed as follows:
    .def("__iter__", [](py::object s) { return PySequenceIterator(s.cast<const Sequence &>(), s); })
#endif
Dean Moldovan's avatar
Dean Moldovan committed
379

380
    // test_python_iterator_in_cpp
381
    m.def("object_to_list", [](const py::object &o) {
Dean Moldovan's avatar
Dean Moldovan committed
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
        auto l = py::list();
        for (auto item : o) {
            l.append(item);
        }
        return l;
    });

    m.def("iterator_to_list", [](py::iterator it) {
        auto l = py::list();
        while (it != py::iterator::sentinel()) {
            l.append(*it);
            ++it;
        }
        return l;
    });
397

398
    // test_sequence_length: check that Python sequences can be converted to py::sequence.
399
    m.def("sequence_length", [](const py::sequence &seq) { return seq.size(); });
400

401
    // Make sure that py::iterator works with std algorithms
402
    m.def("count_none", [](const py::object &o) {
403
404
405
        return std::count_if(o.begin(), o.end(), [](py::handle h) { return h.is_none(); });
    });

406
    m.def("find_none", [](const py::object &o) {
407
408
409
        auto it = std::find_if(o.begin(), o.end(), [](py::handle h) { return h.is_none(); });
        return it->is_none();
    });
410

411
412
413
414
    m.def("count_nonzeros", [](const py::dict &d) {
        return std::count_if(d.begin(), d.end(), [](std::pair<py::handle, py::handle> p) {
            return p.second.cast<int>() != 0;
        });
415
416
    });

417
418
419
    m.def("tuple_iterator", &test_random_access_iterator<py::tuple>);
    m.def("list_iterator", &test_random_access_iterator<py::list>);
    m.def("sequence_iterator", &test_random_access_iterator<py::sequence>);
420

421
    // test_iterator_passthrough
422
423
424
425
426
    // #181: iterator passthrough did not compile
    m.def("iterator_passthrough", [](py::iterator s) -> py::iterator {
        return py::make_iterator(std::begin(s), std::end(s));
    });

427
    // test_iterator_rvp
428
429
430
431
    // #388: Can't make iterators via make_iterator() with different r/v policies
    static std::vector<int> list = { 1, 2, 3 };
    m.def("make_iterator_1", []() { return py::make_iterator<py::return_value_policy::copy>(list); });
    m.def("make_iterator_2", []() { return py::make_iterator<py::return_value_policy::automatic>(list); });
432
}