Browse Source

linmath: Implement read-only buffer protocol support for vectors

Fixes #1194
rdb 4 years ago
parent
commit
a4ea476cce

+ 40 - 0
panda/src/linmath/lvecBase2_ext_src.I

@@ -357,5 +357,45 @@ __ceil__(PyObject *self) const {
   return py_vec;
 }
 
+/**
+ *
+ */
+INLINE_LINMATH int Extension<FLOATNAME(LVecBase2)>::
+__getbuffer__(PyObject *self, Py_buffer *view, int flags) const {
+  if ((flags & PyBUF_WRITABLE) == PyBUF_WRITABLE) {
+    PyErr_SetString(PyExc_BufferError,
+                    "Object is not writable.");
+    return -1;
+  }
+
+  static const char format[2] = {FLOATTOKEN, 0};
+  static const Py_ssize_t shape = FLOATNAME(LVecBase2)::num_components;
+
+  Py_INCREF(self);
+
+  view->buf = (void *)_this->get_data();
+  view->obj = self;
+  view->len = 2 * sizeof(FLOATTYPE);
+  view->readonly = 1;
+  view->itemsize = sizeof(FLOATTYPE);
+  view->format = nullptr;
+  if ((flags & PyBUF_FORMAT) == PyBUF_FORMAT) {
+    view->format = (char *)format;
+  }
+  view->ndim = 1;
+  view->shape = nullptr;
+  if ((flags & PyBUF_ND) == PyBUF_ND) {
+    view->shape = (Py_ssize_t *)&shape;
+  }
+  view->strides = nullptr;
+  if ((flags & PyBUF_STRIDES) == PyBUF_STRIDES) {
+    view->strides = &view->itemsize;
+  }
+  view->suboffsets = nullptr;
+  view->internal = nullptr;
+
+  return 0;
+}
+
 #undef PYNUMBER_FLOATTYPE
 #undef PY_AS_FLOATTYPE

+ 2 - 0
panda/src/linmath/lvecBase2_ext_src.h

@@ -34,6 +34,8 @@ public:
   INLINE_LINMATH PyObject *__round__(PyObject *self) const;
   INLINE_LINMATH PyObject *__floor__(PyObject *self) const;
   INLINE_LINMATH PyObject *__ceil__(PyObject *self) const;
+
+  INLINE_LINMATH int __getbuffer__(PyObject *self, Py_buffer *view, int flags) const;
 };
 
 #include "lvecBase2_ext_src.I"

+ 2 - 0
panda/src/linmath/lvecBase2_src.h

@@ -160,6 +160,8 @@ PUBLISHED:
   INLINE_LINMATH void write_datagram(Datagram &destination) const;
   INLINE_LINMATH void read_datagram(DatagramIterator &source);
 
+  EXTENSION(INLINE_LINMATH int __getbuffer__(PyObject *self, Py_buffer *view, int flags) const);
+
 public:
   // The underlying implementation is via the Eigen library, if available.
 

+ 40 - 0
panda/src/linmath/lvecBase3_ext_src.I

@@ -370,5 +370,45 @@ __ceil__(PyObject *self) const {
   return py_vec;
 }
 
+/**
+ *
+ */
+INLINE_LINMATH int Extension<FLOATNAME(LVecBase3)>::
+__getbuffer__(PyObject *self, Py_buffer *view, int flags) const {
+  if ((flags & PyBUF_WRITABLE) == PyBUF_WRITABLE) {
+    PyErr_SetString(PyExc_BufferError,
+                    "Object is not writable.");
+    return -1;
+  }
+
+  static const char format[2] = {FLOATTOKEN, 0};
+  static const Py_ssize_t shape = FLOATNAME(LVecBase3)::num_components;
+
+  Py_INCREF(self);
+
+  view->buf = (void *)_this->get_data();
+  view->obj = self;
+  view->len = 3 * sizeof(FLOATTYPE);
+  view->readonly = 1;
+  view->itemsize = sizeof(FLOATTYPE);
+  view->format = nullptr;
+  if ((flags & PyBUF_FORMAT) == PyBUF_FORMAT) {
+    view->format = (char *)format;
+  }
+  view->ndim = 1;
+  view->shape = nullptr;
+  if ((flags & PyBUF_ND) == PyBUF_ND) {
+    view->shape = (Py_ssize_t *)&shape;
+  }
+  view->strides = nullptr;
+  if ((flags & PyBUF_STRIDES) == PyBUF_STRIDES) {
+    view->strides = &view->itemsize;
+  }
+  view->suboffsets = nullptr;
+  view->internal = nullptr;
+
+  return 0;
+}
+
 #undef PYNUMBER_FLOATTYPE
 #undef PY_AS_FLOATTYPE

+ 2 - 0
panda/src/linmath/lvecBase3_ext_src.h

@@ -34,6 +34,8 @@ public:
   INLINE_LINMATH PyObject *__round__(PyObject *self) const;
   INLINE_LINMATH PyObject *__floor__(PyObject *self) const;
   INLINE_LINMATH PyObject *__ceil__(PyObject *self) const;
