Merge pull request #282 from jagerman/key-iterators

Add pybind11::make_key_iterator for map iteration

docs/advanced.rst

@@ -957,6 +957,12 @@ exceptions:
 |                                      | indicate wrong value passed  |
 |                                      | in ``container.remove(...)`` |
 +--------------------------------------+------------------------------+
+| :class:`pybind11::key_error`         | ``KeyError`` (used to        |
+|                                      | indicate out of bounds       |
+|                                      | accesses in ``__getitem__``, |
+|                                      | ``__setitem__`` in dict-like |
+|                                      | objects, etc.)               |
++--------------------------------------+------------------------------+
 | :class:`pybind11::error_already_set` | Indicates that the Python    |
 |                                      | exception flag has already   |
 |                                      | been initialized             |

example/example-sequences-and-iterators.cpp

@@ -116,6 +116,34 @@ private:
     float *m_data;
 };
 
+// Interface of a map-like object that isn't (directly) an unordered_map, but provides some basic
+// map-like functionality.
+class StringMap {
+public:
+    StringMap(std::unordered_map<std::string, std::string> init = {})
+        : map(std::move(init)) {}
+
+    void set(std::string key, std::string val) {
+        map[key] = val;
+    }
+
+    std::string get(std::string key) const {
+        return map.at(key);
+    }
+
+    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(); }
+};
+
+
 void init_ex_sequences_and_iterators(py::module &m) {
     py::class_<Sequence> seq(m, "Sequence");
 
@@ -164,6 +192,25 @@ void init_ex_sequences_and_iterators(py::module &m) {
        .def(py::self != py::self);
        // Could also define py::self + py::self for concatenation, etc.
 
+    py::class_<StringMap> map(m, "StringMap");
+
+    map .def(py::init<>())
+        .def(py::init<std::unordered_map<std::string, std::string>>())
+        .def("__getitem__", [](const StringMap &map, std::string key) {
+                try { return map.get(key); }
+                catch (const std::out_of_range&) {
+                    throw py::key_error("key '" + key + "' does not exist");
+                }
+                })
+        .def("__setitem__", &StringMap::set)
+        .def("__len__", &StringMap::size)
+        .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()); },
+                py::keep_alive<0, 1>())
+        ;
+
+
 #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

example/example-sequences-and-iterators.py

@@ -3,7 +3,7 @@ from __future__ import print_function
 import sys
 sys.path.append('.')
 
-from example import Sequence
+from example import Sequence, StringMap
 
 s = Sequence(5)
 print("s = " + str(s))
@@ -29,6 +29,24 @@ for i in rev:
     print(i, end=' ')
 print('')
 
+m = StringMap({ 'hi': 'bye', 'black': 'white' })
+print(m['hi'])
+print(len(m))
+print(m['black'])
+try:
+    print(m['orange'])
+    print('Error: should have thrown exception')
+except KeyError:
+    pass
+m['orange'] = 'banana'
+print(m['orange'])
+
+for k in m:
+    print("key = %s, value = %s" % (k, m[k]))
+
+for k,v in m.items():
+    print("item: (%s, %s)" % (k,v))
+
 from example import ConstructorStats
 cstats = ConstructorStats.get(Sequence)
 print("Instances not destroyed:", cstats.alive())

example/example-sequences-and-iterators.ref

@@ -13,9 +13,19 @@ rev[0], rev[1], rev[2], rev[3], rev[4] = 0.000000 56.779999 0.000000 0.000000 12
 0.0 56.779998779296875 0.0 0.0 12.34000015258789 
 0.0 56.779998779296875 0.0 0.0 12.34000015258789 
 True
-### Sequence @ 0x153c4b0 created of size 3  from std::vector
-### Sequence @ 0x153c4b0 destroyed
+### Sequence @ 0x1b4d1f0 created of size 3 from std::vector
+### Sequence @ 0x1b4d1f0 destroyed
 2.0 56.779998779296875 2.0 0.0 2.0 
