buffer_protocol for vectors and matrices in python bindings (#524)

his change enables the buffer_protocol for `VectorX` and `MatrixYxY` so that `memoryiew` and `numpy.array` can be used on them

```
np.array(math7.Vector3())
memoryview(math7.Matrix3d())

list(math7.Vector2d())
>>> math7.Quaterniond().xyzw() # == [0.0, 0.0, 0.0, 1.0]
```
Additionally:
- IndexError is raised on invalid vector __getitem__/__setitem__ access. This way list(vector) works correctly instead of spinning
- A convenience Quaternion.xyzw() -> list is added

---------

Signed-off-by: Benoit Maurin <[email protected]>
Signed-off-by: Addisu Z. Taddese <[email protected]>
Co-authored-by: Addisu Z. Taddese <[email protected]>

.github/ci/packages.apt

@@ -2,6 +2,7 @@ libeigen3-dev
 libgz-cmake4-dev
 libgz-utils3-dev
 libpython3-dev
+python3-numpy
 python3-pybind11
 python3-pytest
 ruby-dev

include/gz/math/Matrix3.hh

@@ -576,6 +576,16 @@ namespace gz::math
       return _in;
     }
 
+    /// \brief Underlying data pointer
+    /// \remarks This method is intended for python bindings (numpy).
+    /// \remarks The bounds-checking array subscript operator is much preferred
+    /// for element access from C++.
+    /// \return A pointer to the underlying data array.
+    public: T* Data()
+    {
+      return this->data[0];
+    }
+
     /// \brief the 3x3 matrix
     private: T data[3][3];
   };

include/gz/math/Matrix4.hh

@@ -829,6 +829,16 @@ namespace gz::math
                 0,      0,         0,        1);
     }
 
+    /// \brief Underlying data pointer
+    /// \remarks This method is intended for python bindings (numpy).
+    /// \remarks The bounds-checking array subscript operator is much preferred
+    /// for element access from C++.
+    /// \return A pointer to the underlying data array.
+    public: T* Data()
+    {
+      return this->data[0];
+    }
+
     /// \brief The 4x4 matrix
     private: T data[4][4];
   };

include/gz/math/Matrix6.hh

@@ -541,6 +541,16 @@ namespace gz::math
       return _in;
     }
 
+    /// \brief Underlying data pointer
+    /// \remarks This method is intended for python bindings (numpy).
+    /// \remarks The bounds-checking array subscript operator is much preferred
+    /// for element access from C++.
+    /// \return A pointer to the underlying data array.
+    public: T* Data()
+    {
+      return this->data[0];
+    }
+
     /// \brief The 6x6 matrix
     private: T data[MatrixSize][MatrixSize];
   };

include/gz/math/Vector2.hh

@@ -556,6 +556,16 @@ namespace gz::math
       return _in;
     }
 
+    /// \brief Underlying data pointer
+    /// \remarks This method is intended for python bindings (numpy).
+    /// \remarks The bounds-checking array subscript operator is much preferred
+    /// for element access from C++.
+    /// \return A pointer to the underlying data array.
+    public: T* Data()
+    {
+      return this->data;
+    }
+
     /// \brief The x and y values.
     private: T data[2];
   };

include/gz/math/Vector3.hh

@@ -746,6 +746,16 @@ namespace gz::math
       return _in;
     }
 
+    /// \brief Underlying data pointer
+    /// \remarks This method is intended for python bindings (numpy).
+    /// \remarks The bounds-checking array subscript operator is much preferred
+    /// for element access from C++.
+    /// \return A pointer to the underlying data array.
+    public: T* Data()
+    {
+      return this->data;
+    }
+
     /// \brief The x, y, and z values
     private: T data[3];
   };

include/gz/math/Vector4.hh

@@ -706,6 +706,16 @@ namespace gz::math
       return _in;
     }
 
+    /// \brief Underlying data pointer
+    /// \remarks This method is intended for python bindings (numpy).
+    /// \remarks The bounds-checking array subscript operator is much preferred
+    /// for element access from C++.
+    /// \return A pointer to the underlying data array.
+    public: T* Data()
+    {
+      return this->data;
+    }
+
     /// \brief Data values, 0==x, 1==y, 2==z, 3==w
     private: T data[4];
   };

