PyAPI: Fast buffer access to id property arrays
authorJacques Lucke <mail@jlucke.com>
Mon, 10 Apr 2017 11:06:00 +0000 (21:06 +1000)
committerCampbell Barton <ideasman42@gmail.com>
Mon, 10 Apr 2017 11:12:02 +0000 (21:12 +1000)
Support Python's buffer protocol for ID-properties.

source/blender/blenkernel/intern/idprop.c
source/blender/python/generic/idprop_py_api.c
tests/python/bl_pyapi_idprop.py

index 0f1716dcaf7808d11e0f9aae3a6381a1a7cdc57a..5e4e8eb34ad93e4257e07653f2c3e21aeaaa0443 100644 (file)
@@ -890,7 +890,7 @@ IDProperty *IDP_New(const char type, const IDPropertyTemplate *val, const char *
                        *(float *)&prop->data.val = val->f;
                        break;
                case IDP_DOUBLE:
-                       prop = MEM_callocN(sizeof(IDProperty), "IDProperty float");
+                       prop = MEM_callocN(sizeof(IDProperty), "IDProperty double");
                        *(double *)&prop->data.val = val->d;
                        break;
                case IDP_ARRAY:
index 2e15b7b14132bf2e8831207e67f9371a56c2eb97..0a9931f2683e4abf2beb2c586e62077cc140e109 100644 (file)
@@ -335,19 +335,9 @@ static char idp_sequence_type(PyObject *seq_fast)
        return type;
 }
 
-/**
- * \note group can be a pointer array or a group.
- * assume we already checked key is a string.
- *
- * \return success.
- */
-bool BPy_IDProperty_Map_ValidateAndCreate(PyObject *name_obj, IDProperty *group, PyObject *ob)
+static const char *idp_try_read_name(PyObject *name_obj)
 {
-       IDProperty *prop = NULL;
-       IDPropertyTemplate val = {0};
-
-       const char *name;
-
+       const char *name = NULL;
        if (name_obj) {
                Py_ssize_t name_size;
                name = _PyUnicode_AsStringAndSize(name_obj, &name_size);
@@ -356,168 +346,297 @@ bool BPy_IDProperty_Map_ValidateAndCreate(PyObject *name_obj, IDProperty *group,
                        PyErr_Format(PyExc_KeyError,
                                     "invalid id-property key, expected a string, not a %.200s",
                                     Py_TYPE(name_obj)->tp_name);
-                       return false;
+                       return NULL;
                }
 
                if (name_size > MAX_IDPROP_NAME) {
                        PyErr_SetString(PyExc_KeyError, "the length of IDProperty names is limited to 63 characters");
-                       return false;
+                       return NULL;
                }
        }
        else {
                name = "";
        }
+       return name;
+}
 
