From b856b6010ea331ffb568d7bd9eab64f5501bb263 Mon Sep 17 00:00:00 2001 From: Oxicid Date: Sat, 16 Aug 2025 06:14:19 +0000 Subject: [PATCH] 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 --- source/blender/python/mathutils/mathutils.cc | 35 ++++++++++ source/blender/python/mathutils/mathutils.hh | 20 ++++++ .../python/mathutils/mathutils_Color.cc | 55 ++++++++++++++- .../python/mathutils/mathutils_Euler.cc | 55 ++++++++++++++- .../python/mathutils/mathutils_Matrix.cc | 67 ++++++++++++++++++- .../python/mathutils/mathutils_Quaternion.cc | 55 ++++++++++++++- .../python/mathutils/mathutils_Vector.cc | 56 +++++++++++++++- tests/python/bl_pyapi_mathutils.py | 65 +++++++++++++++++- 8 files changed, 402 insertions(+), 6 deletions(-) diff --git a/source/blender/python/mathutils/mathutils.cc b/source/blender/python/mathutils/mathutils.cc index 9059bd8b1e5..0a478160a61 100644 --- a/source/blender/python/mathutils/mathutils.cc +++ b/source/blender/python/mathutils/mathutils.cc @@ -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); diff --git a/source/blender/python/mathutils/mathutils.hh b/source/blender/python/mathutils/mathutils.hh index 93ef21d12c0..9b30d474c1f 100644 --- a/source/blender/python/mathutils/mathutils.hh +++ b/source/blender/python/mathutils/mathutils.hh @@ -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. diff --git a/source/blender/python/mathutils/mathutils_Color.cc b/source/blender/python/mathutils/mathutils_Color.cc index 446d326333b..e945073c603 100644 --- a/source/blender/python/mathutils/mathutils_Color.cc +++ b/source/blender/python/mathutils/mathutils_Color.cc @@ -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, diff --git a/source/blender/python/mathutils/mathutils_Euler.cc b/source/blender/python/mathutils/mathutils_Euler.cc index 7f7de3f609a..a33a25b31bb 100644 --- a/source/blender/python/mathutils/mathutils_Euler.cc +++ b/source/blender/python/mathutils/mathutils_Euler.cc @@ -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, diff --git a/source/blender/python/mathutils/mathutils_Matrix.cc b/source/blender/python/mathutils/mathutils_Matrix.cc index aa540207f68..e00d0a449b3 100644 --- a/source/blender/python/mathutils/mathutils_Matrix.cc +++ b/source/blender/python/mathutils/mathutils_Matrix.cc @@ -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(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(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, diff --git a/source/blender/python/mathutils/mathutils_Quaternion.cc b/source/blender/python/mathutils/mathutils_Quaternion.cc index fe5448c0981..2b8941423c3 100644 --- a/source/blender/python/mathutils/mathutils_Quaternion.cc +++ b/source/blender/python/mathutils/mathutils_Quaternion.cc @@ -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, diff --git a/source/blender/python/mathutils/mathutils_Vector.cc b/source/blender/python/mathutils/mathutils_Vector.cc index 8d80aa9d11a..3387e126e46 100644 --- a/source/blender/python/mathutils/mathutils_Vector.cc +++ b/source/blender/python/mathutils/mathutils_Vector.cc @@ -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, diff --git a/tests/python/bl_pyapi_mathutils.py b/tests/python/bl_pyapi_mathutils.py index d4c5bc58ad2..f3f3a269d09 100644 --- a/tests/python/bl_pyapi_mathutils.py +++ b/tests/python/bl_pyapi_mathutils.py @@ -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