src/Matrix3_TEST.cc

@@ -417,3 +417,16 @@ TEST(Matrix3dTest, From2Axes)
   m1.SetFrom2Axes(v1, v2);
   EXPECT_EQ(math::Matrix3d::Zero - math::Matrix3d::Identity, m1);
 }
+
+/////////////////////////////////////////////////
+TEST(Matrix3dTest, Data)
+{
+  math::Matrix3d m1(0, 1, 2,
+                    3, 4, 5,
+                    6, 7, 8);
+
+  for (int i = 0; i < 9; ++i) EXPECT_EQ(m1.Data()[i], i);
+  math::Matrix3d m2 = math::Matrix3d::Zero;
+  for (int i = 0; i < 9; ++i) m2.Data()[i] = i;
+  EXPECT_EQ(m1, m2);
+}

src/Matrix4_TEST.cc

@@ -741,3 +741,17 @@ TEST(Matrix4dTest, LookAt)
                 .Pose(),
             math::Pose3d(1, 1, 1, GZ_PI_4, 0, GZ_PI));
 }
+
+/////////////////////////////////////////////////
+TEST(Matrix4dTest, Data)
+{
+  math::Matrix4d m1(0, 1, 2, 3,
+                    4, 5, 6, 7,
+                    8, 9, 10, 11,
+                    12, 13, 14, 15);
+
+  for (int i = 0; i < 16; ++i) EXPECT_EQ(m1.Data()[i], i);
+  math::Matrix4d m2 = math::Matrix4d::Zero;
+  for (int i = 0; i < 16; ++i) m2.Data()[i] = i;
+  EXPECT_EQ(m1, m2);
+}

src/Matrix6_TEST.cc

@@ -393,3 +393,19 @@ TEST(Matrix6dTest, SetValue)
       4, 3, 2, 1, 0, -1,
       5, 4, 3, 2, 1, 0));
 }
+
+/////////////////////////////////////////////////
+TEST(Matrix6dTest, Data)
+{
+  math::Matrix6d m1(0, 1, 2, 3, 4, 5,
+                    6, 7, 8, 9, 10, 11,
+                    12, 13, 14, 15, 16, 17,
+                    18, 19, 20, 21, 22, 23,
+                    24, 25, 26, 27, 28, 29,
+                    30, 31, 32, 33, 34, 35);
+
+  for (int i = 0; i < 36; ++i) EXPECT_EQ(m1.Data()[i], i);
+  math::Matrix6d m2 = math::Matrix6d::Zero;
+  for (int i = 0; i < 36; ++i) m2.Data()[i] = i;
+  EXPECT_EQ(m1, m2);
+}

src/Vector2_TEST.cc

@@ -476,3 +476,15 @@ TEST(Vector2Test, NaN)
   EXPECT_EQ(math::Vector2f::Zero, nanVecF);
   EXPECT_TRUE(nanVecF.IsFinite());
 }
+
+/////////////////////////////////////////////////
+TEST(Vector2Test, Data)
+{
+  math::Vector2d v(0, 1);
+  for (int i = 0; i < 2; ++i) EXPECT_EQ(v.Data()[i], v[i]);
+
+  auto v2 = math::Vector2d::Zero;
+  for (int i = 0; i < 2; ++i) v2.Data()[i] = i;
+
+  EXPECT_EQ(v, v2);
+}

src/Vector3_TEST.cc

@@ -576,3 +576,15 @@ TEST(Vector3dTest, OperatorStreamOut)
   stream << std::setprecision(1) << std::fixed << v;
   EXPECT_EQ(stream.str(), "0.1 1.2 2.3");
 }
+
+/////////////////////////////////////////////////
+TEST(Vector3Test, Data)
+{
+  math::Vector3d v(0, 1, 2);
+  for (int i = 0; i < 3; ++i) EXPECT_EQ(v.Data()[i], v[i]);
+
+  auto v2 = math::Vector3d::Zero;
+  for (int i = 0; i < 3; ++i) v2.Data()[i] = i;
+
+  EXPECT_EQ(v, v2);
+}

