Line data    Source code

       1             : /* PickleBuffer object implementation */
       2             : 
       3             : #define PY_SSIZE_T_CLEAN
       4             : #include "Python.h"
       5             : #include <stddef.h>
       6             : 
       7             : typedef struct {
       8             :     PyObject_HEAD
       9             :     /* The view exported by the original object */
      10             :     Py_buffer view;
      11             :     PyObject *weakreflist;
      12             : } PyPickleBufferObject;
      13             : 
      14             : /* C API */
      15             : 
      16             : PyObject *
      17           0 : PyPickleBuffer_FromObject(PyObject *base)
      18             : {
      19           0 :     PyTypeObject *type = &PyPickleBuffer_Type;
      20             :     PyPickleBufferObject *self;
      21             : 
      22           0 :     self = (PyPickleBufferObject *) type->tp_alloc(type, 0);
      23           0 :     if (self == NULL) {
      24           0 :         return NULL;
      25             :     }
      26           0 :     self->view.obj = NULL;
      27           0 :     self->weakreflist = NULL;
      28           0 :     if (PyObject_GetBuffer(base, &self->view, PyBUF_FULL_RO) < 0) {
      29           0 :         Py_DECREF(self);
      30           0 :         return NULL;
      31             :     }
      32           0 :     return (PyObject *) self;
      33             : }
      34             : 
      35             : const Py_buffer *
      36         110 : PyPickleBuffer_GetBuffer(PyObject *obj)
      37             : {
      38         110 :     PyPickleBufferObject *self = (PyPickleBufferObject *) obj;
      39             : 
      40         110 :     if (!PyPickleBuffer_Check(obj)) {
      41           0 :         PyErr_Format(PyExc_TypeError,
      42             :                      "expected PickleBuffer, %.200s found",
      43           0 :                      Py_TYPE(obj)->tp_name);
      44           0 :         return NULL;
      45             :     }
      46         110 :     if (self->view.obj == NULL) {
      47           0 :         PyErr_SetString(PyExc_ValueError,
      48             :                         "operation forbidden on released PickleBuffer object");
      49           0 :         return NULL;
      50             :     }
      51         110 :     return &self->view;
      52             : }
      53             : 
      54             : int
      55           0 : PyPickleBuffer_Release(PyObject *obj)
      56             : {
      57           0 :     PyPickleBufferObject *self = (PyPickleBufferObject *) obj;
      58             : 
      59           0 :     if (!PyPickleBuffer_Check(obj)) {
      60           0 :         PyErr_Format(PyExc_TypeError,
      61             :                      "expected PickleBuffer, %.200s found",
      62           0 :                      Py_TYPE(obj)->tp_name);
      63           0 :         return -1;
      64             :     }
      65           0 :     PyBuffer_Release(&self->view);
      66           0 :     return 0;
      67             : }
      68             : 
      69             : static PyObject *
      70         371 : picklebuf_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
      71             : {
      72             :     PyPickleBufferObject *self;
      73             :     PyObject *base;
      74         371 :     char *keywords[] = {"", NULL};
      75             : 
      76         371 :     if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:PickleBuffer",
      77             :                                      keywords, &base)) {
      78           1 :         return NULL;
      79             :     }
      80             : 
      81         370 :     self = (PyPickleBufferObject *) type->tp_alloc(type, 0);
      82         370 :     if (self == NULL) {
      83           0 :         return NULL;
      84             :     }
      85         370 :     self->view.obj = NULL;
      86         370 :     self->weakreflist = NULL;
      87         370 :     if (PyObject_GetBuffer(base, &self->view, PyBUF_FULL_RO) < 0) {
      88           2 :         Py_DECREF(self);
      89           2 :         return NULL;
      90             :     }
      91         368 :     return (PyObject *) self;
      92             : }
      93             : 
      94             : static int
      95           2 : picklebuf_traverse(PyPickleBufferObject *self, visitproc visit, void *arg)
      96             : {
      97           2 :     Py_VISIT(self->view.obj);
      98           2 :     return 0;
      99             : }
     100             : 
     101             : static int
     102           0 : picklebuf_clear(PyPickleBufferObject *self)
     103             : {
     104           0 :     PyBuffer_Release(&self->view);
     105           0 :     return 0;
     106             : }
     107             : 
     108             : static void
     109         370 : picklebuf_dealloc(PyPickleBufferObject *self)
     110             : {
     111         370 :     PyObject_GC_UnTrack(self);
     112         370 :     if (self->weakreflist != NULL)
     113           0 :         PyObject_ClearWeakRefs((PyObject *) self);
     114         370 :     PyBuffer_Release(&self->view);
     115         370 :     Py_TYPE(self)->tp_free((PyObject *) self);
     116         370 : }
     117             : 
     118             : /* Buffer API */
     119             : 
     120             : static int
     121         317 : picklebuf_getbuf(PyPickleBufferObject *self, Py_buffer *view, int flags)
     122             : {
     123         317 :     if (self->view.obj == NULL) {
     124           1 :         PyErr_SetString(PyExc_ValueError,
     125             :                         "operation forbidden on released PickleBuffer object");
     126           1 :         return -1;
     127             :     }
     128         316 :     return PyObject_GetBuffer(self->view.obj, view, flags);
     129             : }
     130             : 
     131             : static void
     132           0 : picklebuf_releasebuf(PyPickleBufferObject *self, Py_buffer *view)
     133             : {
     134             :     /* Since our bf_getbuffer redirects to the original object, this
     135             :      * implementation is never called.  It only exists to signal that
     136             :      * buffers exported by PickleBuffer have non-trivial releasing
     137             :      * behaviour (see check in Python/getargs.c).
     138             :      */
     139           0 : }
     140             : 
     141             : static PyBufferProcs picklebuf_as_buffer = {
     142             :     .bf_getbuffer = (getbufferproc) picklebuf_getbuf,
     143             :     .bf_releasebuffer = (releasebufferproc) picklebuf_releasebuf,
     144             : };
     145             : 
     146             : /* Methods */
     147             : 
     148             : static PyObject *
     149         293 : picklebuf_raw(PyPickleBufferObject *self, PyObject *Py_UNUSED(ignored))
     150             : {
     151         293 :     if (self->view.obj == NULL) {
     152           1 :         PyErr_SetString(PyExc_ValueError,
     153             :                         "operation forbidden on released PickleBuffer object");
     154           1 :         return NULL;
     155             :     }
     156         292 :     if (self->view.suboffsets != NULL
     157         292 :         || !PyBuffer_IsContiguous(&self->view, 'A')) {
     158           2 :         PyErr_SetString(PyExc_BufferError,
     159             :                         "cannot extract raw buffer from non-contiguous buffer");
     160           2 :         return NULL;
     161             :     }
     162         290 :     PyObject *m = PyMemoryView_FromObject((PyObject *) self);
     163         290 :     if (m == NULL) {
     164           0 :         return NULL;
     165             :     }
     166         290 :     PyMemoryViewObject *mv = (PyMemoryViewObject *) m;
     167         290 :     assert(mv->view.suboffsets == NULL);
     168             :     /* Mutate memoryview instance to make it a "raw" memoryview */
     169         290 :     mv->view.format = "B";
     170         290 :     mv->view.ndim = 1;
     171         290 :     mv->view.itemsize = 1;
     172             :     /* shape = (length,) */
     173         290 :     mv->view.shape = &mv->view.len;
     174             :     /* strides = (1,) */
     175         290 :     mv->view.strides = &mv->view.itemsize;
     176             :     /* Fix memoryview state flags */
     177             :     /* XXX Expose memoryobject.c's init_flags() instead? */
     178         290 :     mv->flags = _Py_MEMORYVIEW_C | _Py_MEMORYVIEW_FORTRAN;
     179         290 :     return m;
     180             : }
     181             : 
     182             : PyDoc_STRVAR(picklebuf_raw_doc,
     183             : "raw($self, /)\n--\n\
     184             : \n\
     185             : Return a memoryview of the raw memory underlying this buffer.\n\
     186             : Will raise BufferError is the buffer isn't contiguous.");
     187             : 
     188             : static PyObject *
     189           3 : picklebuf_release(PyPickleBufferObject *self, PyObject *Py_UNUSED(ignored))
     190             : {
     191           3 :     PyBuffer_Release(&self->view);
     192           3 :     Py_RETURN_NONE;
     193             : }
     194             : 
     195             : PyDoc_STRVAR(picklebuf_release_doc,
     196             : "release($self, /)\n--\n\
     197             : \n\
     198             : Release the underlying buffer exposed by the PickleBuffer object.");
     199             : 
     200             : static PyMethodDef picklebuf_methods[] = {
     201             :     {"raw",     (PyCFunction) picklebuf_raw,     METH_NOARGS, picklebuf_raw_doc},
     202             :     {"release", (PyCFunction) picklebuf_release, METH_NOARGS, picklebuf_release_doc},
     203             :     {NULL,      NULL}
     204             : };
     205             : 
     206             : PyTypeObject PyPickleBuffer_Type = {
     207             :     PyVarObject_HEAD_INIT(NULL, 0)
     208             :     .tp_name = "pickle.PickleBuffer",
     209             :     .tp_doc = PyDoc_STR("Wrapper for potentially out-of-band buffers"),
     210             :     .tp_basicsize = sizeof(PyPickleBufferObject),
     211             :     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
     212             :     .tp_new = picklebuf_new,
     213             :     .tp_dealloc = (destructor) picklebuf_dealloc,
     214             :     .tp_traverse = (traverseproc) picklebuf_traverse,
     215             :     .tp_clear = (inquiry) picklebuf_clear,
     216             :     .tp_weaklistoffset = offsetof(PyPickleBufferObject, weakreflist),
     217             :     .tp_as_buffer = &picklebuf_as_buffer,
     218             :     .tp_methods = picklebuf_methods,
     219             : };