Branch data     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                 :            :     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                 :            : };