+bye
+2
+white
+banana
+key = orange, value = banana
+key = hi, value = bye
+key = black, value = white
+item: (orange, banana)
+item: (hi, bye)
+item: (black, white)
 Instances not destroyed: 3
 ### Sequence @ 0x1535b00 destroyed
 Instances not destroyed: 2

include/pybind11/cast.h

@@ -55,6 +55,7 @@ PYBIND11_NOINLINE inline internals &get_internals() {
                     if (p) std::rethrow_exception(p);
                 } catch (const error_already_set &)      {                                                 return;
                 } catch (const index_error &e)           { PyErr_SetString(PyExc_IndexError,    e.what()); return;
+                } catch (const key_error &e)             { PyErr_SetString(PyExc_KeyError,      e.what()); return;
                 } catch (const value_error &e)           { PyErr_SetString(PyExc_ValueError,    e.what()); return;
                 } catch (const stop_iteration &e)        { PyErr_SetString(PyExc_StopIteration, e.what()); return;
                 } catch (const std::bad_alloc &e)        { PyErr_SetString(PyExc_MemoryError,   e.what()); return;

include/pybind11/common.h

@@ -314,6 +314,7 @@ NAMESPACE_END(detail)
 class error_already_set : public std::runtime_error { public: error_already_set() : std::runtime_error(detail::error_string())  {} };
 PYBIND11_RUNTIME_EXCEPTION(stop_iteration)
 PYBIND11_RUNTIME_EXCEPTION(index_error)
+PYBIND11_RUNTIME_EXCEPTION(key_error)
 PYBIND11_RUNTIME_EXCEPTION(value_error)
 PYBIND11_RUNTIME_EXCEPTION(cast_error) /// Thrown when pybind11::cast or handle::call fail due to a type casting error
 PYBIND11_RUNTIME_EXCEPTION(reference_cast_error) /// Used internally

include/pybind11/pybind11.h

@@ -1117,7 +1117,7 @@ PYBIND11_NOINLINE inline void keep_alive_impl(int Nurse, int Patient, handle arg
     keep_alive_impl(nurse, patient);
 }
 
-template <typename Iterator> struct iterator_state {
+template <typename Iterator, bool KeyIterator = false> struct iterator_state {
     Iterator it, end;
     bool first;
 };
@@ -1148,11 +1148,37 @@ iterator make_iterator(Iterator first, Iterator last, Extra &&... extra) {
 
     return (iterator) cast(state { first, last, true });
 }
+template <typename Iterator,
+          typename KeyType = decltype(std::declval<Iterator>()->first),
+          typename... Extra>
+iterator make_key_iterator(Iterator first, Iterator last, Extra &&... extra) {
+    typedef detail::iterator_state<Iterator, true> state;
+
+    if (!detail::get_type_info(typeid(state))) {
+        class_<state>(handle(), "")
+            .def("__iter__", [](state &s) -> state& { return s; })
+            .def("__next__", [](state &s) -> KeyType {
+                if (!s.first)
+                    ++s.it;
+                else
+                    s.first = false;
+                if (s.it == s.end)
+                    throw stop_iteration();
+                return s.it->first;
+            }, return_value_policy::reference_internal, std::forward<Extra>(extra)...);
+    }
+
+    return (iterator) cast(state { first, last, true });
+}
 
 template <typename Type, typename... Extra> iterator make_iterator(Type &value, Extra&&... extra) {
     return make_iterator(std::begin(value), std::end(value), extra...);
 }
 
+template <typename Type, typename... Extra> iterator make_key_iterator(Type &value, Extra&&... extra) {
+    return make_key_iterator(std::begin(value), std::end(value), extra...);
+}
+
 template <typename InputType, typename OutputType> void implicitly_convertible() {
     auto implicit_caster = [](PyObject *obj, PyTypeObject *type) -> PyObject * {
         if (!detail::type_caster<InputType>().load(obj, false))