+
+  INLINE_LINMATH int __getbuffer__(PyObject *self, Py_buffer *view, int flags) const;
 };
 
 #include "lvecBase3_ext_src.I"

+ 2 - 0
panda/src/linmath/lvecBase3_src.h

@@ -181,6 +181,8 @@ PUBLISHED:
   INLINE_LINMATH void write_datagram(Datagram &destination) const;
   INLINE_LINMATH void read_datagram(DatagramIterator &source);
 
+  EXTENSION(INLINE_LINMATH int __getbuffer__(PyObject *self, Py_buffer *view, int flags) const);
+
 public:
   // The underlying implementation is via the Eigen library, if available.
 

+ 40 - 0
panda/src/linmath/lvecBase4_ext_src.I

@@ -388,5 +388,45 @@ __ceil__(PyObject *self) const {
   return py_vec;
 }
 
+/**
+ *
+ */
+INLINE_LINMATH int Extension<FLOATNAME(LVecBase4)>::
+__getbuffer__(PyObject *self, Py_buffer *view, int flags) const {
+  if ((flags & PyBUF_WRITABLE) == PyBUF_WRITABLE) {
+    PyErr_SetString(PyExc_BufferError,
+                    "Object is not writable.");
+    return -1;
+  }
+
+  static const char format[2] = {FLOATTOKEN, 0};
+  static const Py_ssize_t shape = FLOATNAME(LVecBase4)::num_components;
+
+  Py_INCREF(self);
+
+  view->buf = (void *)_this->get_data();
+  view->obj = self;
+  view->len = 4 * sizeof(FLOATTYPE);
+  view->readonly = 1;
+  view->itemsize = sizeof(FLOATTYPE);
+  view->format = nullptr;
+  if ((flags & PyBUF_FORMAT) == PyBUF_FORMAT) {
+    view->format = (char *)format;
+  }
+  view->ndim = 1;
+  view->shape = nullptr;
+  if ((flags & PyBUF_ND) == PyBUF_ND) {
+    view->shape = (Py_ssize_t *)&shape;
+  }
+  view->strides = nullptr;
+  if ((flags & PyBUF_STRIDES) == PyBUF_STRIDES) {
+    view->strides = &view->itemsize;
+  }
+  view->suboffsets = nullptr;
+  view->internal = nullptr;
+
+  return 0;
+}
+
 #undef PYNUMBER_FLOATTYPE
 #undef PY_AS_FLOATTYPE

+ 2 - 0
panda/src/linmath/lvecBase4_ext_src.h

@@ -34,6 +34,8 @@ public:
   INLINE_LINMATH PyObject *__round__(PyObject *self) const;
   INLINE_LINMATH PyObject *__floor__(PyObject *self) const;
   INLINE_LINMATH PyObject *__ceil__(PyObject *self) const;
+
+  INLINE_LINMATH int __getbuffer__(PyObject *self, Py_buffer *view, int flags) const;
 };
 
 #include "lvecBase4_ext_src.I"

+ 2 - 0
panda/src/linmath/lvecBase4_src.h

@@ -187,6 +187,8 @@ PUBLISHED:
   INLINE_LINMATH void write_datagram(Datagram &destination) const;
   INLINE_LINMATH void read_datagram(DatagramIterator &source);
 
+  EXTENSION(INLINE_LINMATH int __getbuffer__(PyObject *self, Py_buffer *view, int flags) const);
+
 public:
   // The underlying implementation is via the Eigen library, if available.
 

+ 8 - 0
tests/linmath/test_lvector2.py

@@ -143,3 +143,11 @@ def test_vec2_floordiv(type):
             v = type(i)
             v //= -j
             assert v.x == i // -j
+
+
+def test_vec2_buffer():
+    v = Vec2(1.5, -10.0)
+    m = memoryview(v)
+    assert len(m) == 2
+    assert m[0] == 1.5
+    assert m[1] == -10.0

+ 9 - 0
tests/linmath/test_lvector3.py

@@ -128,3 +128,12 @@ def test_vec3_floordiv(type):
             v = type(i)
             v //= -j
             assert v.x == i // -j
+
+
+def test_vec3_buffer():
+    v = Vec3(0.5, 2.0, -10.0)
+    m = memoryview(v)
+    assert len(m) == 3
+    assert m[0] == 0.5
+    assert m[1] == 2.0
+    assert m[2] == -10.0

+ 10 - 0
tests/linmath/test_lvector4.py

@@ -144,3 +144,13 @@ def test_vec4_floordiv(type):
             v = type(i)
             v //= -j
             assert v.x == i // -j
+
+
+def test_vec4_buffer():
+    v = Vec4(0, 0.5, 2.0, -4.0)
+    m = memoryview(v)
+    assert len(m) == 4
+    assert m[0] == 0
+    assert m[1] == 0.5
+    assert m[2] == 2.0
+    assert m[3] == -4.0