-       if (PyFloat_Check(ob)) {
-               val.d = PyFloat_AsDouble(ob);
-               prop = IDP_New(IDP_DOUBLE, &val, name);
-       }
-       else if (PyLong_Check(ob)) {
-               val.i = _PyLong_AsInt(ob);
-               if (val.i == -1 && PyErr_Occurred()) {
-                       return false;
-               }
-               prop = IDP_New(IDP_INT, &val, name);
+/* -------------------------------------------------------------------------- */
+
+/**
+ * The 'idp_from_Py*' functions expect that the input type has been checked before
+ * and return NULL if the IDProperty can't be created.
+ */
+
+static IDProperty *idp_from_PyFloat(const char *name, PyObject *ob)
+{
+       IDPropertyTemplate val = {0};
+       val.d = PyFloat_AsDouble(ob);
+       return IDP_New(IDP_DOUBLE, &val, name);
+}
+
+static IDProperty *idp_from_PyLong(const char *name, PyObject *ob)
+{
+       IDPropertyTemplate val = {0};
+       val.i = _PyLong_AsInt(ob);
+       if (val.i == -1 && PyErr_Occurred()) {
+               return NULL;
        }
-       else if (PyUnicode_Check(ob)) {
+       return IDP_New(IDP_INT, &val, name);
+}
+
+static IDProperty *idp_from_PyUnicode(const char *name, PyObject *ob)
+{
+       IDProperty *prop;
+       IDPropertyTemplate val = {0};
 #ifdef USE_STRING_COERCE
-               Py_ssize_t value_size;
-               PyObject *value_coerce = NULL;
-               val.string.str = PyC_UnicodeAsByteAndSize(ob, &value_size, &value_coerce);
-               val.string.len = (int)value_size + 1;
-               val.string.subtype = IDP_STRING_SUB_UTF8;
-               prop = IDP_New(IDP_STRING, &val, name);
-               Py_XDECREF(value_coerce);
+       Py_ssize_t value_size;
+       PyObject *value_coerce = NULL;
+       val.string.str = PyC_UnicodeAsByteAndSize(ob, &value_size, &value_coerce);
+       val.string.len = (int)value_size + 1;
+       val.string.subtype = IDP_STRING_SUB_UTF8;
+       prop = IDP_New(IDP_STRING, &val, name);
+       Py_XDECREF(value_coerce);
 #else
-               val.str = _PyUnicode_AsString(ob);
-               prop = IDP_New(IDP_STRING, val, name);
+       val.str = _PyUnicode_AsString(ob);
+       prop = IDP_New(IDP_STRING, val, name);
 #endif
-       }
-       else if (PyBytes_Check(ob)) {
-               val.string.str = PyBytes_AS_STRING(ob);
-               val.string.len = PyBytes_GET_SIZE(ob);
-               val.string.subtype = IDP_STRING_SUB_BYTE;
+       return prop;
+}
 
-               prop = IDP_New(IDP_STRING, &val, name);
-               //prop = IDP_NewString(PyBytes_AS_STRING(ob), name, PyBytes_GET_SIZE(ob));
-               //prop->subtype = IDP_STRING_SUB_BYTE;
+static IDProperty *idp_from_PyBytes(const char *name, PyObject *ob)
+{
+       IDPropertyTemplate val = {0};
+       val.string.str = PyBytes_AS_STRING(ob);
+       val.string.len = PyBytes_GET_SIZE(ob);
+       val.string.subtype = IDP_STRING_SUB_BYTE;
+       return IDP_New(IDP_STRING, &val, name);
+}
+
+static int idp_array_type_from_format_char(char format)
+{
+       if (format == 'i') return IDP_INT;
+       if (format == 'f') return IDP_FLOAT;
+       if (format == 'd') return IDP_DOUBLE;
+       return -1;
+}
+
+static const char *idp_format_from_array_type(int type)
+{
+       if (type == IDP_INT) return "i";
+       if (type == IDP_FLOAT) return "f";
+       if (type == IDP_DOUBLE) return "d";
+       return NULL;
+}
+
+static IDProperty *idp_from_PySequence_Buffer(const char *name, Py_buffer *buffer)
+{
+       IDProperty *prop;
+       IDPropertyTemplate val = {0};
+
+       int format = idp_array_type_from_format_char(*buffer->format);
+       if (format == -1) {
+               /* should never happen as the type has been checked before */
+               return NULL;
        }
-       else if (PySequence_Check(ob)) {
-               PyObject *ob_seq_fast;
-               PyObject **ob_seq_fast_items;
-               PyObject *item;
-               int i;
+       else {
+               val.array.type = format;
+               val.array.len = buffer->len / buffer->itemsize;
+       }
+       prop = IDP_New(IDP_ARRAY, &val, name);
+       memcpy(IDP_Array(prop), buffer->buf, buffer->len);
+       return prop;
+}
 
-               if (!(ob_seq_fast = PySequence_Fast(ob, "py -> idprop"))) {
-                       return false;
-               }
+static IDProperty *idp_from_PySequence_Fast(const char *name, PyObject *ob)
+{
+       IDProperty *prop;
+       IDPropertyTemplate val = {0};
 
-               ob_seq_fast_items = PySequence_Fast_ITEMS(ob_seq_fast);
+       PyObject **ob_seq_fast_items;
+       PyObject *item;
+       int i;
 
-               if ((val.array.type = idp_sequence_type(ob_seq_fast)) == (char)-1) {
-                       Py_DECREF(ob_seq_fast);
-                       PyErr_SetString(PyExc_TypeError, "only floats, ints and dicts are allowed in ID property arrays");
-                       return false;
-               }
+       ob_seq_fast_items = PySequence_Fast_ITEMS(ob);
 
-               /* validate sequence and derive type.
-                * we assume IDP_INT unless we hit a float
-                * number; then we assume it's */
+       if ((val.array.type = idp_sequence_type(ob)) == (char)-1) {
+               PyErr_SetString(PyExc_TypeError, "only floats, ints and dicts are allowed in ID property arrays");
+               return NULL;
+       }
 
-               val.array.len = PySequence_Fast_GET_SIZE(ob_seq_fast);
+       /* validate sequence and derive type.
+        * we assume IDP_INT unless we hit a float
+        * number; then we assume it's */
 
-               switch (val.array.type) {
-                       case IDP_DOUBLE:
-                       {
-                               double *prop_data;
-
-                               prop = IDP_New(IDP_ARRAY, &val, name);
-                               prop_data = IDP_Array(prop);
-                               for (i = 0; i < val.array.len; i++) {
-                                       item = ob_seq_fast_items[i];
-                                       if (((prop_data[i] = PyFloat_AsDouble(item)) == -1.0) && PyErr_Occurred()) {
-                                               Py_DECREF(ob_seq_fast);
-                                               return false;
-                                       }
+       val.array.len = PySequence_Fast_GET_SIZE(ob);
+
+       switch (val.array.type) {
+               case IDP_DOUBLE:
+               {
+                       double *prop_data;
+                       prop = IDP_New(IDP_ARRAY, &val, name);
+                       prop_data = IDP_Array(prop);
+                       for (i = 0; i < val.array.len; i++) {
+                               item = ob_seq_fast_items[i];
+                               if (((prop_data[i] = PyFloat_AsDouble(item)) == -1.0) && PyErr_Occurred()) {
+                                       return NULL;
                                }
-                               break;
                        }
-                       case IDP_INT:
-                       {
-                               int *prop_data;
-                               prop = IDP_New(IDP_ARRAY, &val, name);
-                               prop_data = IDP_Array(prop);
-                               for (i = 0; i < val.array.len; i++) {
-                                       item = ob_seq_fast_items[i];
-                                       if (((prop_data[i] = _PyLong_AsInt(item)) == -1) && PyErr_Occurred()) {
-                                               Py_DECREF(ob_seq_fast);
-                                               return false;
-                                       }
+                       break;
+               }
+               case IDP_INT:
+               {
+                       int *prop_data;
+                       prop = IDP_New(IDP_ARRAY, &val, name);
+                       prop_data = IDP_Array(prop);
+                       for (i = 0; i < val.array.len; i++) {
+                               item = ob_seq_fast_items[i];
+                               if (((prop_data[i] = _PyLong_AsInt(item)) == -1) && PyErr_Occurred()) {
+                                       return NULL;
                                }
-                               break;
                        }
-                       case IDP_IDPARRAY:
-                       {
-                               prop = IDP_NewIDPArray(name);
-                               for (i = 0; i < val.array.len; i++) {
-                                       item = ob_seq_fast_items[i];
-
-                                       if (BPy_IDProperty_Map_ValidateAndCreate(NULL, prop, item) == false) {
-                                               Py_DECREF(ob_seq_fast);
-                                               return false;
-                                       }
+                       break;
+               }
+               case IDP_IDPARRAY:
+               {
+                       prop = IDP_NewIDPArray(name);
+                       for (i = 0; i < val.array.len; i++) {
+                               item = ob_seq_fast_items[i];
+                               if (BPy_IDProperty_Map_ValidateAndCreate(NULL, prop, item) == false) {
+                                       return NULL;
                                }
-                               break;
                        }
-                       default:
-                               /* should never happen */
-                               Py_DECREF(ob_seq_fast);
-                               PyErr_SetString(PyExc_RuntimeError, "internal error with idp array.type");
-                               return false;
+                       break;
+               }
+               default:
+                       /* should never happen */
+                       PyErr_SetString(PyExc_RuntimeError, "internal error with idp array.type");
+                       return NULL;
+       }
+       return prop;
+}
+
+
+static IDProperty *idp_from_PySequence(const char *name, PyObject *ob)
+{
+       Py_buffer buffer;
+       bool use_buffer = false;
+
+       if (PyObject_CheckBuffer(ob)) {
+               PyObject_GetBuffer(ob, &buffer, PyBUF_SIMPLE | PyBUF_FORMAT);
+               char format = *buffer.format;
+               if (ELEM(format, 'i', 'f', 'd')) {
+                       use_buffer = true;
                }
+               else {
+                       PyBuffer_Release(&buffer);
+               }
+       }
 
-               Py_DECREF(ob_seq_fast);
+       if (use_buffer) {
+               IDProperty *prop = idp_from_PySequence_Buffer(name, &buffer);
+               PyBuffer_Release(&buffer);
+               return prop;
        }
-       else if (PyMapping_Check(ob)) {
-               PyObject *keys, *vals, *key, *pval;
-               int i, len;
-               /*yay! we get into recursive stuff now!*/
-               keys = PyMapping_Keys(ob);
-               vals = PyMapping_Values(ob);
-
-               /* we allocate the group first; if we hit any invalid data,
-                * we can delete it easily enough.*/
-               prop = IDP_New(IDP_GROUP, &val, name);
-               len = PyMapping_Length(ob);
-               for (i = 0; i < len; i++) {
-                       key = PySequence_GetItem(keys, i);
-                       pval = PySequence_GetItem(vals, i);
-                       if (BPy_IDProperty_Map_ValidateAndCreate(key, prop, pval) == false) {
-                               IDP_FreeProperty(prop);
-                               MEM_freeN(prop);
-                               Py_XDECREF(keys);
-                               Py_XDECREF(vals);
-                               Py_XDECREF(key);
-                               Py_XDECREF(pval);
-                               /* error is already set */
-                               return false;
-                       }
+       else {
+               PyObject *ob_seq_fast = PySequence_Fast(ob, "py -> idprop");
+               if (ob_seq_fast != NULL) {
+                       IDProperty *prop = idp_from_PySequence_Fast(name, ob_seq_fast);
+                       Py_DECREF(ob_seq_fast);
+                       return prop;
+               }
+               else {
+                       return NULL;
+               }
+       }
+}
+
+static IDProperty *idp_from_PyMapping(const char *name, PyObject *ob)
+{
+       IDProperty *prop;
+       IDPropertyTemplate val = {0};
+
+       PyObject *keys, *vals, *key, *pval;
+       int i, len;
+       /* yay! we get into recursive stuff now! */
+       keys = PyMapping_Keys(ob);
+       vals = PyMapping_Values(ob);
+
+       /* we allocate the group first; if we hit any invalid data,
+        * we can delete it easily enough.*/
+       prop = IDP_New(IDP_GROUP, &val, name);
+       len = PyMapping_Length(ob);
+       for (i = 0; i < len; i++) {
+               key = PySequence_GetItem(keys, i);
+               pval = PySequence_GetItem(vals, i);
+               if (BPy_IDProperty_Map_ValidateAndCreate(key, prop, pval) == false) {
+                       IDP_FreeProperty(prop);
+                       MEM_freeN(prop);
+                       Py_XDECREF(keys);
+                       Py_XDECREF(vals);
                        Py_XDECREF(key);
                        Py_XDECREF(pval);
+                       /* error is already set */
+                       return NULL;
                }
-               Py_XDECREF(keys);
-               Py_XDECREF(vals);
+               Py_XDECREF(key);
+               Py_XDECREF(pval);
+       }
+       Py_XDECREF(keys);
+       Py_XDECREF(vals);
+       return prop;
+}
+
+static IDProperty *idp_from_PyObject(PyObject *name_obj, PyObject *ob)
+{
+       const char *name = idp_try_read_name(name_obj);
+       if (name == NULL) {
+               return NULL;
+       }
+
+       if (PyFloat_Check(ob)) {
+               return idp_from_PyFloat(name, ob);
+       }
+       else if (PyLong_Check(ob)) {
+               return idp_from_PyLong(name, ob);
+       }
+       else if (PyUnicode_Check(ob)) {
+               return idp_from_PyUnicode(name, ob);
+       }
+       else if (PyBytes_Check(ob)) {
+               return idp_from_PyBytes(name, ob);
+       }
+       else if (PySequence_Check(ob)) {
+               return idp_from_PySequence(name, ob);
+       }
+       else if (PyMapping_Check(ob)) {
+               return idp_from_PyMapping(name, ob);
        }
        else {
                PyErr_Format(PyExc_TypeError,
                             "invalid id-property type %.200s not supported",
                             Py_TYPE(ob)->tp_name);
+               return NULL;
+       }
+}
+
+/* -------------------------------------------------------------------------- */
+
+/**
+ * \note group can be a pointer array or a group.
+ * assume we already checked key is a string.
+ *
+ * \return success.
+ */
+bool BPy_IDProperty_Map_ValidateAndCreate(PyObject *name_obj, IDProperty *group, PyObject *ob)
+{
+       IDProperty *prop = idp_from_PyObject(name_obj, ob);
+       if (prop == NULL) {
                return false;
        }
 
        if (group->type == IDP_IDPARRAY) {
                IDP_AppendArray(group, prop);
-               // IDP_FreeProperty(item);  /* IDP_AppendArray does a shallow copy (memcpy), only free memory */
+               /* IDP_AppendArray does a shallow copy (memcpy), only free memory */
                MEM_freeN(prop);
        }
        else {
@@ -1371,6 +1490,44 @@ static PyMappingMethods BPy_IDArray_AsMapping = {
        (objobjargproc)BPy_IDArray_ass_subscript
 };
 
+static int itemsize_by_idarray_type(int array_type)
+{
+       if (array_type == IDP_INT) return sizeof(int);
+       if (array_type == IDP_FLOAT) return sizeof(float);
+       if (array_type == IDP_DOUBLE) return sizeof(double);
+       return -1;  /* should never happen */
+}
+
+static int BPy_IDArray_getbuffer(BPy_IDArray *self, Py_buffer *view, int flags)
+{
+       IDProperty *prop = self->prop;
+       int itemsize = itemsize_by_idarray_type(prop->subtype);
+       int length = itemsize * prop->len;
+
+       if (PyBuffer_FillInfo(view, (PyObject *)self, IDP_Array(prop), length, false, flags) == -1) {
+               return -1;
+       }
+
+       view->itemsize = itemsize;
+       view->format = (char *)idp_format_from_array_type(prop->subtype);
+
+       Py_ssize_t *shape = MEM_mallocN(sizeof(Py_ssize_t), __func__);
+       shape[0] = prop->len;
+       view->shape = shape;
+
+       return 0;
+}
+
+static void BPy_IDArray_releasebuffer(BPy_IDArray *UNUSED(self), Py_buffer *view)
+{
+       MEM_freeN(view->shape);
+}
+
+static PyBufferProcs BPy_IDArray_Buffer = {
+       (getbufferproc)BPy_IDArray_getbuffer,
+       (releasebufferproc)BPy_IDArray_releasebuffer,
+};
+
 
 PyTypeObject BPy_IDArray_Type = {
        PyVarObject_HEAD_INIT(NULL, 0)
@@ -1403,7 +1560,7 @@ PyTypeObject BPy_IDArray_Type = {
        NULL,                       /* setattrofunc tp_setattro; */
 
        /* Functions to access object as input/output buffer */
-       NULL,                       /* PyBufferProcs *tp_as_buffer; */
+       &BPy_IDArray_Buffer,        /* PyBufferProcs *tp_as_buffer; */
 
        /*** Flags to define presence of optional/expanded features ***/
        Py_TPFLAGS_DEFAULT,         /* long tp_flags; */
index 0a9cb044571c9cfebf8a3060efa877b46e16b977..7bf68c16cc72966068de391b355084f12dd7acaa 100644 (file)
@@ -3,6 +3,7 @@
 # ./blender.bin --background -noaudio --python tests/python/bl_pyapi_idprop.py -- --verbose
 import bpy
 import unittest
+import numpy as np
 from array import array
 
 
@@ -75,7 +76,7 @@ class TestIdPropertyCreation(TestHelper, unittest.TestCase):
         mylist = [1.2, 3.4, 5.6]
         self.id["a"] = array("f", mylist)
         self.assertAlmostEqualSeq(self.id["a"].to_list(), mylist)
-        self.assertEqual(self.id["a"].typecode, "d")
+        self.assertEqual(self.id["a"].typecode, "f")
 
     def test_sequence_double_array(self):
         mylist = [1.2, 3.4, 5.6]
@@ -138,6 +139,65 @@ class TestIdPropertyCreation(TestHelper, unittest.TestCase):
             self.id["a"] = self
 
 
+class TestBufferProtocol(TestHelper, unittest.TestCase):
+
+    def test_int(self):
+        self.id["a"] = array("i", [1, 2, 3, 4, 5])
+        a = np.frombuffer(self.id["a"], self.id["a"].typecode)
+        self.assertEqual(len(a), 5)
+        a[2] = 10
+        self.assertEqual(self.id["a"].to_list(), [1, 2, 10, 4, 5])
+
+    def test_float(self):
+        self.id["a"] = array("f", [1.0, 2.0, 3.0, 4.0])
+        a = np.frombuffer(self.id["a"], self.id["a"].typecode)
+        self.assertEqual(len(a), 4)
+        a[-1] = 10
+        self.assertEqual(self.id["a"].to_list(), [1.0, 2.0, 3.0, 10.0])
+
+    def test_double(self):
+        self.id["a"] = array("d", [1.0, 2.0, 3.0, 4.0])
+        a = np.frombuffer(self.id["a"], self.id["a"].typecode)
+        a[1] = 10
+        self.assertEqual(self.id["a"].to_list(), [1.0, 10.0, 3.0, 4.0])
+
+    def test_full_update(self):
+        self.id["a"] = array("i", [1, 2, 3, 4, 5, 6])
+        a = np.frombuffer(self.id["a"], self.id["a"].typecode)
+        a[:] = [10, 20, 30, 40, 50, 60]
+        self.assertEqual(self.id["a"].to_list(), [10, 20, 30, 40, 50, 60])
+
+    def test_partial_update(self):
+        self.id["a"] = array("i", [1, 2, 3, 4, 5, 6, 7, 8])
+        a = np.frombuffer(self.id["a"], self.id["a"].typecode)
+        a[1:5] = [10, 20, 30, 40]
+        self.assertEqual(self.id["a"].to_list(), [1, 10, 20, 30, 40, 6, 7, 8])
+
+    def test_copy(self):
+        self.id["a"] = array("i", [1, 2, 3, 4, 5])
+        self.id["b"] = self.id["a"]
+        self.assertEqual(self.id["a"].to_list(), self.id["b"].to_list())
+
+    def test_memview_attributes(self):
+        mylist = [1, 2, 3]
+        self.id["a"] = mylist
+
+        view1 = memoryview(self.id["a"])
+        view2 = memoryview(array("i", mylist))
+
+        self.assertEqualMemviews(view1, view2)
+
+    def assertEqualMemviews(self, view1, view2):
+        props_to_compare = (
+            "contiguous", "format", "itemsize", "nbytes", "ndim",
+            "readonly", "shape", "strides", "suboffsets"
+        )
+        for attr in props_to_compare:
+            self.assertEqual(getattr(view1, attr), getattr(view2, attr))
+
+        self.assertEqual(list(view1), list(view2))
+        self.assertEqual(view1.tobytes(), view2.tobytes())
+
 if __name__ == '__main__':
     import sys
     sys.argv = [__file__] + (sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else [])