Add the buffer interface for wrapped STL vectors

Allows use of vectors as python buffers, so for example they can be adopted without a copy by numpy.asarray
Allows faster conversion of numeric buffers to vectors with memcpy instead of individually casting the elements

Author: patstew

include/pybind11/stl_bind.h

@@ -64,6 +64,7 @@ template <typename, typename, typename... Args> void vector_if_copy_constructibl
 template <typename, typename, typename... Args> void vector_if_equal_operator(const Args &...) { }
 template <typename, typename, typename... Args> void vector_if_insertion_operator(const Args &...) { }
 template <typename, typename, typename... Args> void vector_modifiers(const Args &...) { }
+template <typename, typename, typename... Args> void vector_buffer(const Args&...) { }
 
 template<typename Vector, typename Class_>
 void vector_if_copy_constructible(enable_if_t<
@@ -326,6 +327,36 @@ template <typename Vector, typename Class_> auto vector_if_insertion_operator(Cl
     );
 }
 
+// Provide the buffer interface for vectors if we have data() and we have a format for it
+// GCC seems to have "void std::vector<bool>::data()" - doing SFINAE on the existence of data() is insufficient, we need to check it doesn't return void
+template <typename Vector, typename Class_>
+enable_if_t<!std::is_same<decltype(py::format_descriptor<typename Vector::value_type>::format(), std::declval<Vector>().data()), void>::value>
+vector_buffer(Class_& cl) {
+    using T = typename Vector::value_type;
+
+    try {
+        //numpy.h declares this for arbitrary types, but it may raise an exception if PYBIND11_NUMPY_DTYPE hasn't been called
+        py::format_descriptor<T>::format();
+    } catch (std::runtime_error&) {
+        return;
+    }
+
+    cl.def_buffer([](Vector& v) -> py::buffer_info {
+        return py::buffer_info(v.data(), sizeof(T), py::format_descriptor<T>::format(), 1, {v.size()}, {sizeof(T)});
+    });
+
+    cl.def("__init__", [](Vector& vec, py::buffer buf) {
+        auto info = buf.request();
+        if (info.ndim != 1)
+            throw pybind11::type_error("Only 1D buffers can be copied to a vector");
+        if (info.strides[0] != sizeof(T))
+            throw pybind11::type_error("Item size mismatch (Python: " + std::to_string(info.strides[0]) + " C++: " + std::to_string(sizeof(T)) + ")");
+        if (info.format != py::format_descriptor<T>::format())
+            throw pybind11::type_error("Format mismatch (Python: " + info.format + " C++: " + py::format_descriptor<T>::format() + ")");
+        new (&vec) Vector(static_cast<T*>(info.ptr), static_cast<T*>(info.ptr) + info.shape[0]);
+    });
+}
+
 NAMESPACE_END(detail)
 
 //
@@ -348,6 +379,9 @@ pybind11::class_<Vector, holder_type> bind_vector(pybind11::module &m, std::stri
     // Register stream insertion operator (if possible)
     detail::vector_if_insertion_operator<Vector, Class_>(cl, name);
 
+    // Register the buffer interface (if possible)
+    detail::vector_buffer<Vector, Class_>(cl);
+
     // Modifiers require copyable vector value type
     detail::vector_modifiers<Vector, Class_>(cl);
 

tests/test_stl_binders.cpp

@@ -10,6 +10,7 @@
 #include "pybind11_tests.h"
 
 #include <pybind11/stl_bind.h>
+#include <pybind11/numpy.h>
 #include <map>
 #include <deque>
 #include <unordered_map>
@@ -53,17 +54,44 @@ template <class Map> Map *times_ten(int n) {
     return m;
 }
 
+struct VStruct {
+    bool w;
+    uint32_t x;
+    double y;
+    bool z;
+};
+
+struct VUndeclStruct { //dtype not declared for this version
+    bool w;
+    uint32_t x;
+    double y;
+    bool z;
+};
+
 test_initializer stl_binder_vector([](py::module &m) {
     py::class_<El>(m, "El")
         .def(py::init<int>());
 
+    py::bind_vector<std::vector<unsigned char>>(m, "VectorUChar");
     py::bind_vector<std::vector<unsigned int>>(m, "VectorInt");
     py::bind_vector<std::vector<bool>>(m, "VectorBool");
 
     py::bind_vector<std::vector<El>>(m, "VectorEl");
 
     py::bind_vector<std::vector<std::vector<El>>>(m, "VectorVectorEl");
 
+    py::bind_vector<std::vector<VUndeclStruct>>(m, "VectorUndeclStruct");
+    m.def("get_undeclstruct", [] {return std::vector<VUndeclStruct> {{0, 5, 3.0, 1}, {1, 30, -1e4, 0}};});
+
+    try {
+        py::module::import("numpy");
+    } catch (...) {
+        return;
+    }
+    PYBIND11_NUMPY_DTYPE(VStruct, w, x, y, z);
+    py::class_<VStruct>(m, "VStruct").def_readwrite("x", &VStruct::x);
+    py::bind_vector<std::vector<VStruct>>(m, "VectorStruct");
+    m.def("get_vectorstruct", [] {return std::vector<VStruct> {{0, 5, 3.0, 1}, {1, 30, -1e4, 0}};});
 });
 
 test_initializer stl_binder_map([](py::module &m) {
@@ -92,4 +120,3 @@ test_initializer stl_binder_noncopyable([](py::module &m) {
     py::bind_map<std::unordered_map<int, E_nc>>(m, "UmapENC");
     m.def("get_umnc", &times_ten<std::unordered_map<int, E_nc>>, py::return_value_policy::reference);
 });
-

tests/test_stl_binders.py

@@ -1,3 +1,9 @@
+import pytest
+import sys
+
+with pytest.suppress(ImportError):
+    import numpy as np
+
 def test_vector_int():
     from pybind11_tests import VectorInt
 
@@ -25,6 +31,45 @@ def test_vector_int():
     del v_int2[0]
     assert v_int2 == VectorInt([0, 99, 2, 3])
 
+def test_vector_buffer():
+    from pybind11_tests import VectorUChar, get_undeclstruct
+    b = bytearray([1,2,3,4])
+    v = VectorUChar(b);
+    assert v[1] == 2
+    v[2] = 5
+    m = memoryview(v) #We expose the buffer interface
+    if sys.version_info.major > 2:
+        assert m[2] == 5
+        m[2] = 6
+    else:
+        assert m[2] == '\x05'
+        m[2] = '\x06'
+    assert v[2] == 6
+
+    v = get_undeclstruct()
+    with pytest.raises(TypeError):
+        memoryview(v) #Undeclared struct contents, no buffer interface
+
+@pytest.requires_numpy
+def test_vector_buffer_numpy():
+    from pybind11_tests import VectorInt, VectorStruct, get_vectorstruct
+
+    a = np.array([1,2,3,4], dtype=np.int32)
+    with pytest.raises(TypeError):
+        VectorInt(a)
+
+    a = np.array([1,2,3,4], dtype=np.uintc)
+    v = VectorInt(a)
+    assert v[2] == 3
+    m = np.asarray(v)
+    m[2] = 5
+    assert v[2] == 5
+
+    v = get_vectorstruct()
+    assert v[0].x == 5
+    m = np.asarray(v)
+    m[1]['x'] = 99
+    assert v[1].x == 99
 
 def test_vector_custom():
     from pybind11_tests import El, VectorEl, VectorVectorEl
@@ -150,4 +195,3 @@ def test_noncopyable_unordered_map():
         vsum += v.value
 
     assert(vsum == 150)
-