[vector] Supported list/tuple indexing in vector slicing

Author: cmastalli

include/eigenpy/std-vector.hpp

@@ -86,7 +86,9 @@ struct overload_base_get_item_for_std_vector
   template <class Class>
   void visit(Class &cl) const {
     cl.def("__getitem__", &base_get_item_int)
-      .def("__getitem__", &base_get_item_slice);
+      .def("__getitem__", &base_get_item_slice)
+      .def("__getitem__", &base_get_item_list)
+      .def("__getitem__", &base_get_item_tuple);
   }
 
  private:
@@ -131,6 +133,54 @@ struct overload_base_get_item_for_std_vector
     return out;
   }
 
+  static bp::object base_get_item_list(bp::back_reference<Container&> c, bp::list idxs) {
+    const Py_ssize_t m = bp::len(idxs);
+    bp::list out;
+    for (Py_ssize_t k = 0; k < m; ++k) {
+      bp::object obj = idxs[k];
+      bp::extract<long> ei(obj);
+      if (!ei.check()) {
+        PyErr_SetString(PyExc_TypeError, "indices must be integers");
+        bp::throw_error_already_set();
+      }
+      auto idx = normalize_index(c.get().size(), ei());
+      out.append(elem_ref(c.get(), idx));
+    }
+    return out;
+  }
+
+  static bp::object base_get_item_tuple(bp::back_reference<Container&> c, bp::tuple idxs) {
+    const Py_ssize_t m = bp::len(idxs);
+    bp::list out;
+    for (Py_ssize_t k = 0; k < m; ++k) {
+      bp::object obj = idxs[k];
+      bp::extract<long> ei(obj);
+      if (!ei.check()) {
+        PyErr_SetString(PyExc_TypeError, "indices must be integers");
+        bp::throw_error_already_set();
+      }
+      auto idx = normalize_index(c.get().size(), ei());
+      out.append(elem_ref(c.get(), idx));
+    }
+    return out;
+  }
+
+  static index_type normalize_index(std::size_t n, long i) {
+    long idx = i;
+    if (idx < 0) idx += static_cast<long>(n);
+    if (idx < 0 || idx >= static_cast<long>(n)) {
+      PyErr_SetString(PyExc_IndexError, "index out of range");
+      bp::throw_error_already_set();
+    }
+    return static_cast<index_type>(idx);
+  }
+
+  static bp::object elem_ref(Container& c, index_type i) {
+    typename bp::to_python_indirect<value_type&,
+                                    bp::detail::make_reference_holder> conv;
+    return bp::object(bp::handle<>(conv(c[i])));
+  }
+
   static index_type convert_index(Container &container, PyObject *i_) {
     bp::extract<long> i(i_);
     if (i.check()) {

unittest/python/test_std_vector.py

@@ -107,3 +107,11 @@ def checkZero(v):
 checkAllValues(l6[1:], l6.tolist()[1:])
 checkAllValues(l6[:-1], l6.tolist()[:-1])
 checkAllValues(l6[::2], l6.tolist()[::2])
+L = [0, 2]
+L6_copy = l6[L]
+for k, i in enumerate(L):
+    checkAllValues(L6_copy[k], l6[i])
+T = (0, 2)
+L6_copy = l6[T]
+for k, i in enumerate(L):
+    checkAllValues(L6_copy[k], l6[i])