bpo-36785: PEP 574 implementation (GH-7076)

1 files changed, 219 insertions, 0 deletions

diff --git a/Objects/picklebufobject.c b/Objects/picklebufobject.c
new file mode 100644
index 0000000000..a135e5575e
--- /dev/null
+++ b/Objects/picklebufobject.c
@@ -0,0 +1,219 @@
+/* PickleBuffer object implementation */
+
+#define PY_SSIZE_T_CLEAN
+#include "Python.h"
+#include <stddef.h>
+
+typedef struct {
+    PyObject_HEAD
+    /* The view exported by the original object */
+    Py_buffer view;
+    PyObject *weakreflist;
+} PyPickleBufferObject;
+
+/* C API */
+
+PyObject *
+PyPickleBuffer_FromObject(PyObject *base)
+{
+    PyTypeObject *type = &PyPickleBuffer_Type;
+    PyPickleBufferObject *self;
+
+    self = (PyPickleBufferObject *) type->tp_alloc(type, 0);
+    if (self == NULL) {
+        return NULL;
+    }
+    self->view.obj = NULL;
+    self->weakreflist = NULL;
+    if (PyObject_GetBuffer(base, &self->view, PyBUF_FULL_RO) < 0) {
+        Py_DECREF(self);
+        return NULL;
+    }
+    return (PyObject *) self;
+}
+
+const Py_buffer *
+PyPickleBuffer_GetBuffer(PyObject *obj)
+{
+    PyPickleBufferObject *self = (PyPickleBufferObject *) obj;
+
+    if (!PyPickleBuffer_Check(obj)) {
+        PyErr_Format(PyExc_TypeError,
+                     "expected PickleBuffer, %.200s found",
+                     Py_TYPE(obj)->tp_name);
+        return NULL;
+    }
+    if (self->view.obj == NULL) {
+        PyErr_SetString(PyExc_ValueError,
+                        "operation forbidden on released PickleBuffer object");
+        return NULL;
+    }
+    return &self->view;
+}
+
+int
+PyPickleBuffer_Release(PyObject *obj)
+{
+    PyPickleBufferObject *self = (PyPickleBufferObject *) obj;
+
+    if (!PyPickleBuffer_Check(obj)) {
+        PyErr_Format(PyExc_TypeError,
+                     "expected PickleBuffer, %.200s found",
+                     Py_TYPE(obj)->tp_name);
+        return -1;
+    }
+    PyBuffer_Release(&self->view);
+    return 0;
+}
+
+static PyObject *
+picklebuf_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+    PyPickleBufferObject *self;
+    PyObject *base;
+    char *keywords[] = {"", NULL};
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:PickleBuffer",
+                                     keywords, &base)) {
+        return NULL;
+    }
+
+    self = (PyPickleBufferObject *) type->tp_alloc(type, 0);
+    if (self == NULL) {
+        return NULL;
+    }
+    self->view.obj = NULL;
+    self->weakreflist = NULL;
+    if (PyObject_GetBuffer(base, &self->view, PyBUF_FULL_RO) < 0) {
+        Py_DECREF(self);
+        return NULL;
+    }
+    return (PyObject *) self;
+}
+
+static int
+picklebuf_traverse(PyPickleBufferObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(self->view.obj);
+    return 0;
+}
+
+static int
+picklebuf_clear(PyPickleBufferObject *self)
+{
+    PyBuffer_Release(&self->view);
+    return 0;
+}
+
+static void
+picklebuf_dealloc(PyPickleBufferObject *self)
+{
+    PyObject_GC_UnTrack(self);
+    if (self->weakreflist != NULL)
+        PyObject_ClearWeakRefs((PyObject *) self);
+    PyBuffer_Release(&self->view);
+    Py_TYPE(self)->tp_free((PyObject *) self);
+}
+
+/* Buffer API */
+
+static int
+picklebuf_getbuf(PyPickleBufferObject *self, Py_buffer *view, int flags)
+{
+    if (self->view.obj == NULL) {
+        PyErr_SetString(PyExc_ValueError,
+                        "operation forbidden on released PickleBuffer object");
+        return -1;
+    }
+    return PyObject_GetBuffer(self->view.obj, view, flags);
+}
+
+static void
+picklebuf_releasebuf(PyPickleBufferObject *self, Py_buffer *view)
+{
+    /* Since our bf_getbuffer redirects to the original object, this
+     * implementation is never called.  It only exists to signal that
+     * buffers exported by PickleBuffer have non-trivial releasing
+     * behaviour (see check in Python/getargs.c).
+     */
+}
+
+static PyBufferProcs picklebuf_as_buffer = {
+    .bf_getbuffer = (getbufferproc) picklebuf_getbuf,
+    .bf_releasebuffer = (releasebufferproc) picklebuf_releasebuf,
+};
+
+/* Methods */
+
+static PyObject *
+picklebuf_raw(PyPickleBufferObject *self, PyObject *Py_UNUSED(ignored))
+{
+    if (self->view.obj == NULL) {
+        PyErr_SetString(PyExc_ValueError,
+                        "operation forbidden on released PickleBuffer object");
+        return NULL;
+    }
+    if (self->view.suboffsets != NULL
+        || !PyBuffer_IsContiguous(&self->view, 'A')) {
+        PyErr_SetString(PyExc_BufferError,
+                        "cannot extract raw buffer from non-contiguous buffer");
+        return NULL;
+    }
+    PyObject *m = PyMemoryView_FromObject((PyObject *) self);
+    if (m == NULL) {
+        return NULL;
+    }
+    PyMemoryViewObject *mv = (PyMemoryViewObject *) m;
+    assert(mv->view.suboffsets == NULL);
+    /* Mutate memoryview instance to make it a "raw" memoryview */
+    mv->view.format = "B";
+    mv->view.ndim = 1;
+    mv->view.itemsize = 1;
+    /* shape = (length,) */
+    mv->view.shape = &mv->view.len;
+    /* strides = (1,) */
+    mv->view.strides = &mv->view.itemsize;
+    /* Fix memoryview state flags */
+    /* XXX Expose memoryobject.c's init_flags() instead? */
+    mv->flags = _Py_MEMORYVIEW_C | _Py_MEMORYVIEW_FORTRAN;
+    return m;
+}
+
+PyDoc_STRVAR(picklebuf_raw_doc,
+"raw($self, /)\n--\n\
+\n\
+Return a memoryview of the raw memory underlying this buffer.\n\
+Will raise BufferError is the buffer isn't contiguous.");
+
+static PyObject *
+picklebuf_release(PyPickleBufferObject *self, PyObject *Py_UNUSED(ignored))
+{
+    PyBuffer_Release(&self->view);
+    Py_RETURN_NONE;
+}
+
+PyDoc_STRVAR(picklebuf_release_doc,
+"release($self, /)\n--\n\
+\n\
+Release the underlying buffer exposed by the PickleBuffer object.");
+
+static PyMethodDef picklebuf_methods[] = {
+    {"raw",     (PyCFunction) picklebuf_raw,     METH_NOARGS, picklebuf_raw_doc},
+    {"release", (PyCFunction) picklebuf_release, METH_NOARGS, picklebuf_release_doc},
+    {NULL,      NULL}
+};
+
+PyTypeObject PyPickleBuffer_Type = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    .tp_name = "pickle.PickleBuffer",
+    .tp_doc = "Wrapper for potentially out-of-band buffers",
+    .tp_basicsize = sizeof(PyPickleBufferObject),
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
+    .tp_new = picklebuf_new,
+    .tp_dealloc = (destructor) picklebuf_dealloc,
+    .tp_traverse = (traverseproc) picklebuf_traverse,
+    .tp_clear = (inquiry) picklebuf_clear,
+    .tp_weaklistoffset = offsetof(PyPickleBufferObject, weakreflist),
+    .tp_as_buffer = &picklebuf_as_buffer,
+    .tp_methods = picklebuf_methods,
+};