src/Vector4_TEST.cc

@@ -480,3 +480,15 @@ TEST(Vector4Test, NaN)
   EXPECT_EQ(math::Vector4f::Zero, nanVecF);
   EXPECT_TRUE(nanVecF.IsFinite());
 }
+
+/////////////////////////////////////////////////
+TEST(Vector4Test, Data)
+{
+  math::Vector4d v(0, 1, 2, 3);
+  for (int i = 0; i < 4; ++i) EXPECT_EQ(v.Data()[i], v[i]);
+
+  auto v2 = math::Vector4d::Zero;
+  for (int i = 0; i < 4; ++i) v2.Data()[i] = i;
+
+  EXPECT_EQ(v, v2);
+}

src/python_pybind11/src/Matrix3.hh

@@ -51,6 +51,7 @@ template<typename T>
 void helpDefineMathMatrix3(py::module &m, const std::string &typestr)
 {
   using Class = gz::math::Matrix3<T>;
+  static constexpr int mat_size = 3;
   auto toString = [](const Class &si) {
     std::stringstream stream;
     stream << si;
@@ -122,6 +123,16 @@ void helpDefineMathMatrix3(py::module &m, const std::string &typestr)
     }, "memo"_a)
     .def_readonly_static("IDENTITY", &Class::Identity, "Identity matrix")
     .def_readonly_static("ZERO", &Class::Zero, "Zero matrix")
+    .def_buffer([](Class &self) -> py::buffer_info {
+        return py::buffer_info(
+            self.Data(),
+            sizeof(T),
+            py::format_descriptor<T>::format(),
+            2,
+            { mat_size, mat_size },
+            { mat_size * sizeof(T), sizeof(T) }
+        );
+    })
     .def("__str__", toString)
     .def("__repr__", toString);
 }

src/python_pybind11/src/Matrix4.hh

@@ -51,6 +51,7 @@ template<typename T>
 void helpDefineMathMatrix4(py::module &m, const std::string &typestr)
 {
   using Class = gz::math::Matrix4<T>;
+  static constexpr int mat_size = 4;
   auto toString = [](const Class &si) {
     std::stringstream stream;
     stream << si;
@@ -144,6 +145,16 @@ void helpDefineMathMatrix4(py::module &m, const std::string &typestr)
     }, "memo"_a)
     .def_readonly_static("IDENTITY", &Class::Identity, "Identity matrix")
     .def_readonly_static("ZERO", &Class::Zero, "Zero matrix")
+    .def_buffer([](Class &self) -> py::buffer_info {
+        return py::buffer_info(
+            self.Data(),
+            sizeof(T),
+            py::format_descriptor<T>::format(),
+            2,
+            { mat_size, mat_size },
+            { mat_size * sizeof(T), sizeof(T) }
+        );
+    })
     .def("__str__", toString)
     .def("__repr__", toString);
 }

src/python_pybind11/src/Matrix6.hh

@@ -51,6 +51,7 @@ template<typename T>
 void helpDefineMathMatrix6(py::module &_m, const std::string &_typestr)
 {
   using Class = Matrix6<T>;
+  static const int mat_size = 6;
   auto toString = [](const Class &_si) {
     std::stringstream stream;
     stream << _si;
@@ -98,6 +99,16 @@ void helpDefineMathMatrix6(py::module &_m, const std::string &_typestr)
     }, "memo"_a)
     .def_readonly_static("IDENTITY", &Class::Identity, "Identity matrix")
     .def_readonly_static("ZERO", &Class::Zero, "Zero matrix")
+    .def_buffer([](Class &self) -> py::buffer_info {
+        return py::buffer_info(
+            self.Data(),
+            sizeof(T),
+            py::format_descriptor<T>::format(),
+            2,
+            { mat_size, mat_size },
+            { mat_size * sizeof(T), sizeof(T) }
+        );
+    })
     .def("__str__", toString)
     .def("__repr__", toString);
 }

src/python_pybind11/src/Quaternion.hh

