Fix buffer classes using super()

python · JelleZijlstra · May 8, 2023 · May 6, 2023 · May 6, 2023 · May 6, 2023
commit b7fb7a47071bb85c054576d6e0cb17b6622505e2
diff --git a/Include/internal/pycore_memoryobject.h b/Include/internal/pycore_memoryobject.h
@@ -9,7 +9,8 @@ extern "C" {
 #endif
 
 PyObject *
-PyMemoryView_FromObjectAndFlags(PyObject *v, int flags);
+_PyMemoryView_FromBufferProc(PyObject *v, int flags,
+                             getbufferproc bufferproc);
 
 #ifdef __cplusplus
 }

diff --git a/Lib/test/test_buffer.py b/Lib/test/test_buffer.py
@@ -4579,6 +4579,15 @@ def test_c_buffer(self):
             buf.__release_buffer__(mv)
         self.assertEqual(buf.references, 0)
 
+    def test_inheritance(self):
+        class A(bytearray):
+            def __buffer__(self, flags):
+                return super().__buffer__(flags)
+
+        a = A(b"hello")
+        mv = memoryview(a)
+        self.assertEqual(mv.tobytes(), b"hello")
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Objects/bytearrayobject.c b/Objects/bytearrayobject.c
@@ -61,6 +61,7 @@ static void
 bytearray_releasebuffer(PyByteArrayObject *obj, Py_buffer *view)
 {
     obj->ob_exports--;
+    assert(obj->ob_exports >= 0);
 }
 
 static int

diff --git a/Objects/memoryobject.c b/Objects/memoryobject.c
@@ -781,7 +781,7 @@ PyMemoryView_FromBuffer(const Py_buffer *info)
    using the given flags.
    If the object is a memoryview, the new memoryview must be registered
    with the same managed buffer. Otherwise, a new managed buffer is created. */
-PyObject *
+static PyObject *
 PyMemoryView_FromObjectAndFlags(PyObject *v, int flags)
 {
     _PyManagedBufferObject *mbuf;
@@ -806,6 +806,30 @@ PyMemoryView_FromObjectAndFlags(PyObject *v, int flags)
         Py_TYPE(v)->tp_name);
     return NULL;
 }
+
+/* Create a memoryview from an object that implements the buffer protocol,
+   using the given flags.
+   If the object is a memoryview, the new memoryview must be registered
+   with the same managed buffer. Otherwise, a new managed buffer is created. */
+PyObject *
+_PyMemoryView_FromBufferProc(PyObject *v, int flags, getbufferproc bufferproc)
+{
+    _PyManagedBufferObject *mbuf = mbuf_alloc();
+    if (mbuf == NULL)
+        return NULL;
+
+    int res = bufferproc(v, &mbuf->master, flags);
+    if (res < 0) {
+        mbuf->master.obj = NULL;
+        Py_DECREF(mbuf);
+        return NULL;
+    }
+
+    PyObject *ret = mbuf_add_view(mbuf, NULL);
+    Py_DECREF(mbuf);
+    return ret;
+}
+
 /* Create a memoryview from an object that implements the buffer protocol.
    If the object is a memoryview, the new memoryview must be registered
    with the same managed buffer. Otherwise, a new managed buffer is created. */

@@ -6,7 +6,7 @@
 #include "pycore_symtable.h"      // _Py_Mangle()
 #include "pycore_dict.h"          // _PyDict_KeysSize()
 #include "pycore_initconfig.h"    // _PyStatus_OK()
-#include "pycore_memoryobject.h"  // PyMemoryView_FromObjectAndFlags()
+#include "pycore_memoryobject.h"  // _PyMemoryView_FromBufferProc()
 #include "pycore_moduleobject.h"  // _PyModule_GetDef()
 #include "pycore_object.h"        // _PyType_HasFeature()
 #include "pycore_long.h"          // _PyLong_IsNegative()
@@ -56,6 +56,8 @@ typedef struct PySlot_Offset {
     short slot_offset;
 } PySlot_Offset;
 
+static void
+slot_bf_releasebuffer(PyObject *self, Py_buffer *buffer);
 
 static PyObject *
 slot_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds);
@@ -8078,7 +8080,8 @@ wrap_buffer(PyObject *self, PyObject *args, void *wrapped)
         return NULL;
     }
 
-    return PyMemoryView_FromObjectAndFlags(self, Py_SAFE_DOWNCAST(flags, Py_ssize_t, int));
+    return _PyMemoryView_FromBufferProc(self, Py_SAFE_DOWNCAST(flags, Py_ssize_t, int),
+                                        (getbufferproc)wrapped);
 }
 
 static PyObject *
@@ -8980,8 +8983,10 @@ bufferwrapper_releasebuf(PyObject *self, Py_buffer *view)
 
     assert(PyMemoryView_Check(bw->mv));
     Py_TYPE(bw->mv)->tp_as_buffer->bf_releasebuffer(bw->mv, view);
+    // We only need to call bf_releasebuffer if it's a Python function. If it's a C
+    // bf_releasebuf, it will be called when the memoryview is released.
     if (Py_TYPE(bw->obj)->tp_as_buffer != NULL
-        && Py_TYPE(bw->obj)->tp_as_buffer->bf_releasebuffer != NULL) {
+        && Py_TYPE(bw->obj)->tp_as_buffer->bf_releasebuffer == slot_bf_releasebuffer) {
         Py_TYPE(bw->obj)->tp_as_buffer->bf_releasebuffer(bw->obj, view);
     }
 }