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