PyAPI: buffer protocol support for mathutils types

Adding buffer protocol support increases the speed of copying a Vector
(3D) array into a `numpy.array` by up to x3.8.

Ref !144401
This commit is contained in:
Oxicid
2025-08-16 06:14:19 +00:00
committed by Campbell Barton
parent ebfa7edeb1
commit b856b6010e
8 changed files with 402 additions and 6 deletions

View File

@@ -622,6 +622,12 @@ int _BaseMathObject_ResizeOkOrRaiseExc(BaseMathObject *self, const char *error_p
PyErr_Format(PyExc_ValueError, "%s: cannot resize wrapped data", error_prefix);
return -1;
}
if (UNLIKELY(self->flag & BASE_MATH_FLAG_HAS_BUFFER_VIEW)) {
PyErr_Format(PyExc_BufferError,
"%s: cannot resize data while exported to buffer protocol",
error_prefix);
return -1;
}
if (UNLIKELY(self->cb_user)) {
PyErr_Format(PyExc_ValueError, "%s: cannot resize owned data", error_prefix);
return -1;
@@ -629,6 +635,30 @@ int _BaseMathObject_ResizeOkOrRaiseExc(BaseMathObject *self, const char *error_p
return 0;
}
int _BaseMathObject_RaiseBufferViewExc(BaseMathObject *self, Py_buffer *view, int flags)
{
if (UNLIKELY(view == nullptr)) {
PyErr_SetString(PyExc_BufferError, "null view in getbuffer is obsolete");
return -1;
}
if (UNLIKELY(self->flag & BASE_MATH_FLAG_HAS_BUFFER_VIEW)) {
PyErr_SetString(PyExc_BufferError,
"Data is already exported via buffer protocol, "
"multiple simultaneous exports are not allowed.");
return -1;
}
if (flags & PyBUF_WRITABLE) {
if (UNLIKELY(BaseMath_WriteCallback(self) == -1)) {
return -1;
}
if (UNLIKELY(self->flag & BASE_MATH_FLAG_IS_FROZEN)) {
PyErr_Format(PyExc_BufferError, "Data is frozen, cannot get a writable buffer");
return -1;
}
}
return 0;
}
/* #BaseMathObject generic functions for all mathutils types. */
char BaseMathObject_owner_doc[] = "The item this is wrapping or None (read-only).";
@@ -673,6 +703,11 @@ PyObject *BaseMathObject_freeze(BaseMathObject *self)
return nullptr;
}
if (self->flag & BASE_MATH_FLAG_HAS_BUFFER_VIEW) {
PyErr_SetString(PyExc_BufferError, "Cannot freeze data while exported to buffer protocol");
return nullptr;
}
self->flag |= BASE_MATH_FLAG_IS_FROZEN;
return Py_NewRef(self);

View File