@@ -20,6 +20,7 @@
 
 #include <sstream>
 #include <string>
+#include <vector>
 
 #include <pybind11/pybind11.h>
 #include <pybind11/operators.h>
@@ -62,7 +63,6 @@ void helpDefineMathQuaternion(py::module &m, const std::string &typestr)
   std::string pyclass_name = typestr;
   py::class_<Class>(m,
                     pyclass_name.c_str(),
-                    py::buffer_protocol(),
                     py::dynamic_attr())
     .def(py::init<>())
     .def(py::init<T, T, T, T>())
@@ -210,6 +210,9 @@ void helpDefineMathQuaternion(py::module &m, const std::string &typestr)
     }, "memo"_a)
     .def_readonly_static("IDENTITY", &Class::Identity, "Identity matrix")
     .def_readonly_static("ZERO", &Class::Zero, "Zero matrix")
+    .def("xyzw", [](const Class &self){
+        return std::vector<T> { self.X(), self.Y(), self.Z(), self.W() };
+      }, "Number of elements in the vector")
     .def("__str__", toString)
     .def("__repr__", toString);
 }

src/python_pybind11/src/Vector2.hh

@@ -43,6 +43,7 @@ template<typename T>
 void helpDefineMathVector2(py::module &m, const std::string &typestr)
 {
   using Class = gz::math::Vector2<T>;
+  static constexpr int vec_size = 2;
   auto toString = [](const Class &si) {
     std::stringstream stream;
     stream << si;
@@ -132,10 +133,24 @@ void helpDefineMathVector2(py::module &m, const std::string &typestr)
     .def("__deepcopy__", [](const Class &self, py::dict) {
       return Class(self);
     }, "memo"_a)
-    .def("__getitem__",
-         py::overload_cast<const std::size_t>(&Class::operator[], py::const_))
-    .def("__setitem__",
-         [](Class* vec, unsigned index, T val) { (*vec)[index] = val; })
+    .def("__getitem__", [](const Class &self, unsigned index) {
+        if (index >= vec_size) throw py::index_error("Invalid index");
+        return self[index];
+      })
+    .def("__setitem__", [](Class* vec, unsigned index, T val) {
+        if (index >= vec_size) throw py::index_error("Invalid index");
+        (*vec)[index] = val;
+      })
+    .def_buffer([](Class &self) -> py::buffer_info {
+        return py::buffer_info(
+            self.Data(),
+            sizeof(T),
+            py::format_descriptor<T>::format(),
+            1,
+            { vec_size },
+            { sizeof(T) }
+        );
+    })
     .def("__str__", toString)
     .def("__repr__", toString);
 }

src/python_pybind11/src/Vector3.hh

@@ -43,6 +43,7 @@ template<typename T>
 void helpDefineMathVector3(py::module &m, const std::string &typestr)
 {
   using Class = gz::math::Vector3<T>;
+  static constexpr int vec_size = 3;
   auto toString = [](const Class &si) {
     std::stringstream stream;
     stream << si;
@@ -159,10 +160,24 @@ void helpDefineMathVector3(py::module &m, const std::string &typestr)
     .def("__deepcopy__", [](const Class &self, py::dict) {
       return Class(self);
     }, "memo"_a)
-    .def("__getitem__",
-         py::overload_cast<const std::size_t>(&Class::operator[], py::const_))
-    .def("__setitem__",
-         [](Class* vec, unsigned index, T val) { (*vec)[index] = val; })
+    .def("__getitem__", [](const Class &self, unsigned index) {
+        if (index >= vec_size) throw py::index_error("Invalid index");
+        return self[index];
+      })
+    .def("__setitem__", [](Class* vec, unsigned index, T val) {
+        if (index >= vec_size) throw py::index_error("Invalid index");
+        (*vec)[index] = val;
+      })
+    .def_buffer([](Class &self) -> py::buffer_info {
+        return py::buffer_info(
+            self.Data(),
+            sizeof(T),
+            py::format_descriptor<T>::format(),
+            1,
+            { vec_size },
+            { sizeof(T) }
+        );
+    })
     .def("__str__", toString)
     .def("__repr__", toString);
 }