1/* PickleBuffer object implementation */
2
3#define PY_SSIZE_T_CLEAN
4#include "Python.h"
5#include <stddef.h>
6
7typedef 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
16PyObject *
17PyPickleBuffer_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
35const Py_buffer *
36PyPickleBuffer_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
54int
55PyPickleBuffer_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
69static PyObject *
70picklebuf_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
94static int
95picklebuf_traverse(PyPickleBufferObject *self, visitproc visit, void *arg)
96{
97 Py_VISIT(self->view.obj);
98 return 0;
99}
100
101static int
102picklebuf_clear(PyPickleBufferObject *self)
103{
104 PyBuffer_Release(&self->view);
105 return 0;
106}
107
108static void
109picklebuf_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
120static int
121picklebuf_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
131static void
132picklebuf_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
141static PyBufferProcs picklebuf_as_buffer = {
142 .bf_getbuffer = (getbufferproc) picklebuf_getbuf,
143 .bf_releasebuffer = (releasebufferproc) picklebuf_releasebuf,
144};
145
146/* Methods */
147
148static PyObject *
149picklebuf_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
182PyDoc_STRVAR(picklebuf_raw_doc,
183"raw($self, /)\n--\n\
184\n\
185Return a memoryview of the raw memory underlying this buffer.\n\
186Will raise BufferError is the buffer isn't contiguous.");
187
188static PyObject *
189picklebuf_release(PyPickleBufferObject *self, PyObject *Py_UNUSED(ignored))
190{
191 PyBuffer_Release(&self->view);
192 Py_RETURN_NONE;
193}
194
195PyDoc_STRVAR(picklebuf_release_doc,
196"release($self, /)\n--\n\
197\n\
198Release the underlying buffer exposed by the PickleBuffer object.");
199
200static 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
206PyTypeObject 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};
220