@@ -40,6 +40,14 @@ enum {
* (typical use cases for tuple).
*/
BASE_MATH_FLAG_IS_FROZEN = (1 << 1),
/**
* When set, prevents calling freeze() and resize() while using the buffer protocol.
*
* \note `memoryview` & `np.frombuffer` pass the `PyBUF_FORMAT | PyBUF_INDIRECT` flags,
* and the object can be mutated, so `PyBUF_WRITABLE` can't be handled.
* That's why it's always necessary to check for write access.
*/
BASE_MATH_FLAG_HAS_BUFFER_VIEW = (1 << 2),
};
#define BASE_MATH_FLAG_DEFAULT 0
@@ -121,6 +129,9 @@ struct Mathutils_Callback {
/** To implement #BaseMath_Prepare_ForResize. */
[[nodiscard]] int _BaseMathObject_ResizeOkOrRaiseExc(BaseMathObject *self,
const char *error_prefix);
[[nodiscard]] int _BaseMathObject_RaiseBufferViewExc(BaseMathObject *self,
Py_buffer *view,
int flags);
void _BaseMathObject_RaiseFrozenExc(const BaseMathObject *self);
void _BaseMathObject_RaiseNotFrozenExc(const BaseMathObject *self);
@@ -164,6 +175,15 @@ void _BaseMathObject_RaiseNotFrozenExc(const BaseMathObject *self);
#define BaseMathObject_Prepare_ForResize(_self, error_prefix) \
_BaseMathObject_ResizeOkOrRaiseExc((BaseMathObject *)_self, error_prefix)
/**
* Ensure #BASE_MATH_FLAG_HAS_BUFFER_VIEW is supported.
* \param _view: The `view` argument forwarded from #PyBufferProcs::bf_getbuffer.
* \param _flags: The `flags` argument forwarded from #PyBufferProcs::bf_getbuffer.
* \return -1 and set an exception if the vector `_self` does not support buffer access.
*/
#define BaseMath_Prepare_ForBufferAccess(_self, _view, _flags) \
_BaseMathObject_RaiseBufferViewExc((BaseMathObject *)_self, _view, _flags)
/* utility func */
/**
* Helper function.

View File

@@ -305,6 +305,59 @@ static PyObject *Color_str(ColorObject *self)
/** \} */
/* -------------------------------------------------------------------- */
/** \name Color Type: Buffer Protocol
* \{ */
static int Color_getbuffer(PyObject *obj, Py_buffer *view, int flags)
{
ColorObject *self = (ColorObject *)obj;
if (UNLIKELY(BaseMath_Prepare_ForBufferAccess(self, view, flags) == -1)) {
return -1;
}
if (UNLIKELY(BaseMath_ReadCallback(self) == -1)) {
return -1;
}
memset(view, 0, sizeof(*view));
view->obj = (PyObject *)self;
view->buf = (void *)self->col;
view->len = Py_ssize_t(COLOR_SIZE * sizeof(float));
view->itemsize = sizeof(float);
view->ndim = 1;
if ((flags & PyBUF_WRITABLE) == 0) {
view->readonly = 1;
}
if (flags & PyBUF_FORMAT) {
view->format = (char *)"f";
}
self->flag |= BASE_MATH_FLAG_HAS_BUFFER_VIEW;
Py_INCREF(self);
return 0;
}
static void Color_releasebuffer(PyObject * /*exporter*/, Py_buffer *view)
{
ColorObject *self = (ColorObject *)view->obj;
self->flag &= ~BASE_MATH_FLAG_HAS_BUFFER_VIEW;
if (view->readonly == 0) {
if (UNLIKELY(BaseMath_WriteCallback(self) == -1)) {
PyErr_Print();
}
}
}
static PyBufferProcs Color_as_buffer = {
(getbufferproc)Color_getbuffer,
(releasebufferproc)Color_releasebuffer,
};
/** \} */
/* -------------------------------------------------------------------- */
/** \name Color Type: Rich Compare
* \{ */
@@ -1230,7 +1283,7 @@ PyTypeObject color_Type = {
/*tp_str*/ (reprfunc)Color_str,
/*tp_getattro*/ nullptr,
/*tp_setattro*/ nullptr,
/*tp_as_buffer*/ nullptr,
/*tp_as_buffer*/ &Color_as_buffer,
/*tp_flags*/ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
/*tp_doc*/ color_doc,
/*tp_traverse*/ (traverseproc)BaseMathObject_traverse,

View File

@@ -383,6 +383,59 @@ static PyObject *Euler_str(EulerObject *self)
/** \} */
/* -------------------------------------------------------------------- */
/** \name Euler Type: Buffer Protocol
* \{ */
static int Euler_getbuffer(PyObject *obj, Py_buffer *view, int flags)
{
EulerObject *self = (EulerObject *)obj;
if (UNLIKELY(BaseMath_Prepare_ForBufferAccess(self, view, flags) == -1)) {
return -1;
}
if (UNLIKELY(BaseMath_ReadCallback(self) == -1)) {
return -1;
}
memset(view, 0, sizeof(*view));
view->obj = (PyObject *)self;
view->buf = (void *)self->eul;
view->len = Py_ssize_t(EULER_SIZE * sizeof(float));
view->itemsize = sizeof(float);
view->ndim = 1;
if ((flags & PyBUF_WRITABLE) == 0) {
view->readonly = 1;
}
if (flags & PyBUF_FORMAT) {
view->format = (char *)"f";
}
self->flag |= BASE_MATH_FLAG_HAS_BUFFER_VIEW;
Py_INCREF(self);
return 0;
}
static void Euler_releasebuffer(PyObject * /*exporter*/, Py_buffer *view)
{
EulerObject *self = (EulerObject *)view->obj;
self->flag &= ~BASE_MATH_FLAG_HAS_BUFFER_VIEW;
if (view->readonly == 0) {
if (UNLIKELY(BaseMath_WriteCallback(self) == -1)) {
PyErr_Print();
}
}
}
static PyBufferProcs Euler_as_buffer = {
(getbufferproc)Euler_getbuffer,
(releasebufferproc)Euler_releasebuffer,
};
/** \} */
/* -------------------------------------------------------------------- */
/** \name Euler Type: Rich Compare
* \{ */
@@ -854,7 +907,7 @@ PyTypeObject euler_Type = {
/*tp_str*/ (reprfunc)Euler_str,
/*tp_getattro*/ nullptr,
/*tp_setattro*/ nullptr,
/*tp_as_buffer*/ nullptr,
/*tp_as_buffer*/ &Euler_as_buffer,
/*tp_flags*/ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
/*tp_doc*/ euler_doc,
/*tp_traverse*/ (traverseproc)BaseMathObject_traverse,

View File

@@ -2370,6 +2370,71 @@ static PyObject *Matrix_str(MatrixObject *self)
/** \} */
/* -------------------------------------------------------------------- */
/** \name Matrix Type: Buffer Protocol
* \{ */
static int Matrix_getbuffer(PyObject *obj, Py_buffer *view, int flags)
{
MatrixObject *self = (MatrixObject *)obj;
if (UNLIKELY(BaseMath_Prepare_ForBufferAccess(self, view, flags) == -1)) {
return -1;
}
if (UNLIKELY(BaseMath_ReadCallback(self) == -1)) {
return -1;
}
memset(view, 0, sizeof(*view));
view->obj = (PyObject *)self;
view->buf = (void *)self->matrix;
view->len = Py_ssize_t(self->row_num * self->col_num * sizeof(float));
view->itemsize = sizeof(float);
if ((flags & PyBUF_WRITABLE) == 0) {
view->readonly = 1;
}
if (flags & PyBUF_FORMAT) {
view->format = (char *)"f";
}
if (flags & PyBUF_ND) {
view->ndim = 2;
view->shape = MEM_malloc_arrayN<Py_ssize_t>(size_t(view->ndim), __func__);
view->shape[0] = self->row_num;
view->shape[1] = self->col_num;
}
if (flags & PyBUF_STRIDES) {
view->strides = MEM_malloc_arrayN<Py_ssize_t>(size_t(view->ndim), __func__);
view->strides[0] = sizeof(float); /* step between lines in column-major */
view->strides[1] = Py_ssize_t(self->row_num) * sizeof(float); /* step between columns */
}
self->flag |= BASE_MATH_FLAG_HAS_BUFFER_VIEW;
Py_INCREF(self);
return 0;
}
static void Matrix_releasebuffer(PyObject * /*exporter*/, Py_buffer *view)
{
MatrixObject *self = (MatrixObject *)view->obj;
self->flag &= ~BASE_MATH_FLAG_HAS_BUFFER_VIEW;
if (view->readonly == 0) {
if (UNLIKELY(BaseMath_WriteCallback(self) == -1)) {
PyErr_Print();
}
}
MEM_SAFE_FREE(view->shape);
MEM_SAFE_FREE(view->strides);
}
static PyBufferProcs Matrix_as_buffer = {
(getbufferproc)Matrix_getbuffer,
(releasebufferproc)Matrix_releasebuffer,
};
/** \} */
/* -------------------------------------------------------------------- */
/** \name Matrix Type: Rich Compare
* \{ */
@@ -3478,7 +3543,7 @@ PyTypeObject matrix_Type = {
/*tp_str*/ (reprfunc)Matrix_str,
/*tp_getattro*/ nullptr,
/*tp_setattro*/ nullptr,
/*tp_as_buffer*/ nullptr,
/*tp_as_buffer*/ &Matrix_as_buffer,
/*tp_flags*/ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
/*tp_doc*/ matrix_doc,
/*tp_traverse*/ (traverseproc)BaseMathObject_traverse,

View File

@@ -863,6 +863,59 @@ static PyObject *Quaternion_str(QuaternionObject *self)
/** \} */
/* -------------------------------------------------------------------- */
/** \name Quaternion Type: Buffer Protocol
* \{ */
static int Quaternion_getbuffer(PyObject *obj, Py_buffer *view, int flags)
{
QuaternionObject *self = (QuaternionObject *)obj;
if (UNLIKELY(BaseMath_Prepare_ForBufferAccess(self, view, flags) == -1)) {
return -1;
}
if (UNLIKELY(BaseMath_ReadCallback(self) == -1)) {
return -1;
}
memset(view, 0, sizeof(*view));
view->obj = (PyObject *)self;
view->buf = (void *)self->quat;
view->len = Py_ssize_t(QUAT_SIZE * sizeof(float));
view->itemsize = sizeof(float);
view->ndim = 1;
if ((flags & PyBUF_WRITABLE) == 0) {
view->readonly = 1;
}
if (flags & PyBUF_FORMAT) {
view->format = (char *)"f";
}
self->flag |= BASE_MATH_FLAG_HAS_BUFFER_VIEW;
Py_INCREF(self);
return 0;
}
static void Quaternion_releasebuffer(PyObject * /*exporter*/, Py_buffer *view)
{
QuaternionObject *self = (QuaternionObject *)view->obj;
self->flag &= ~BASE_MATH_FLAG_HAS_BUFFER_VIEW;
if (view->readonly == 0) {
if (UNLIKELY(BaseMath_WriteCallback(self) == -1)) {
PyErr_Print();
}
}
}
static PyBufferProcs Quaternion_as_buffer = {
(getbufferproc)Quaternion_getbuffer,
(releasebufferproc)Quaternion_releasebuffer,
};
/** \} */
/* -------------------------------------------------------------------- */
/** \name Quaternion Type: Rich Compare
* \{ */
@@ -1801,7 +1854,7 @@ PyTypeObject quaternion_Type = {
/*tp_str*/ (reprfunc)Quaternion_str,
/*tp_getattro*/ nullptr,
/*tp_setattro*/ nullptr,
/*tp_as_buffer*/ nullptr,
/*tp_as_buffer*/ &Quaternion_as_buffer,
/*tp_flags*/ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
/*tp_doc*/ quaternion_doc,
/*tp_traverse*/ (traverseproc)BaseMathObject_traverse,

View File

@@ -499,6 +499,7 @@ static PyObject *Vector_resize(VectorObject *self, PyObject *value)
if (UNLIKELY(BaseMathObject_Prepare_ForResize(self, "Vector.resize()") == -1)) {
/* An exception has been raised. */
return nullptr;
}
@@ -1595,6 +1596,59 @@ static PyObject *Vector_str(VectorObject *self)
/** \} */
/* -------------------------------------------------------------------- */
/** \name Vector Type: Buffer Protocol
* \{ */
static int Vector_getbuffer(PyObject *obj, Py_buffer *view, int flags)
{
VectorObject *self = (VectorObject *)obj;
if (UNLIKELY(BaseMath_Prepare_ForBufferAccess(self, view, flags) == -1)) {
return -1;
}
if (UNLIKELY(BaseMath_ReadCallback(self) == -1)) {
return -1;
}
memset(view, 0, sizeof(*view));
view->obj = (PyObject *)self;
view->buf = (void *)self->vec;
view->len = Py_ssize_t(self->vec_num * sizeof(float));
view->itemsize = sizeof(float);
view->ndim = 1;
if ((flags & PyBUF_WRITABLE) == 0) {
view->readonly = 1;
}
if (flags & PyBUF_FORMAT) {
view->format = (char *)"f";
}
self->flag |= BASE_MATH_FLAG_HAS_BUFFER_VIEW;
Py_INCREF(self);
return 0;
}
static void Vector_releasebuffer(PyObject * /*exporter*/, Py_buffer *view)
{
VectorObject *self = (VectorObject *)view->obj;
self->flag &= ~BASE_MATH_FLAG_HAS_BUFFER_VIEW;
if (view->readonly == 0) {
if (UNLIKELY(BaseMath_WriteCallback(self) == -1)) {
PyErr_Print();
}
}
}
static PyBufferProcs Vector_as_buffer = {
(getbufferproc)Vector_getbuffer,
(releasebufferproc)Vector_releasebuffer,
};
/** \} */
/* -------------------------------------------------------------------- */
/** \name Vector Type: Rich Compare
* \{ */
@@ -3411,7 +3465,7 @@ PyTypeObject vector_Type = {
/*tp_str*/ (reprfunc)Vector_str,
/*tp_getattro*/ nullptr,
/*tp_setattro*/ nullptr,
/*tp_as_buffer*/ nullptr,
/*tp_as_buffer*/ &Vector_as_buffer,
/*tp_flags*/ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
/*tp_doc*/ vector_doc,
/*tp_traverse*/ (traverseproc)BaseMathObject_traverse,

View File

@@ -4,7 +4,7 @@
# ./blender.bin --background --python tests/python/bl_pyapi_mathutils.py -- --verbose
import unittest
from mathutils import Matrix, Vector, Quaternion, Euler
from mathutils import Matrix, Vector, Quaternion, Euler, Color
from mathutils import kdtree, geometry
import math
@@ -33,6 +33,42 @@ vector_data = sum(
for sign in (1.0, -1.0))), ()) + ((0.0, 0.0, 0.0),)
def _test_flat_buffer_protocol(self, ty, n):
expected = list(range(n))
data = ty(expected)
view = memoryview(data)
self.assertEqual(view.shape, (n,))
self.assertEqual(view.format, "f")
self.assertEqual(view.tolist(), expected)
# Check multiple simultaneous.
with self.assertRaises(BufferError):
memoryview(data)
# Check frozen.
with self.assertRaises(BufferError):
data.freeze()
# Check resize.
if ty is Vector:
with self.assertRaises(BufferError):
data.resize(100)
_incref = view # For potential changes in GC.
# Check for a release buffer call, GC releases the buffer if it's not referenced.
data = ty(expected)
memoryview(data)
memoryview(data)
vec = ty(expected)
vec.freeze()
with self.assertRaises(TypeError):
view = memoryview(vec)
view[0] = 1
class MatrixTesting(unittest.TestCase):
def test_matrix_column_access(self):
# mat =
@@ -257,6 +293,15 @@ class MatrixTesting(unittest.TestCase):
with self.assertRaises(TypeError):
mat[0][0] = 0.0
def test_buffer_protocol(self):
expected = [list(range(i * 4, (i * 4) + 4)) for i in range(4)]
m = Matrix(expected)
view = memoryview(m)
self.assertEqual(view.shape, (4, 4))
self.assertEqual(view.format, "f")
self.assertEqual(view.tolist(), expected)
def assertAlmostEqualMatrix(self, first, second, size, *, places=6, msg=None, delta=None):
for i in range(size):
for j in range(size):
@@ -323,6 +368,9 @@ class VectorTesting(unittest.TestCase):
with self.assertRaises(TypeError):
vec[0] = 0.0
def test_buffer_protocol(self):
_test_flat_buffer_protocol(self, Vector, 10)
class QuaternionTesting(unittest.TestCase):
@@ -352,6 +400,21 @@ class QuaternionTesting(unittest.TestCase):
self.assertAlmostEqual(axis.y, math.sqrt(0.5), 6)
self.assertAlmostEqual(axis.z, 0)
def test_buffer_protocol(self):
_test_flat_buffer_protocol(self, Quaternion, 4)
class EulerTesting(unittest.TestCase):
def test_buffer_protocol(self):
_test_flat_buffer_protocol(self, Euler, 3)
class ColorTesting(unittest.TestCase):
def test_buffer_protocol(self):
_test_flat_buffer_protocol(self, Color, 3)
class KDTreeTesting(unittest.TestCase):
@staticmethod