From 4c5e21b0cbe3fcf9c9972d385215c199598aad86 Mon Sep 17 00:00:00 2001
From: Ivan Smirnov <i.s.smirnov@gmail.com>
Date: Wed, 24 Aug 2016 23:30:00 +0100
Subject: [PATCH] Add tests for generalized iterators

---
 tests/test_sequences_and_iterators.cpp | 34 ++++++++++++++++++++++++++
 tests/test_sequences_and_iterators.py  | 12 +++++++++
 2 files changed, 46 insertions(+)

diff --git a/tests/test_sequences_and_iterators.cpp b/tests/test_sequences_and_iterators.cpp
index a92c6bf62..39e342ba6 100644
--- a/tests/test_sequences_and_iterators.cpp
+++ b/tests/test_sequences_and_iterators.cpp
@@ -116,6 +116,15 @@ private:
     float *m_data;
 };
 
+class IntPairs {
+public:
+    IntPairs(std::vector<std::pair<int, int>> data) : data_(std::move(data)) {}
+    const std::pair<int, int>* begin() const { return data_.data(); }
+
+private:
+    std::vector<std::pair<int, int>> data_;
+};
+
 // Interface of a map-like object that isn't (directly) an unordered_map, but provides some basic
 // map-like functionality.
 class StringMap {
@@ -143,8 +152,24 @@ public:
     decltype(map.cend()) end() const { return map.cend(); }
 };
 
+template<typename T>
+class NonZeroIterator {
+    const T* ptr_;
+public:
+    NonZeroIterator(const T* ptr) : ptr_(ptr) {}
+    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;
+}
 
 void init_ex_sequences_and_iterators(py::module &m) {
+
     py::class_<Sequence> seq(m, "Sequence");
 
     seq.def(py::init<size_t>())
@@ -210,6 +235,15 @@ void init_ex_sequences_and_iterators(py::module &m) {
                 py::keep_alive<0, 1>())
         ;
 
+    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());
+            }, py::keep_alive<0, 1>())
+        .def("nonzero_keys", [](const IntPairs& s) {
+            return py::make_key_iterator(NonZeroIterator<std::pair<int, int>>(s.begin()), NonZeroSentinel());
+        }, py::keep_alive<0, 1>());
+
 
 #if 0
     // Obsolete: special data structure for exposing custom iterator types to python
diff --git a/tests/test_sequences_and_iterators.py b/tests/test_sequences_and_iterators.py
index a35dc584b..c83c4e57c 100644
--- a/tests/test_sequences_and_iterators.py
+++ b/tests/test_sequences_and_iterators.py
@@ -10,6 +10,18 @@ def allclose(a_list, b_list, rel_tol=1e-05, abs_tol=0.0):
     return all(isclose(a, b, rel_tol=rel_tol, abs_tol=abs_tol) for a, b in zip(a_list, b_list))
 
 
+def test_generalized_iterators():
+    from pybind11_tests import IntPairs
+
+    assert list(IntPairs([(1, 2), (3, 4), (0, 5)]).nonzero()) == [(1, 2), (3, 4)]
+    assert list(IntPairs([(1, 2), (2, 0), (0, 3), (4, 5)]).nonzero()) == [(1, 2)]
+    assert list(IntPairs([(0, 3), (1, 2), (3, 4)]).nonzero()) == []
+
+    assert list(IntPairs([(1, 2), (3, 4), (0, 5)]).nonzero_keys()) == [1, 3]
+    assert list(IntPairs([(1, 2), (2, 0), (0, 3), (4, 5)]).nonzero_keys()) == [1]
+    assert list(IntPairs([(0, 3), (1, 2), (3, 4)]).nonzero_keys()) == []
+
+
 def test_sequence():
     from pybind11_tests import Sequence, ConstructorStats