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 *
PyPickleBuffer_FromObject(PyObject * base)17 PyPickleBuffer_FromObject(PyObject *base)
18 {
19 PyTypeObject *type = &PyPickleBuffer_Type;
20 PyPickleBufferObject *self;
21
22 self = (PyPickleBufferObject *) type->tp_alloc(type, 0);
23 if (self == NULL) {
24 return NULL;
25 }
26 self->view.obj = NULL;
27 self->weakreflist = NULL;
28 if (PyObject_GetBuffer(base, &self->view, PyBUF_FULL_RO) < 0) {
29 Py_DECREF(self);
30 return NULL;
31 }
32 return (PyObject *) self;
33 }
34
35 const Py_buffer *
PyPickleBuffer_GetBuffer(PyObject * obj)36 PyPickleBuffer_GetBuffer(PyObject *obj)
37 {
38 PyPickleBufferObject *self = (PyPickleBufferObject *) obj;
39
40 if (!PyPickleBuffer_Check(obj)) {
41 PyErr_Format(PyExc_TypeError,
42 "expected PickleBuffer, %.200s found",
43 Py_TYPE(obj)->tp_name);
44 return NULL;
45 }
46 if (self->view.obj == NULL) {
47 PyErr_SetString(PyExc_ValueError,
48 "operation forbidden on released PickleBuffer object");
49 return NULL;
50 }
51 return &self->view;
52 }
53
54 int
PyPickleBuffer_Release(PyObject * obj)55 PyPickleBuffer_Release(PyObject *obj)
56 {
57 PyPickleBufferObject *self = (PyPickleBufferObject *) obj;
58
59 if (!PyPickleBuffer_Check(obj)) {
60 PyErr_Format(PyExc_TypeError,
61 "expected PickleBuffer, %.200s found",
62 Py_TYPE(obj)->tp_name);
63 return -1;
64 }
65 PyBuffer_Release(&self->view);
66 return 0;
67 }
68
69 static PyObject *
picklebuf_new(PyTypeObject * type,PyObject * args,PyObject * kwds)70 picklebuf_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
71 {
72 PyPickleBufferObject *self;
73 PyObject *base;
74 char *keywords[] = {"", NULL};
75
76 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:PickleBuffer",
77 keywords, &base)) {
78 return NULL;
79 }
80
81 self = (PyPickleBufferObject *) type->tp_alloc(type, 0);
82 if (self == NULL) {
83 return NULL;
84 }
85 self->view.obj = NULL;
86 self->weakreflist = NULL;
87 if (PyObject_GetBuffer(base, &self->view, PyBUF_FULL_RO) < 0) {
88 Py_DECREF(self);
89 return NULL;
90 }
91 return (PyObject *) self;
92 }
93
94 static int
picklebuf_traverse(PyPickleBufferObject * self,visitproc visit,void * arg)95 picklebuf_traverse(PyPickleBufferObject *self, visitproc visit, void *arg)
96 {
97 Py_VISIT(self->view.obj);
98 return 0;
99 }
100
101 static int
picklebuf_clear(PyPickleBufferObject * self)102 picklebuf_clear(PyPickleBufferObject *self)
103 {
104 PyBuffer_Release(&self->view);
105 return 0;
106 }
107
108 static void
picklebuf_dealloc(PyPickleBufferObject * self)109 picklebuf_dealloc(PyPickleBufferObject *self)
110 {
111 PyObject_GC_UnTrack(self);
112 if (self->weakreflist != NULL)
113 PyObject_ClearWeakRefs((PyObject *) self);
114 PyBuffer_Release(&self->view);
115 Py_TYPE(self)->tp_free((PyObject *) self);
116 }
117
118 /* Buffer API */
119
120 static int
picklebuf_getbuf(PyPickleBufferObject * self,Py_buffer * view,int flags)121 picklebuf_getbuf(PyPickleBufferObject *self, Py_buffer *view, int flags)
122 {
123 if (self->view.obj == NULL) {
124 PyErr_SetString(PyExc_ValueError,
125 "operation forbidden on released PickleBuffer object");
126 return -1;
127 }
128 return PyObject_GetBuffer(self->view.obj, view, flags);
129 }
130
131 static void
picklebuf_releasebuf(PyPickleBufferObject * self,Py_buffer * view)132 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 }
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 *
picklebuf_raw(PyPickleBufferObject * self,PyObject * Py_UNUSED (ignored))149 picklebuf_raw(PyPickleBufferObject *self, PyObject *Py_UNUSED(ignored))
150 {
151 if (self->view.obj == NULL) {
152 PyErr_SetString(PyExc_ValueError,
153 "operation forbidden on released PickleBuffer object");
154 return NULL;
155 }
156 if (self->view.suboffsets != NULL
157 || !PyBuffer_IsContiguous(&self->view, 'A')) {
158 PyErr_SetString(PyExc_BufferError,
159 "cannot extract raw buffer from non-contiguous buffer");
160 return NULL;
161 }
162 PyObject *m = PyMemoryView_FromObject((PyObject *) self);
163 if (m == NULL) {
164 return NULL;
165 }
166 PyMemoryViewObject *mv = (PyMemoryViewObject *) m;
167 assert(mv->view.suboffsets == NULL);
168 /* Mutate memoryview instance to make it a "raw" memoryview */
169 mv->view.format = "B";
170 mv->view.ndim = 1;
171 mv->view.itemsize = 1;
172 /* shape = (length,) */
173 mv->view.shape = &mv->view.len;
174 /* strides = (1,) */
175 mv->view.strides = &mv->view.itemsize;
176 /* Fix memoryview state flags */
177 /* XXX Expose memoryobject.c's init_flags() instead? */
178 mv->flags = _Py_MEMORYVIEW_C | _Py_MEMORYVIEW_FORTRAN;
179 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 *
picklebuf_release(PyPickleBufferObject * self,PyObject * Py_UNUSED (ignored))189 picklebuf_release(PyPickleBufferObject *self, PyObject *Py_UNUSED(ignored))
190 {
191 PyBuffer_Release(&self->view);
192 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 = "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 };
220