blob: 7a234f1dfff5590dd129d39191bfbe74cbc4e96b [file] [log] [blame]
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001#include "Python.h"
2#include "structmember.h"
3
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08004PyDoc_STRVAR(pickle_module_doc,
5"Optimized C implementation for the Python pickle module.");
6
Larry Hastings61272b72014-01-07 12:41:53 -08007/*[clinic input]
Serhiy Storchaka3c1f0f12014-01-27 10:34:22 +02008output preset file
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08009module _pickle
Larry Hastingsc2047262014-01-25 20:43:29 -080010class _pickle.Pickler "PicklerObject *" "&Pickler_Type"
11class _pickle.PicklerMemoProxy "PicklerMemoProxyObject *" "&PicklerMemoProxyType"
12class _pickle.Unpickler "UnpicklerObject *" "&Unpickler_Type"
13class _pickle.UnpicklerMemoProxy "UnpicklerMemoProxyObject *" "&UnpicklerMemoProxyType"
Larry Hastings61272b72014-01-07 12:41:53 -080014[clinic start generated code]*/
Larry Hastings581ee362014-01-28 05:00:08 -080015/*[clinic end generated code: output=da39a3ee5e6b4b0d input=11c45248a41dd3fc]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -080016
Alexandre Vassalottica2d6102008-06-12 18:26:05 +000017/* Bump this when new opcodes are added to the pickle protocol. */
18enum {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +010019 HIGHEST_PROTOCOL = 4,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +000020 DEFAULT_PROTOCOL = 3
21};
22
Alexandre Vassalottica2d6102008-06-12 18:26:05 +000023/* Pickle opcodes. These must be kept updated with pickle.py.
24 Extensive docs are in pickletools.py. */
25enum opcode {
26 MARK = '(',
27 STOP = '.',
28 POP = '0',
29 POP_MARK = '1',
30 DUP = '2',
31 FLOAT = 'F',
32 INT = 'I',
33 BININT = 'J',
34 BININT1 = 'K',
35 LONG = 'L',
36 BININT2 = 'M',
37 NONE = 'N',
38 PERSID = 'P',
39 BINPERSID = 'Q',
40 REDUCE = 'R',
41 STRING = 'S',
42 BINSTRING = 'T',
43 SHORT_BINSTRING = 'U',
44 UNICODE = 'V',
45 BINUNICODE = 'X',
46 APPEND = 'a',
47 BUILD = 'b',
48 GLOBAL = 'c',
49 DICT = 'd',
50 EMPTY_DICT = '}',
51 APPENDS = 'e',
52 GET = 'g',
53 BINGET = 'h',
54 INST = 'i',
55 LONG_BINGET = 'j',
56 LIST = 'l',
57 EMPTY_LIST = ']',
58 OBJ = 'o',
59 PUT = 'p',
60 BINPUT = 'q',
61 LONG_BINPUT = 'r',
62 SETITEM = 's',
63 TUPLE = 't',
64 EMPTY_TUPLE = ')',
65 SETITEMS = 'u',
66 BINFLOAT = 'G',
67
68 /* Protocol 2. */
69 PROTO = '\x80',
70 NEWOBJ = '\x81',
71 EXT1 = '\x82',
72 EXT2 = '\x83',
73 EXT4 = '\x84',
74 TUPLE1 = '\x85',
75 TUPLE2 = '\x86',
76 TUPLE3 = '\x87',
77 NEWTRUE = '\x88',
78 NEWFALSE = '\x89',
79 LONG1 = '\x8a',
80 LONG4 = '\x8b',
81
82 /* Protocol 3 (Python 3.x) */
83 BINBYTES = 'B',
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +010084 SHORT_BINBYTES = 'C',
85
86 /* Protocol 4 */
87 SHORT_BINUNICODE = '\x8c',
88 BINUNICODE8 = '\x8d',
89 BINBYTES8 = '\x8e',
90 EMPTY_SET = '\x8f',
91 ADDITEMS = '\x90',
92 FROZENSET = '\x91',
93 NEWOBJ_EX = '\x92',
94 STACK_GLOBAL = '\x93',
95 MEMOIZE = '\x94',
96 FRAME = '\x95'
Alexandre Vassalottica2d6102008-06-12 18:26:05 +000097};
98
Alexandre Vassalottica2d6102008-06-12 18:26:05 +000099enum {
100 /* Keep in synch with pickle.Pickler._BATCHSIZE. This is how many elements
101 batch_list/dict() pumps out before doing APPENDS/SETITEMS. Nothing will
102 break if this gets out of synch with pickle.py, but it's unclear that would
103 help anything either. */
104 BATCHSIZE = 1000,
105
106 /* Nesting limit until Pickler, when running in "fast mode", starts
107 checking for self-referential data-structures. */
108 FAST_NESTING_LIMIT = 50,
109
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000110 /* Initial size of the write buffer of Pickler. */
111 WRITE_BUF_SIZE = 4096,
112
Antoine Pitrou04248a82010-10-12 20:51:21 +0000113 /* Prefetch size when unpickling (disabled on unpeekable streams) */
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100114 PREFETCH = 8192 * 16,
115
116 FRAME_SIZE_TARGET = 64 * 1024,
117
118 FRAME_HEADER_SIZE = 9
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000119};
120
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800121/*************************************************************************/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000122
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800123/* State of the pickle module, per PEP 3121. */
124typedef struct {
125 /* Exception classes for pickle. */
126 PyObject *PickleError;
127 PyObject *PicklingError;
128 PyObject *UnpicklingError;
Larry Hastings61272b72014-01-07 12:41:53 -0800129
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800130 /* copyreg.dispatch_table, {type_object: pickling_function} */
131 PyObject *dispatch_table;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000132
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800133 /* For the extension opcodes EXT1, EXT2 and EXT4. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000134
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800135 /* copyreg._extension_registry, {(module_name, function_name): code} */
136 PyObject *extension_registry;
137 /* copyreg._extension_cache, {code: object} */
138 PyObject *extension_cache;
139 /* copyreg._inverted_registry, {code: (module_name, function_name)} */
140 PyObject *inverted_registry;
141
142 /* Import mappings for compatibility with Python 2.x */
143
144 /* _compat_pickle.NAME_MAPPING,
145 {(oldmodule, oldname): (newmodule, newname)} */
146 PyObject *name_mapping_2to3;
147 /* _compat_pickle.IMPORT_MAPPING, {oldmodule: newmodule} */
148 PyObject *import_mapping_2to3;
149 /* Same, but with REVERSE_NAME_MAPPING / REVERSE_IMPORT_MAPPING */
150 PyObject *name_mapping_3to2;
151 PyObject *import_mapping_3to2;
152
153 /* codecs.encode, used for saving bytes in older protocols */
154 PyObject *codecs_encode;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800155} PickleState;
156
157/* Forward declaration of the _pickle module definition. */
158static struct PyModuleDef _picklemodule;
159
160/* Given a module object, get its per-module state. */
161static PickleState *
162_Pickle_GetState(PyObject *module)
163{
164 return (PickleState *)PyModule_GetState(module);
165}
166
167/* Find the module instance imported in the currently running sub-interpreter
168 and get its state. */
169static PickleState *
170_Pickle_GetGlobalState(void)
171{
172 return _Pickle_GetState(PyState_FindModule(&_picklemodule));
173}
174
175/* Clear the given pickle module state. */
176static void
177_Pickle_ClearState(PickleState *st)
178{
179 Py_CLEAR(st->PickleError);
180 Py_CLEAR(st->PicklingError);
181 Py_CLEAR(st->UnpicklingError);
182 Py_CLEAR(st->dispatch_table);
183 Py_CLEAR(st->extension_registry);
184 Py_CLEAR(st->extension_cache);
185 Py_CLEAR(st->inverted_registry);
186 Py_CLEAR(st->name_mapping_2to3);
187 Py_CLEAR(st->import_mapping_2to3);
188 Py_CLEAR(st->name_mapping_3to2);
189 Py_CLEAR(st->import_mapping_3to2);
190 Py_CLEAR(st->codecs_encode);
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800191}
192
193/* Initialize the given pickle module state. */
194static int
195_Pickle_InitState(PickleState *st)
196{
197 PyObject *copyreg = NULL;
198 PyObject *compat_pickle = NULL;
199 PyObject *codecs = NULL;
200
201 copyreg = PyImport_ImportModule("copyreg");
202 if (!copyreg)
203 goto error;
204 st->dispatch_table = PyObject_GetAttrString(copyreg, "dispatch_table");
205 if (!st->dispatch_table)
206 goto error;
207 if (!PyDict_CheckExact(st->dispatch_table)) {
208 PyErr_Format(PyExc_RuntimeError,
209 "copyreg.dispatch_table should be a dict, not %.200s",
210 Py_TYPE(st->dispatch_table)->tp_name);
211 goto error;
212 }
213 st->extension_registry = \
214 PyObject_GetAttrString(copyreg, "_extension_registry");
215 if (!st->extension_registry)
216 goto error;
217 if (!PyDict_CheckExact(st->extension_registry)) {
218 PyErr_Format(PyExc_RuntimeError,
219 "copyreg._extension_registry should be a dict, "
220 "not %.200s", Py_TYPE(st->extension_registry)->tp_name);
221 goto error;
222 }
223 st->inverted_registry = \
224 PyObject_GetAttrString(copyreg, "_inverted_registry");
225 if (!st->inverted_registry)
226 goto error;
227 if (!PyDict_CheckExact(st->inverted_registry)) {
228 PyErr_Format(PyExc_RuntimeError,
229 "copyreg._inverted_registry should be a dict, "
230 "not %.200s", Py_TYPE(st->inverted_registry)->tp_name);
231 goto error;
232 }
233 st->extension_cache = PyObject_GetAttrString(copyreg, "_extension_cache");
234 if (!st->extension_cache)
235 goto error;
236 if (!PyDict_CheckExact(st->extension_cache)) {
237 PyErr_Format(PyExc_RuntimeError,
238 "copyreg._extension_cache should be a dict, "
239 "not %.200s", Py_TYPE(st->extension_cache)->tp_name);
240 goto error;
241 }
242 Py_CLEAR(copyreg);
243
244 /* Load the 2.x -> 3.x stdlib module mapping tables */
245 compat_pickle = PyImport_ImportModule("_compat_pickle");
246 if (!compat_pickle)
247 goto error;
248 st->name_mapping_2to3 = \
249 PyObject_GetAttrString(compat_pickle, "NAME_MAPPING");
250 if (!st->name_mapping_2to3)
251 goto error;
252 if (!PyDict_CheckExact(st->name_mapping_2to3)) {
253 PyErr_Format(PyExc_RuntimeError,
254 "_compat_pickle.NAME_MAPPING should be a dict, not %.200s",
255 Py_TYPE(st->name_mapping_2to3)->tp_name);
256 goto error;
257 }
258 st->import_mapping_2to3 = \
259 PyObject_GetAttrString(compat_pickle, "IMPORT_MAPPING");
260 if (!st->import_mapping_2to3)
261 goto error;
262 if (!PyDict_CheckExact(st->import_mapping_2to3)) {
263 PyErr_Format(PyExc_RuntimeError,
264 "_compat_pickle.IMPORT_MAPPING should be a dict, "
265 "not %.200s", Py_TYPE(st->import_mapping_2to3)->tp_name);
266 goto error;
267 }
268 /* ... and the 3.x -> 2.x mapping tables */
269 st->name_mapping_3to2 = \
270 PyObject_GetAttrString(compat_pickle, "REVERSE_NAME_MAPPING");
271 if (!st->name_mapping_3to2)
272 goto error;
273 if (!PyDict_CheckExact(st->name_mapping_3to2)) {
274 PyErr_Format(PyExc_RuntimeError,
275 "_compat_pickle.REVERSE_NAME_MAPPING should be a dict, "
276 "not %.200s", Py_TYPE(st->name_mapping_3to2)->tp_name);
277 goto error;
278 }
279 st->import_mapping_3to2 = \
280 PyObject_GetAttrString(compat_pickle, "REVERSE_IMPORT_MAPPING");
281 if (!st->import_mapping_3to2)
282 goto error;
283 if (!PyDict_CheckExact(st->import_mapping_3to2)) {
284 PyErr_Format(PyExc_RuntimeError,
285 "_compat_pickle.REVERSE_IMPORT_MAPPING should be a dict, "
286 "not %.200s", Py_TYPE(st->import_mapping_3to2)->tp_name);
287 goto error;
288 }
289 Py_CLEAR(compat_pickle);
290
291 codecs = PyImport_ImportModule("codecs");
292 if (codecs == NULL)
293 goto error;
294 st->codecs_encode = PyObject_GetAttrString(codecs, "encode");
295 if (st->codecs_encode == NULL) {
296 goto error;
297 }
298 if (!PyCallable_Check(st->codecs_encode)) {
299 PyErr_Format(PyExc_RuntimeError,
300 "codecs.encode should be a callable, not %.200s",
301 Py_TYPE(st->codecs_encode)->tp_name);
302 goto error;
303 }
304 Py_CLEAR(codecs);
305
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800306 return 0;
307
308 error:
309 Py_CLEAR(copyreg);
310 Py_CLEAR(compat_pickle);
311 Py_CLEAR(codecs);
312 _Pickle_ClearState(st);
313 return -1;
314}
315
316/* Helper for calling a function with a single argument quickly.
317
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800318 This function steals the reference of the given argument. */
319static PyObject *
320_Pickle_FastCall(PyObject *func, PyObject *obj)
321{
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800322 PyObject *result;
Alexandre Vassalottib13e6bc2013-11-28 14:56:09 -0800323 PyObject *arg_tuple = PyTuple_New(1);
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800324
Alexandre Vassalottib13e6bc2013-11-28 14:56:09 -0800325 /* Note: this function used to reuse the argument tuple. This used to give
326 a slight performance boost with older pickle implementations where many
327 unbuffered reads occurred (thus needing many function calls).
328
329 However, this optimization was removed because it was too complicated
330 to get right. It abused the C API for tuples to mutate them which led
331 to subtle reference counting and concurrency bugs. Furthermore, the
332 introduction of protocol 4 and the prefetching optimization via peek()
333 significantly reduced the number of function calls we do. Thus, the
334 benefits became marginal at best. */
335
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800336 if (arg_tuple == NULL) {
Alexandre Vassalottib13e6bc2013-11-28 14:56:09 -0800337 Py_DECREF(obj);
338 return NULL;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800339 }
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800340 PyTuple_SET_ITEM(arg_tuple, 0, obj);
341 result = PyObject_Call(func, arg_tuple, NULL);
Alexandre Vassalottib13e6bc2013-11-28 14:56:09 -0800342 Py_CLEAR(arg_tuple);
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800343 return result;
344}
345
346/*************************************************************************/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000347
348static int
349stack_underflow(void)
350{
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800351 PickleState *st = _Pickle_GetGlobalState();
352 PyErr_SetString(st->UnpicklingError, "unpickling stack underflow");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000353 return -1;
354}
355
356/* Internal data type used as the unpickling stack. */
357typedef struct {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000358 PyObject_VAR_HEAD
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000359 PyObject **data;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000360 Py_ssize_t allocated; /* number of slots in data allocated */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000361} Pdata;
362
363static void
364Pdata_dealloc(Pdata *self)
365{
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200366 Py_ssize_t i = Py_SIZE(self);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000367 while (--i >= 0) {
368 Py_DECREF(self->data[i]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000369 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000370 PyMem_FREE(self->data);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000371 PyObject_Del(self);
372}
373
374static PyTypeObject Pdata_Type = {
375 PyVarObject_HEAD_INIT(NULL, 0)
376 "_pickle.Pdata", /*tp_name*/
377 sizeof(Pdata), /*tp_basicsize*/
Serhiy Storchaka5bbd2312014-12-16 19:39:08 +0200378 sizeof(PyObject *), /*tp_itemsize*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000379 (destructor)Pdata_dealloc, /*tp_dealloc*/
380};
381
382static PyObject *
383Pdata_New(void)
384{
385 Pdata *self;
386
387 if (!(self = PyObject_New(Pdata, &Pdata_Type)))
388 return NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000389 Py_SIZE(self) = 0;
390 self->allocated = 8;
391 self->data = PyMem_MALLOC(self->allocated * sizeof(PyObject *));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000392 if (self->data)
393 return (PyObject *)self;
394 Py_DECREF(self);
395 return PyErr_NoMemory();
396}
397
398
399/* Retain only the initial clearto items. If clearto >= the current
400 * number of items, this is a (non-erroneous) NOP.
401 */
402static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200403Pdata_clear(Pdata *self, Py_ssize_t clearto)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000404{
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200405 Py_ssize_t i = Py_SIZE(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000406
407 if (clearto < 0)
408 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000409 if (clearto >= i)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000410 return 0;
411
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000412 while (--i >= clearto) {
413 Py_CLEAR(self->data[i]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000414 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000415 Py_SIZE(self) = clearto;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000416 return 0;
417}
418
419static int
420Pdata_grow(Pdata *self)
421{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000422 PyObject **data = self->data;
Victor Stinnerf13c46c2014-08-17 21:05:55 +0200423 size_t allocated = (size_t)self->allocated;
424 size_t new_allocated;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000425
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000426 new_allocated = (allocated >> 3) + 6;
427 /* check for integer overflow */
Victor Stinnerf13c46c2014-08-17 21:05:55 +0200428 if (new_allocated > (size_t)PY_SSIZE_T_MAX - allocated)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000429 goto nomemory;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000430 new_allocated += allocated;
Victor Stinnerf13c46c2014-08-17 21:05:55 +0200431 if (new_allocated > ((size_t)PY_SSIZE_T_MAX / sizeof(PyObject *)))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000432 goto nomemory;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000433 data = PyMem_REALLOC(data, new_allocated * sizeof(PyObject *));
434 if (data == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000435 goto nomemory;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000436
437 self->data = data;
Victor Stinnerf13c46c2014-08-17 21:05:55 +0200438 self->allocated = (Py_ssize_t)new_allocated;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000439 return 0;
440
441 nomemory:
442 PyErr_NoMemory();
443 return -1;
444}
445
446/* D is a Pdata*. Pop the topmost element and store it into V, which
447 * must be an lvalue holding PyObject*. On stack underflow, UnpicklingError
448 * is raised and V is set to NULL.
449 */
450static PyObject *
451Pdata_pop(Pdata *self)
452{
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800453 PickleState *st = _Pickle_GetGlobalState();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000454 if (Py_SIZE(self) == 0) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -0800455 PyErr_SetString(st->UnpicklingError, "bad pickle data");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000456 return NULL;
457 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000458 return self->data[--Py_SIZE(self)];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000459}
460#define PDATA_POP(D, V) do { (V) = Pdata_pop((D)); } while (0)
461
462static int
463Pdata_push(Pdata *self, PyObject *obj)
464{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000465 if (Py_SIZE(self) == self->allocated && Pdata_grow(self) < 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000466 return -1;
467 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000468 self->data[Py_SIZE(self)++] = obj;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000469 return 0;
470}
471
472/* Push an object on stack, transferring its ownership to the stack. */
473#define PDATA_PUSH(D, O, ER) do { \
474 if (Pdata_push((D), (O)) < 0) return (ER); } while(0)
475
476/* Push an object on stack, adding a new reference to the object. */
477#define PDATA_APPEND(D, O, ER) do { \
478 Py_INCREF((O)); \
479 if (Pdata_push((D), (O)) < 0) return (ER); } while(0)
480
481static PyObject *
482Pdata_poptuple(Pdata *self, Py_ssize_t start)
483{
484 PyObject *tuple;
485 Py_ssize_t len, i, j;
486
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000487 len = Py_SIZE(self) - start;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000488 tuple = PyTuple_New(len);
489 if (tuple == NULL)
490 return NULL;
491 for (i = start, j = 0; j < len; i++, j++)
492 PyTuple_SET_ITEM(tuple, j, self->data[i]);
493
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000494 Py_SIZE(self) = start;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000495 return tuple;
496}
497
498static PyObject *
499Pdata_poplist(Pdata *self, Py_ssize_t start)
500{
501 PyObject *list;
502 Py_ssize_t len, i, j;
503
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000504 len = Py_SIZE(self) - start;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000505 list = PyList_New(len);
506 if (list == NULL)
507 return NULL;
508 for (i = start, j = 0; j < len; i++, j++)
509 PyList_SET_ITEM(list, j, self->data[i]);
510
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000511 Py_SIZE(self) = start;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000512 return list;
513}
514
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000515typedef struct {
516 PyObject *me_key;
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200517 Py_ssize_t me_value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000518} PyMemoEntry;
519
520typedef struct {
521 Py_ssize_t mt_mask;
522 Py_ssize_t mt_used;
523 Py_ssize_t mt_allocated;
524 PyMemoEntry *mt_table;
525} PyMemoTable;
526
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000527typedef struct PicklerObject {
528 PyObject_HEAD
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000529 PyMemoTable *memo; /* Memo table, keep track of the seen
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000530 objects to support self-referential objects
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000531 pickling. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000532 PyObject *pers_func; /* persistent_id() method, can be NULL */
Antoine Pitrou8d3c2902012-03-04 18:31:48 +0100533 PyObject *dispatch_table; /* private dispatch_table, can be NULL */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000534
535 PyObject *write; /* write() method of the output stream. */
536 PyObject *output_buffer; /* Write into a local bytearray buffer before
537 flushing to the stream. */
538 Py_ssize_t output_len; /* Length of output_buffer. */
539 Py_ssize_t max_output_len; /* Allocation size of output_buffer. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000540 int proto; /* Pickle protocol number, >= 0 */
541 int bin; /* Boolean, true if proto > 0 */
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100542 int framing; /* True when framing is enabled, proto >= 4 */
543 Py_ssize_t frame_start; /* Position in output_buffer where the
544 where the current frame begins. -1 if there
545 is no frame currently open. */
546
547 Py_ssize_t buf_size; /* Size of the current buffered pickle data */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000548 int fast; /* Enable fast mode if set to a true value.
549 The fast mode disable the usage of memo,
550 therefore speeding the pickling process by
551 not generating superfluous PUT opcodes. It
552 should not be used if with self-referential
553 objects. */
554 int fast_nesting;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000555 int fix_imports; /* Indicate whether Pickler should fix
556 the name of globals for Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000557 PyObject *fast_memo;
558} PicklerObject;
559
560typedef struct UnpicklerObject {
561 PyObject_HEAD
562 Pdata *stack; /* Pickle data stack, store unpickled objects. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000563
564 /* The unpickler memo is just an array of PyObject *s. Using a dict
565 is unnecessary, since the keys are contiguous ints. */
566 PyObject **memo;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100567 Py_ssize_t memo_size; /* Capacity of the memo array */
568 Py_ssize_t memo_len; /* Number of objects in the memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000569
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000570 PyObject *pers_func; /* persistent_load() method, can be NULL. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000571
572 Py_buffer buffer;
573 char *input_buffer;
574 char *input_line;
575 Py_ssize_t input_len;
576 Py_ssize_t next_read_idx;
Antoine Pitrou04248a82010-10-12 20:51:21 +0000577 Py_ssize_t prefetched_idx; /* index of first prefetched byte */
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100578
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000579 PyObject *read; /* read() method of the input stream. */
580 PyObject *readline; /* readline() method of the input stream. */
Antoine Pitrou04248a82010-10-12 20:51:21 +0000581 PyObject *peek; /* peek() method of the input stream, or NULL */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000582
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000583 char *encoding; /* Name of the encoding to be used for
584 decoding strings pickled using Python
585 2.x. The default value is "ASCII" */
586 char *errors; /* Name of errors handling scheme to used when
587 decoding strings. The default value is
588 "strict". */
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -0500589 Py_ssize_t *marks; /* Mark stack, used for unpickling container
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000590 objects. */
591 Py_ssize_t num_marks; /* Number of marks in the mark stack. */
592 Py_ssize_t marks_size; /* Current allocated size of the mark stack. */
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000593 int proto; /* Protocol of the pickle loaded. */
594 int fix_imports; /* Indicate whether Unpickler should fix
595 the name of globals pickled by Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000596} UnpicklerObject;
597
Serhiy Storchaka3c1f0f12014-01-27 10:34:22 +0200598typedef struct {
599 PyObject_HEAD
600 PicklerObject *pickler; /* Pickler whose memo table we're proxying. */
601} PicklerMemoProxyObject;
602
603typedef struct {
604 PyObject_HEAD
605 UnpicklerObject *unpickler;
606} UnpicklerMemoProxyObject;
607
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000608/* Forward declarations */
609static int save(PicklerObject *, PyObject *, int);
610static int save_reduce(PicklerObject *, PyObject *, PyObject *);
611static PyTypeObject Pickler_Type;
612static PyTypeObject Unpickler_Type;
613
Serhiy Storchaka3c1f0f12014-01-27 10:34:22 +0200614#include "clinic/_pickle.c.h"
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000615
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000616/*************************************************************************
Serhiy Storchaka95949422013-08-27 19:40:23 +0300617 A custom hashtable mapping void* to Python ints. This is used by the pickler
618 for memoization. Using a custom hashtable rather than PyDict allows us to skip
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000619 a bunch of unnecessary object creation. This makes a huge performance
620 difference. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000621
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000622#define MT_MINSIZE 8
623#define PERTURB_SHIFT 5
624
625
626static PyMemoTable *
627PyMemoTable_New(void)
628{
629 PyMemoTable *memo = PyMem_MALLOC(sizeof(PyMemoTable));
630 if (memo == NULL) {
631 PyErr_NoMemory();
632 return NULL;
633 }
634
635 memo->mt_used = 0;
636 memo->mt_allocated = MT_MINSIZE;
637 memo->mt_mask = MT_MINSIZE - 1;
638 memo->mt_table = PyMem_MALLOC(MT_MINSIZE * sizeof(PyMemoEntry));
639 if (memo->mt_table == NULL) {
640 PyMem_FREE(memo);
641 PyErr_NoMemory();
642 return NULL;
643 }
644 memset(memo->mt_table, 0, MT_MINSIZE * sizeof(PyMemoEntry));
645
646 return memo;
647}
648
649static PyMemoTable *
650PyMemoTable_Copy(PyMemoTable *self)
651{
652 Py_ssize_t i;
653 PyMemoTable *new = PyMemoTable_New();
654 if (new == NULL)
655 return NULL;
656
657 new->mt_used = self->mt_used;
658 new->mt_allocated = self->mt_allocated;
659 new->mt_mask = self->mt_mask;
660 /* The table we get from _New() is probably smaller than we wanted.
661 Free it and allocate one that's the right size. */
662 PyMem_FREE(new->mt_table);
663 new->mt_table = PyMem_MALLOC(self->mt_allocated * sizeof(PyMemoEntry));
664 if (new->mt_table == NULL) {
665 PyMem_FREE(new);
Victor Stinner42024562013-07-12 00:53:57 +0200666 PyErr_NoMemory();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000667 return NULL;
668 }
669 for (i = 0; i < self->mt_allocated; i++) {
670 Py_XINCREF(self->mt_table[i].me_key);
671 }
672 memcpy(new->mt_table, self->mt_table,
673 sizeof(PyMemoEntry) * self->mt_allocated);
674
675 return new;
676}
677
678static Py_ssize_t
679PyMemoTable_Size(PyMemoTable *self)
680{
681 return self->mt_used;
682}
683
684static int
685PyMemoTable_Clear(PyMemoTable *self)
686{
687 Py_ssize_t i = self->mt_allocated;
688
689 while (--i >= 0) {
690 Py_XDECREF(self->mt_table[i].me_key);
691 }
692 self->mt_used = 0;
693 memset(self->mt_table, 0, self->mt_allocated * sizeof(PyMemoEntry));
694 return 0;
695}
696
697static void
698PyMemoTable_Del(PyMemoTable *self)
699{
700 if (self == NULL)
701 return;
702 PyMemoTable_Clear(self);
703
704 PyMem_FREE(self->mt_table);
705 PyMem_FREE(self);
706}
707
708/* Since entries cannot be deleted from this hashtable, _PyMemoTable_Lookup()
709 can be considerably simpler than dictobject.c's lookdict(). */
710static PyMemoEntry *
711_PyMemoTable_Lookup(PyMemoTable *self, PyObject *key)
712{
713 size_t i;
714 size_t perturb;
715 size_t mask = (size_t)self->mt_mask;
716 PyMemoEntry *table = self->mt_table;
717 PyMemoEntry *entry;
Benjamin Peterson8f67d082010-10-17 20:54:53 +0000718 Py_hash_t hash = (Py_hash_t)key >> 3;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000719
720 i = hash & mask;
721 entry = &table[i];
722 if (entry->me_key == NULL || entry->me_key == key)
723 return entry;
724
725 for (perturb = hash; ; perturb >>= PERTURB_SHIFT) {
726 i = (i << 2) + i + perturb + 1;
727 entry = &table[i & mask];
728 if (entry->me_key == NULL || entry->me_key == key)
729 return entry;
730 }
731 assert(0); /* Never reached */
732 return NULL;
733}
734
735/* Returns -1 on failure, 0 on success. */
736static int
737_PyMemoTable_ResizeTable(PyMemoTable *self, Py_ssize_t min_size)
738{
739 PyMemoEntry *oldtable = NULL;
740 PyMemoEntry *oldentry, *newentry;
741 Py_ssize_t new_size = MT_MINSIZE;
742 Py_ssize_t to_process;
743
744 assert(min_size > 0);
745
746 /* Find the smallest valid table size >= min_size. */
747 while (new_size < min_size && new_size > 0)
748 new_size <<= 1;
749 if (new_size <= 0) {
750 PyErr_NoMemory();
751 return -1;
752 }
753 /* new_size needs to be a power of two. */
754 assert((new_size & (new_size - 1)) == 0);
755
756 /* Allocate new table. */
757 oldtable = self->mt_table;
758 self->mt_table = PyMem_MALLOC(new_size * sizeof(PyMemoEntry));
759 if (self->mt_table == NULL) {
Victor Stinner8ca72e22013-07-12 00:53:26 +0200760 self->mt_table = oldtable;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000761 PyErr_NoMemory();
762 return -1;
763 }
764 self->mt_allocated = new_size;
765 self->mt_mask = new_size - 1;
766 memset(self->mt_table, 0, sizeof(PyMemoEntry) * new_size);
767
768 /* Copy entries from the old table. */
769 to_process = self->mt_used;
770 for (oldentry = oldtable; to_process > 0; oldentry++) {
771 if (oldentry->me_key != NULL) {
772 to_process--;
773 /* newentry is a pointer to a chunk of the new
774 mt_table, so we're setting the key:value pair
775 in-place. */
776 newentry = _PyMemoTable_Lookup(self, oldentry->me_key);
777 newentry->me_key = oldentry->me_key;
778 newentry->me_value = oldentry->me_value;
779 }
780 }
781
782 /* Deallocate the old table. */
783 PyMem_FREE(oldtable);
784 return 0;
785}
786
787/* Returns NULL on failure, a pointer to the value otherwise. */
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200788static Py_ssize_t *
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000789PyMemoTable_Get(PyMemoTable *self, PyObject *key)
790{
791 PyMemoEntry *entry = _PyMemoTable_Lookup(self, key);
792 if (entry->me_key == NULL)
793 return NULL;
794 return &entry->me_value;
795}
796
797/* Returns -1 on failure, 0 on success. */
798static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200799PyMemoTable_Set(PyMemoTable *self, PyObject *key, Py_ssize_t value)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000800{
801 PyMemoEntry *entry;
802
803 assert(key != NULL);
804
805 entry = _PyMemoTable_Lookup(self, key);
806 if (entry->me_key != NULL) {
807 entry->me_value = value;
808 return 0;
809 }
810 Py_INCREF(key);
811 entry->me_key = key;
812 entry->me_value = value;
813 self->mt_used++;
814
815 /* If we added a key, we can safely resize. Otherwise just return!
816 * If used >= 2/3 size, adjust size. Normally, this quaduples the size.
817 *
818 * Quadrupling the size improves average table sparseness
819 * (reducing collisions) at the cost of some memory. It also halves
820 * the number of expensive resize operations in a growing memo table.
821 *
822 * Very large memo tables (over 50K items) use doubling instead.
823 * This may help applications with severe memory constraints.
824 */
825 if (!(self->mt_used * 3 >= (self->mt_mask + 1) * 2))
826 return 0;
827 return _PyMemoTable_ResizeTable(self,
828 (self->mt_used > 50000 ? 2 : 4) * self->mt_used);
829}
830
831#undef MT_MINSIZE
832#undef PERTURB_SHIFT
833
834/*************************************************************************/
835
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000836
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000837static int
838_Pickler_ClearBuffer(PicklerObject *self)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000839{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000840 Py_CLEAR(self->output_buffer);
841 self->output_buffer =
842 PyBytes_FromStringAndSize(NULL, self->max_output_len);
843 if (self->output_buffer == NULL)
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +0000844 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000845 self->output_len = 0;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100846 self->frame_start = -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000847 return 0;
848}
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +0000849
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100850static void
Antoine Pitrou8f2ee6e2013-11-23 21:05:08 +0100851_write_size64(char *out, size_t value)
852{
Victor Stinnerf13c46c2014-08-17 21:05:55 +0200853 size_t i;
Alexandre Vassalotti1048fb52013-11-25 11:35:46 -0800854
855 assert(sizeof(size_t) <= 8);
856
857 for (i = 0; i < sizeof(size_t); i++) {
858 out[i] = (unsigned char)((value >> (8 * i)) & 0xff);
859 }
860 for (i = sizeof(size_t); i < 8; i++) {
861 out[i] = 0;
Alexandre Vassalottided929b2013-11-24 22:41:13 -0800862 }
Antoine Pitrou8f2ee6e2013-11-23 21:05:08 +0100863}
864
865static void
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100866_Pickler_WriteFrameHeader(PicklerObject *self, char *qdata, size_t frame_len)
867{
Antoine Pitrou8f2ee6e2013-11-23 21:05:08 +0100868 qdata[0] = FRAME;
869 _write_size64(qdata + 1, frame_len);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100870}
871
872static int
873_Pickler_CommitFrame(PicklerObject *self)
874{
875 size_t frame_len;
876 char *qdata;
877
878 if (!self->framing || self->frame_start == -1)
879 return 0;
880 frame_len = self->output_len - self->frame_start - FRAME_HEADER_SIZE;
881 qdata = PyBytes_AS_STRING(self->output_buffer) + self->frame_start;
882 _Pickler_WriteFrameHeader(self, qdata, frame_len);
883 self->frame_start = -1;
884 return 0;
885}
886
887static int
888_Pickler_OpcodeBoundary(PicklerObject *self)
889{
890 Py_ssize_t frame_len;
891
892 if (!self->framing || self->frame_start == -1)
893 return 0;
894 frame_len = self->output_len - self->frame_start - FRAME_HEADER_SIZE;
895 if (frame_len >= FRAME_SIZE_TARGET)
896 return _Pickler_CommitFrame(self);
897 else
898 return 0;
899}
900
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000901static PyObject *
902_Pickler_GetString(PicklerObject *self)
903{
904 PyObject *output_buffer = self->output_buffer;
905
906 assert(self->output_buffer != NULL);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100907
908 if (_Pickler_CommitFrame(self))
909 return NULL;
910
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000911 self->output_buffer = NULL;
912 /* Resize down to exact size */
913 if (_PyBytes_Resize(&output_buffer, self->output_len) < 0)
914 return NULL;
915 return output_buffer;
916}
917
918static int
919_Pickler_FlushToFile(PicklerObject *self)
920{
921 PyObject *output, *result;
922
923 assert(self->write != NULL);
924
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100925 /* This will commit the frame first */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000926 output = _Pickler_GetString(self);
927 if (output == NULL)
928 return -1;
929
Alexandre Vassalotti20c28c12013-11-27 02:26:54 -0800930 result = _Pickle_FastCall(self->write, output);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000931 Py_XDECREF(result);
932 return (result == NULL) ? -1 : 0;
933}
934
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200935static Py_ssize_t
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100936_Pickler_Write(PicklerObject *self, const char *s, Py_ssize_t data_len)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000937{
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100938 Py_ssize_t i, n, required;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000939 char *buffer;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100940 int need_new_frame;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000941
942 assert(s != NULL);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100943 need_new_frame = (self->framing && self->frame_start == -1);
944
945 if (need_new_frame)
946 n = data_len + FRAME_HEADER_SIZE;
947 else
948 n = data_len;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000949
950 required = self->output_len + n;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100951 if (required > self->max_output_len) {
952 /* Make place in buffer for the pickle chunk */
953 if (self->output_len >= PY_SSIZE_T_MAX / 2 - n) {
954 PyErr_NoMemory();
955 return -1;
956 }
957 self->max_output_len = (self->output_len + n) / 2 * 3;
958 if (_PyBytes_Resize(&self->output_buffer, self->max_output_len) < 0)
959 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000960 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000961 buffer = PyBytes_AS_STRING(self->output_buffer);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100962 if (need_new_frame) {
963 /* Setup new frame */
964 Py_ssize_t frame_start = self->output_len;
965 self->frame_start = frame_start;
966 for (i = 0; i < FRAME_HEADER_SIZE; i++) {
967 /* Write an invalid value, for debugging */
968 buffer[frame_start + i] = 0xFE;
969 }
970 self->output_len += FRAME_HEADER_SIZE;
971 }
972 if (data_len < 8) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000973 /* This is faster than memcpy when the string is short. */
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100974 for (i = 0; i < data_len; i++) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000975 buffer[self->output_len + i] = s[i];
976 }
977 }
978 else {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100979 memcpy(buffer + self->output_len, s, data_len);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000980 }
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100981 self->output_len += data_len;
982 return data_len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000983}
984
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000985static PicklerObject *
986_Pickler_New(void)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000987{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000988 PicklerObject *self;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000989
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000990 self = PyObject_GC_New(PicklerObject, &Pickler_Type);
991 if (self == NULL)
992 return NULL;
993
994 self->pers_func = NULL;
Antoine Pitrou8d3c2902012-03-04 18:31:48 +0100995 self->dispatch_table = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000996 self->write = NULL;
997 self->proto = 0;
998 self->bin = 0;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +0100999 self->framing = 0;
1000 self->frame_start = -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001001 self->fast = 0;
1002 self->fast_nesting = 0;
1003 self->fix_imports = 0;
1004 self->fast_memo = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001005 self->max_output_len = WRITE_BUF_SIZE;
1006 self->output_len = 0;
Victor Stinner68c8ea22013-07-11 22:56:25 +02001007
1008 self->memo = PyMemoTable_New();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001009 self->output_buffer = PyBytes_FromStringAndSize(NULL,
1010 self->max_output_len);
Victor Stinner68c8ea22013-07-11 22:56:25 +02001011
1012 if (self->memo == NULL || self->output_buffer == NULL) {
Victor Stinnerc31df042013-07-12 00:08:59 +02001013 Py_DECREF(self);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001014 return NULL;
1015 }
1016 return self;
1017}
1018
1019static int
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08001020_Pickler_SetProtocol(PicklerObject *self, PyObject *protocol, int fix_imports)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001021{
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08001022 long proto;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001023
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08001024 if (protocol == NULL || protocol == Py_None) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001025 proto = DEFAULT_PROTOCOL;
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08001026 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001027 else {
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08001028 proto = PyLong_AsLong(protocol);
1029 if (proto < 0) {
1030 if (proto == -1 && PyErr_Occurred())
1031 return -1;
1032 proto = HIGHEST_PROTOCOL;
1033 }
1034 else if (proto > HIGHEST_PROTOCOL) {
1035 PyErr_Format(PyExc_ValueError, "pickle protocol must be <= %d",
1036 HIGHEST_PROTOCOL);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001037 return -1;
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08001038 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001039 }
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08001040 self->proto = (int)proto;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001041 self->bin = proto > 0;
1042 self->fix_imports = fix_imports && proto < 3;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001043 return 0;
1044}
1045
1046/* Returns -1 (with an exception set) on failure, 0 on success. This may
1047 be called once on a freshly created Pickler. */
1048static int
1049_Pickler_SetOutputStream(PicklerObject *self, PyObject *file)
1050{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02001051 _Py_IDENTIFIER(write);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001052 assert(file != NULL);
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02001053 self->write = _PyObject_GetAttrId(file, &PyId_write);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001054 if (self->write == NULL) {
1055 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1056 PyErr_SetString(PyExc_TypeError,
1057 "file must have a 'write' attribute");
1058 return -1;
1059 }
1060
1061 return 0;
1062}
1063
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001064/* Returns the size of the input on success, -1 on failure. This takes its
1065 own reference to `input`. */
1066static Py_ssize_t
1067_Unpickler_SetStringInput(UnpicklerObject *self, PyObject *input)
1068{
1069 if (self->buffer.buf != NULL)
1070 PyBuffer_Release(&self->buffer);
1071 if (PyObject_GetBuffer(input, &self->buffer, PyBUF_CONTIG_RO) < 0)
1072 return -1;
1073 self->input_buffer = self->buffer.buf;
1074 self->input_len = self->buffer.len;
1075 self->next_read_idx = 0;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001076 self->prefetched_idx = self->input_len;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001077 return self->input_len;
1078}
1079
Antoine Pitrou04248a82010-10-12 20:51:21 +00001080static int
1081_Unpickler_SkipConsumed(UnpicklerObject *self)
1082{
Victor Stinnerb43ad1d2013-10-31 13:38:42 +01001083 Py_ssize_t consumed;
1084 PyObject *r;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001085
Victor Stinnerb43ad1d2013-10-31 13:38:42 +01001086 consumed = self->next_read_idx - self->prefetched_idx;
1087 if (consumed <= 0)
1088 return 0;
1089
1090 assert(self->peek); /* otherwise we did something wrong */
1091 /* This makes an useless copy... */
1092 r = PyObject_CallFunction(self->read, "n", consumed);
1093 if (r == NULL)
1094 return -1;
1095 Py_DECREF(r);
1096
1097 self->prefetched_idx = self->next_read_idx;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001098 return 0;
1099}
1100
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001101static const Py_ssize_t READ_WHOLE_LINE = -1;
1102
1103/* If reading from a file, we need to only pull the bytes we need, since there
1104 may be multiple pickle objects arranged contiguously in the same input
1105 buffer.
1106
1107 If `n` is READ_WHOLE_LINE, read a whole line. Otherwise, read up to `n`
1108 bytes from the input stream/buffer.
1109
1110 Update the unpickler's input buffer with the newly-read data. Returns -1 on
1111 failure; on success, returns the number of bytes read from the file.
1112
1113 On success, self->input_len will be 0; this is intentional so that when
1114 unpickling from a file, the "we've run out of data" code paths will trigger,
1115 causing the Unpickler to go back to the file for more data. Use the returned
1116 size to tell you how much data you can process. */
1117static Py_ssize_t
1118_Unpickler_ReadFromFile(UnpicklerObject *self, Py_ssize_t n)
1119{
1120 PyObject *data;
Serhiy Storchaka6fe39b72013-11-30 23:15:38 +02001121 Py_ssize_t read_size;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001122
1123 assert(self->read != NULL);
Victor Stinner121aab42011-09-29 23:40:53 +02001124
Antoine Pitrou04248a82010-10-12 20:51:21 +00001125 if (_Unpickler_SkipConsumed(self) < 0)
1126 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001127
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08001128 if (n == READ_WHOLE_LINE) {
Alexandre Vassalotti6bf41e52013-11-28 15:17:29 -08001129 PyObject *empty_tuple = PyTuple_New(0);
1130 data = PyObject_Call(self->readline, empty_tuple, NULL);
1131 Py_DECREF(empty_tuple);
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08001132 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001133 else {
Serhiy Storchaka6fe39b72013-11-30 23:15:38 +02001134 PyObject *len;
1135 /* Prefetch some data without advancing the file pointer, if possible */
1136 if (self->peek && n < PREFETCH) {
1137 len = PyLong_FromSsize_t(PREFETCH);
1138 if (len == NULL)
1139 return -1;
1140 data = _Pickle_FastCall(self->peek, len);
1141 if (data == NULL) {
1142 if (!PyErr_ExceptionMatches(PyExc_NotImplementedError))
1143 return -1;
1144 /* peek() is probably not supported by the given file object */
1145 PyErr_Clear();
1146 Py_CLEAR(self->peek);
1147 }
1148 else {
1149 read_size = _Unpickler_SetStringInput(self, data);
1150 Py_DECREF(data);
1151 self->prefetched_idx = 0;
1152 if (n <= read_size)
1153 return n;
1154 }
1155 }
1156 len = PyLong_FromSsize_t(n);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001157 if (len == NULL)
1158 return -1;
Alexandre Vassalotti20c28c12013-11-27 02:26:54 -08001159 data = _Pickle_FastCall(self->read, len);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001160 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001161 if (data == NULL)
1162 return -1;
1163
Serhiy Storchaka6fe39b72013-11-30 23:15:38 +02001164 read_size = _Unpickler_SetStringInput(self, data);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001165 Py_DECREF(data);
1166 return read_size;
1167}
1168
1169/* Read `n` bytes from the unpickler's data source, storing the result in `*s`.
1170
1171 This should be used for all data reads, rather than accessing the unpickler's
1172 input buffer directly. This method deals correctly with reading from input
1173 streams, which the input buffer doesn't deal with.
1174
1175 Note that when reading from a file-like object, self->next_read_idx won't
1176 be updated (it should remain at 0 for the entire unpickling process). You
1177 should use this function's return value to know how many bytes you can
1178 consume.
1179
1180 Returns -1 (with an exception set) on failure. On success, return the
1181 number of chars read. */
1182static Py_ssize_t
Alexandre Vassalottib6a2f2a2013-11-23 20:30:03 -08001183_Unpickler_Read(UnpicklerObject *self, char **s, Py_ssize_t n)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001184{
Antoine Pitrou04248a82010-10-12 20:51:21 +00001185 Py_ssize_t num_read;
1186
Antoine Pitrou04248a82010-10-12 20:51:21 +00001187 if (self->next_read_idx + n <= self->input_len) {
1188 *s = self->input_buffer + self->next_read_idx;
1189 self->next_read_idx += n;
1190 return n;
1191 }
1192 if (!self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001193 PyErr_Format(PyExc_EOFError, "Ran out of input");
Amaury Forgeot d'Arc3e4e72f2008-11-11 20:05:06 +00001194 return -1;
1195 }
Antoine Pitrou04248a82010-10-12 20:51:21 +00001196 num_read = _Unpickler_ReadFromFile(self, n);
1197 if (num_read < 0)
1198 return -1;
1199 if (num_read < n) {
1200 PyErr_Format(PyExc_EOFError, "Ran out of input");
1201 return -1;
1202 }
1203 *s = self->input_buffer;
1204 self->next_read_idx = n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001205 return n;
1206}
1207
1208static Py_ssize_t
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001209_Unpickler_CopyLine(UnpicklerObject *self, char *line, Py_ssize_t len,
1210 char **result)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001211{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001212 char *input_line = PyMem_Realloc(self->input_line, len + 1);
Victor Stinner42024562013-07-12 00:53:57 +02001213 if (input_line == NULL) {
1214 PyErr_NoMemory();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001215 return -1;
Victor Stinner42024562013-07-12 00:53:57 +02001216 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001217
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001218 memcpy(input_line, line, len);
1219 input_line[len] = '\0';
1220 self->input_line = input_line;
1221 *result = self->input_line;
1222 return len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001223}
1224
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001225/* Read a line from the input stream/buffer. If we run off the end of the input
1226 before hitting \n, return the data we found.
1227
1228 Returns the number of chars read, or -1 on failure. */
1229static Py_ssize_t
1230_Unpickler_Readline(UnpicklerObject *self, char **result)
1231{
1232 Py_ssize_t i, num_read;
1233
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001234 for (i = self->next_read_idx; i < self->input_len; i++) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001235 if (self->input_buffer[i] == '\n') {
1236 char *line_start = self->input_buffer + self->next_read_idx;
1237 num_read = i - self->next_read_idx + 1;
1238 self->next_read_idx = i + 1;
1239 return _Unpickler_CopyLine(self, line_start, num_read, result);
1240 }
1241 }
1242 if (self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001243 num_read = _Unpickler_ReadFromFile(self, READ_WHOLE_LINE);
1244 if (num_read < 0)
1245 return -1;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001246 self->next_read_idx = num_read;
Antoine Pitrouf6c7a852011-08-11 21:04:02 +02001247 return _Unpickler_CopyLine(self, self->input_buffer, num_read, result);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001248 }
Victor Stinner121aab42011-09-29 23:40:53 +02001249
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001250 /* If we get here, we've run off the end of the input string. Return the
1251 remaining string and let the caller figure it out. */
1252 *result = self->input_buffer + self->next_read_idx;
1253 num_read = i - self->next_read_idx;
1254 self->next_read_idx = i;
1255 return num_read;
1256}
1257
1258/* Returns -1 (with an exception set) on failure, 0 on success. The memo array
1259 will be modified in place. */
1260static int
1261_Unpickler_ResizeMemoList(UnpicklerObject *self, Py_ssize_t new_size)
1262{
1263 Py_ssize_t i;
1264 PyObject **memo;
1265
1266 assert(new_size > self->memo_size);
1267
1268 memo = PyMem_REALLOC(self->memo, new_size * sizeof(PyObject *));
1269 if (memo == NULL) {
1270 PyErr_NoMemory();
1271 return -1;
1272 }
1273 self->memo = memo;
1274 for (i = self->memo_size; i < new_size; i++)
1275 self->memo[i] = NULL;
1276 self->memo_size = new_size;
1277 return 0;
1278}
1279
1280/* Returns NULL if idx is out of bounds. */
1281static PyObject *
1282_Unpickler_MemoGet(UnpicklerObject *self, Py_ssize_t idx)
1283{
1284 if (idx < 0 || idx >= self->memo_size)
1285 return NULL;
1286
1287 return self->memo[idx];
1288}
1289
1290/* Returns -1 (with an exception set) on failure, 0 on success.
1291 This takes its own reference to `value`. */
1292static int
1293_Unpickler_MemoPut(UnpicklerObject *self, Py_ssize_t idx, PyObject *value)
1294{
1295 PyObject *old_item;
1296
1297 if (idx >= self->memo_size) {
1298 if (_Unpickler_ResizeMemoList(self, idx * 2) < 0)
1299 return -1;
1300 assert(idx < self->memo_size);
1301 }
1302 Py_INCREF(value);
1303 old_item = self->memo[idx];
1304 self->memo[idx] = value;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001305 if (old_item != NULL) {
1306 Py_DECREF(old_item);
1307 }
1308 else {
1309 self->memo_len++;
1310 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001311 return 0;
1312}
1313
1314static PyObject **
1315_Unpickler_NewMemo(Py_ssize_t new_size)
1316{
1317 PyObject **memo = PyMem_MALLOC(new_size * sizeof(PyObject *));
Victor Stinner42024562013-07-12 00:53:57 +02001318 if (memo == NULL) {
1319 PyErr_NoMemory();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001320 return NULL;
Victor Stinner42024562013-07-12 00:53:57 +02001321 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001322 memset(memo, 0, new_size * sizeof(PyObject *));
1323 return memo;
1324}
1325
1326/* Free the unpickler's memo, taking care to decref any items left in it. */
1327static void
1328_Unpickler_MemoCleanup(UnpicklerObject *self)
1329{
1330 Py_ssize_t i;
1331 PyObject **memo = self->memo;
1332
1333 if (self->memo == NULL)
1334 return;
1335 self->memo = NULL;
1336 i = self->memo_size;
1337 while (--i >= 0) {
1338 Py_XDECREF(memo[i]);
1339 }
1340 PyMem_FREE(memo);
1341}
1342
1343static UnpicklerObject *
1344_Unpickler_New(void)
1345{
1346 UnpicklerObject *self;
1347
1348 self = PyObject_GC_New(UnpicklerObject, &Unpickler_Type);
1349 if (self == NULL)
1350 return NULL;
1351
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001352 self->pers_func = NULL;
1353 self->input_buffer = NULL;
1354 self->input_line = NULL;
1355 self->input_len = 0;
1356 self->next_read_idx = 0;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001357 self->prefetched_idx = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001358 self->read = NULL;
1359 self->readline = NULL;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001360 self->peek = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001361 self->encoding = NULL;
1362 self->errors = NULL;
1363 self->marks = NULL;
1364 self->num_marks = 0;
1365 self->marks_size = 0;
1366 self->proto = 0;
1367 self->fix_imports = 0;
Victor Stinner68c8ea22013-07-11 22:56:25 +02001368 memset(&self->buffer, 0, sizeof(Py_buffer));
1369 self->memo_size = 32;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001370 self->memo_len = 0;
Victor Stinner68c8ea22013-07-11 22:56:25 +02001371 self->memo = _Unpickler_NewMemo(self->memo_size);
1372 self->stack = (Pdata *)Pdata_New();
1373
1374 if (self->memo == NULL || self->stack == NULL) {
1375 Py_DECREF(self);
1376 return NULL;
1377 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001378
1379 return self;
1380}
1381
1382/* Returns -1 (with an exception set) on failure, 0 on success. This may
1383 be called once on a freshly created Pickler. */
1384static int
1385_Unpickler_SetInputStream(UnpicklerObject *self, PyObject *file)
1386{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02001387 _Py_IDENTIFIER(peek);
1388 _Py_IDENTIFIER(read);
1389 _Py_IDENTIFIER(readline);
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02001390
1391 self->peek = _PyObject_GetAttrId(file, &PyId_peek);
Antoine Pitrou04248a82010-10-12 20:51:21 +00001392 if (self->peek == NULL) {
1393 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1394 PyErr_Clear();
1395 else
1396 return -1;
1397 }
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02001398 self->read = _PyObject_GetAttrId(file, &PyId_read);
1399 self->readline = _PyObject_GetAttrId(file, &PyId_readline);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001400 if (self->readline == NULL || self->read == NULL) {
1401 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1402 PyErr_SetString(PyExc_TypeError,
1403 "file must have 'read' and 'readline' attributes");
1404 Py_CLEAR(self->read);
1405 Py_CLEAR(self->readline);
Antoine Pitrou04248a82010-10-12 20:51:21 +00001406 Py_CLEAR(self->peek);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001407 return -1;
1408 }
1409 return 0;
1410}
1411
1412/* Returns -1 (with an exception set) on failure, 0 on success. This may
1413 be called once on a freshly created Pickler. */
1414static int
1415_Unpickler_SetInputEncoding(UnpicklerObject *self,
1416 const char *encoding,
1417 const char *errors)
1418{
1419 if (encoding == NULL)
1420 encoding = "ASCII";
1421 if (errors == NULL)
1422 errors = "strict";
1423
Victor Stinner49fc8ec2013-07-07 23:30:24 +02001424 self->encoding = _PyMem_Strdup(encoding);
1425 self->errors = _PyMem_Strdup(errors);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001426 if (self->encoding == NULL || self->errors == NULL) {
1427 PyErr_NoMemory();
1428 return -1;
1429 }
1430 return 0;
1431}
1432
1433/* Generate a GET opcode for an object stored in the memo. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001434static int
1435memo_get(PicklerObject *self, PyObject *key)
1436{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001437 Py_ssize_t *value;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001438 char pdata[30];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001439 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001440
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001441 value = PyMemoTable_Get(self->memo, key);
1442 if (value == NULL) {
1443 PyErr_SetObject(PyExc_KeyError, key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001444 return -1;
1445 }
1446
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001447 if (!self->bin) {
1448 pdata[0] = GET;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001449 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1,
1450 "%" PY_FORMAT_SIZE_T "d\n", *value);
1451 len = strlen(pdata);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001452 }
1453 else {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001454 if (*value < 256) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001455 pdata[0] = BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001456 pdata[1] = (unsigned char)(*value & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001457 len = 2;
1458 }
Serhiy Storchaka67c719b2014-09-05 10:10:23 +03001459 else if ((size_t)*value <= 0xffffffffUL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001460 pdata[0] = LONG_BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001461 pdata[1] = (unsigned char)(*value & 0xff);
1462 pdata[2] = (unsigned char)((*value >> 8) & 0xff);
1463 pdata[3] = (unsigned char)((*value >> 16) & 0xff);
1464 pdata[4] = (unsigned char)((*value >> 24) & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001465 len = 5;
1466 }
1467 else { /* unlikely */
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08001468 PickleState *st = _Pickle_GetGlobalState();
1469 PyErr_SetString(st->PicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001470 "memo id too large for LONG_BINGET");
1471 return -1;
1472 }
1473 }
1474
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001475 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001476 return -1;
1477
1478 return 0;
1479}
1480
1481/* Store an object in the memo, assign it a new unique ID based on the number
1482 of objects currently stored in the memo and generate a PUT opcode. */
1483static int
1484memo_put(PicklerObject *self, PyObject *obj)
1485{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001486 char pdata[30];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001487 Py_ssize_t len;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001488 Py_ssize_t idx;
1489
1490 const char memoize_op = MEMOIZE;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001491
1492 if (self->fast)
1493 return 0;
1494
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001495 idx = PyMemoTable_Size(self->memo);
1496 if (PyMemoTable_Set(self->memo, obj, idx) < 0)
1497 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001498
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001499 if (self->proto >= 4) {
1500 if (_Pickler_Write(self, &memoize_op, 1) < 0)
1501 return -1;
1502 return 0;
1503 }
1504 else if (!self->bin) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001505 pdata[0] = PUT;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001506 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1,
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001507 "%" PY_FORMAT_SIZE_T "d\n", idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001508 len = strlen(pdata);
1509 }
1510 else {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001511 if (idx < 256) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001512 pdata[0] = BINPUT;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001513 pdata[1] = (unsigned char)idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001514 len = 2;
1515 }
Serhiy Storchaka67c719b2014-09-05 10:10:23 +03001516 else if ((size_t)idx <= 0xffffffffUL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001517 pdata[0] = LONG_BINPUT;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001518 pdata[1] = (unsigned char)(idx & 0xff);
1519 pdata[2] = (unsigned char)((idx >> 8) & 0xff);
1520 pdata[3] = (unsigned char)((idx >> 16) & 0xff);
1521 pdata[4] = (unsigned char)((idx >> 24) & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001522 len = 5;
1523 }
1524 else { /* unlikely */
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08001525 PickleState *st = _Pickle_GetGlobalState();
1526 PyErr_SetString(st->PicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001527 "memo id too large for LONG_BINPUT");
1528 return -1;
1529 }
1530 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001531 if (_Pickler_Write(self, pdata, len) < 0)
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001532 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001533
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001534 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001535}
1536
1537static PyObject *
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001538get_dotted_path(PyObject *obj, PyObject *name, int allow_qualname) {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001539 _Py_static_string(PyId_dot, ".");
1540 _Py_static_string(PyId_locals, "<locals>");
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001541 PyObject *dotted_path;
1542 Py_ssize_t i, n;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001543
1544 dotted_path = PyUnicode_Split(name, _PyUnicode_FromId(&PyId_dot), -1);
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001545 if (dotted_path == NULL)
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001546 return NULL;
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001547 n = PyList_GET_SIZE(dotted_path);
1548 assert(n >= 1);
1549 if (!allow_qualname && n > 1) {
Antoine Pitrou6cd5eda2014-12-02 00:20:03 +01001550 if (obj == NULL)
1551 PyErr_Format(PyExc_AttributeError,
1552 "Can't pickle qualified object %R; "
1553 "use protocols >= 4 to enable support",
1554 name);
1555 else
1556 PyErr_Format(PyExc_AttributeError,
1557 "Can't pickle qualified attribute %R on %R; "
1558 "use protocols >= 4 to enable support",
1559 name, obj);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001560 Py_DECREF(dotted_path);
1561 return NULL;
1562 }
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001563 for (i = 0; i < n; i++) {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001564 PyObject *subpath = PyList_GET_ITEM(dotted_path, i);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001565 PyObject *result = PyUnicode_RichCompare(
1566 subpath, _PyUnicode_FromId(&PyId_locals), Py_EQ);
1567 int is_equal = (result == Py_True);
1568 assert(PyBool_Check(result));
1569 Py_DECREF(result);
1570 if (is_equal) {
Antoine Pitrou6cd5eda2014-12-02 00:20:03 +01001571 if (obj == NULL)
1572 PyErr_Format(PyExc_AttributeError,
1573 "Can't pickle local object %R", name);
1574 else
1575 PyErr_Format(PyExc_AttributeError,
1576 "Can't pickle local attribute %R on %R", name, obj);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001577 Py_DECREF(dotted_path);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001578 return NULL;
1579 }
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001580 }
1581 return dotted_path;
1582}
1583
1584static PyObject *
1585get_deep_attribute(PyObject *obj, PyObject *names)
1586{
1587 Py_ssize_t i, n;
1588
1589 assert(PyList_CheckExact(names));
1590 Py_INCREF(obj);
1591 n = PyList_GET_SIZE(names);
1592 for (i = 0; i < n; i++) {
1593 PyObject *name = PyList_GET_ITEM(names, i);
1594 PyObject *tmp;
1595 tmp = PyObject_GetAttr(obj, name);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001596 Py_DECREF(obj);
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001597 if (tmp == NULL)
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001598 return NULL;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001599 obj = tmp;
1600 }
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001601 return obj;
1602}
1603
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001604static void
1605reformat_attribute_error(PyObject *obj, PyObject *name)
1606{
1607 if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
1608 PyErr_Clear();
1609 PyErr_Format(PyExc_AttributeError,
1610 "Can't get attribute %R on %R", name, obj);
1611 }
1612}
1613
1614
1615static PyObject *
1616getattribute(PyObject *obj, PyObject *name, int allow_qualname)
1617{
1618 PyObject *dotted_path, *attr;
1619
1620 dotted_path = get_dotted_path(obj, name, allow_qualname);
1621 if (dotted_path == NULL)
1622 return NULL;
1623 attr = get_deep_attribute(obj, dotted_path);
1624 Py_DECREF(dotted_path);
1625 if (attr == NULL)
1626 reformat_attribute_error(obj, name);
1627 return attr;
1628}
1629
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001630static PyObject *
1631whichmodule(PyObject *global, PyObject *global_name, int allow_qualname)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001632{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001633 PyObject *module_name;
1634 PyObject *modules_dict;
1635 PyObject *module;
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001636 Py_ssize_t i;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001637 _Py_IDENTIFIER(__module__);
1638 _Py_IDENTIFIER(modules);
1639 _Py_IDENTIFIER(__main__);
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001640 PyObject *dotted_path;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001641
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001642 module_name = _PyObject_GetAttrId(global, &PyId___module__);
1643
1644 if (module_name == NULL) {
1645 if (!PyErr_ExceptionMatches(PyExc_AttributeError))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001646 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001647 PyErr_Clear();
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001648 }
1649 else {
1650 /* In some rare cases (e.g., bound methods of extension types),
1651 __module__ can be None. If it is so, then search sys.modules for
1652 the module of global. */
1653 if (module_name != Py_None)
1654 return module_name;
1655 Py_CLEAR(module_name);
1656 }
1657 assert(module_name == NULL);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001658
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001659 /* Fallback on walking sys.modules */
Victor Stinnerbb520202013-11-06 22:40:41 +01001660 modules_dict = _PySys_GetObjectId(&PyId_modules);
Victor Stinner1e53bba2013-07-16 22:26:05 +02001661 if (modules_dict == NULL) {
1662 PyErr_SetString(PyExc_RuntimeError, "unable to get sys.modules");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001663 return NULL;
Victor Stinner1e53bba2013-07-16 22:26:05 +02001664 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001665
Antoine Pitrou6cd5eda2014-12-02 00:20:03 +01001666 dotted_path = get_dotted_path(NULL, global_name, allow_qualname);
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001667 if (dotted_path == NULL)
1668 return NULL;
1669
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001670 i = 0;
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001671 while (PyDict_Next(modules_dict, &i, &module_name, &module)) {
1672 PyObject *candidate;
1673 if (PyUnicode_Check(module_name) &&
1674 !PyUnicode_CompareWithASCIIString(module_name, "__main__"))
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001675 continue;
1676 if (module == Py_None)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001677 continue;
1678
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001679 candidate = get_deep_attribute(module, dotted_path);
1680 if (candidate == NULL) {
1681 if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
1682 Py_DECREF(dotted_path);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001683 return NULL;
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001684 }
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001685 PyErr_Clear();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001686 continue;
1687 }
1688
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001689 if (candidate == global) {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001690 Py_INCREF(module_name);
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001691 Py_DECREF(dotted_path);
1692 Py_DECREF(candidate);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001693 return module_name;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001694 }
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001695 Py_DECREF(candidate);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001696 }
1697
1698 /* If no module is found, use __main__. */
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01001699 module_name = _PyUnicode_FromId(&PyId___main__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001700 Py_INCREF(module_name);
Antoine Pitroufce60ea2014-10-23 22:47:50 +02001701 Py_DECREF(dotted_path);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001702 return module_name;
1703}
1704
1705/* fast_save_enter() and fast_save_leave() are guards against recursive
1706 objects when Pickler is used with the "fast mode" (i.e., with object
1707 memoization disabled). If the nesting of a list or dict object exceed
1708 FAST_NESTING_LIMIT, these guards will start keeping an internal
1709 reference to the seen list or dict objects and check whether these objects
1710 are recursive. These are not strictly necessary, since save() has a
1711 hard-coded recursion limit, but they give a nicer error message than the
1712 typical RuntimeError. */
1713static int
1714fast_save_enter(PicklerObject *self, PyObject *obj)
1715{
1716 /* if fast_nesting < 0, we're doing an error exit. */
1717 if (++self->fast_nesting >= FAST_NESTING_LIMIT) {
1718 PyObject *key = NULL;
1719 if (self->fast_memo == NULL) {
1720 self->fast_memo = PyDict_New();
1721 if (self->fast_memo == NULL) {
1722 self->fast_nesting = -1;
1723 return 0;
1724 }
1725 }
1726 key = PyLong_FromVoidPtr(obj);
1727 if (key == NULL)
1728 return 0;
Alexandre Vassalotti567eba12013-11-28 17:09:16 -08001729 if (PyDict_GetItemWithError(self->fast_memo, key)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001730 Py_DECREF(key);
1731 PyErr_Format(PyExc_ValueError,
1732 "fast mode: can't pickle cyclic objects "
1733 "including object type %.200s at %p",
1734 obj->ob_type->tp_name, obj);
1735 self->fast_nesting = -1;
1736 return 0;
1737 }
Alexandre Vassalotti567eba12013-11-28 17:09:16 -08001738 if (PyErr_Occurred()) {
1739 return 0;
1740 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001741 if (PyDict_SetItem(self->fast_memo, key, Py_None) < 0) {
1742 Py_DECREF(key);
1743 self->fast_nesting = -1;
1744 return 0;
1745 }
1746 Py_DECREF(key);
1747 }
1748 return 1;
1749}
1750
1751static int
1752fast_save_leave(PicklerObject *self, PyObject *obj)
1753{
1754 if (self->fast_nesting-- >= FAST_NESTING_LIMIT) {
1755 PyObject *key = PyLong_FromVoidPtr(obj);
1756 if (key == NULL)
1757 return 0;
1758 if (PyDict_DelItem(self->fast_memo, key) < 0) {
1759 Py_DECREF(key);
1760 return 0;
1761 }
1762 Py_DECREF(key);
1763 }
1764 return 1;
1765}
1766
1767static int
1768save_none(PicklerObject *self, PyObject *obj)
1769{
1770 const char none_op = NONE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001771 if (_Pickler_Write(self, &none_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001772 return -1;
1773
1774 return 0;
1775}
1776
1777static int
1778save_bool(PicklerObject *self, PyObject *obj)
1779{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001780 if (self->proto >= 2) {
Alexandre Vassalotti8a67f522013-11-24 21:40:18 -08001781 const char bool_op = (obj == Py_True) ? NEWTRUE : NEWFALSE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001782 if (_Pickler_Write(self, &bool_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001783 return -1;
1784 }
Alexandre Vassalotti8a67f522013-11-24 21:40:18 -08001785 else {
1786 /* These aren't opcodes -- they're ways to pickle bools before protocol 2
1787 * so that unpicklers written before bools were introduced unpickle them
1788 * as ints, but unpicklers after can recognize that bools were intended.
1789 * Note that protocol 2 added direct ways to pickle bools.
1790 */
1791 const char *bool_str = (obj == Py_True) ? "I01\n" : "I00\n";
1792 if (_Pickler_Write(self, bool_str, strlen(bool_str)) < 0)
1793 return -1;
1794 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001795 return 0;
1796}
1797
1798static int
Alexandre Vassalottided929b2013-11-24 22:41:13 -08001799save_long(PicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001800{
Alexandre Vassalottided929b2013-11-24 22:41:13 -08001801 PyObject *repr = NULL;
1802 Py_ssize_t size;
1803 long val;
1804 int status = 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001805
Alexandre Vassalottided929b2013-11-24 22:41:13 -08001806 const char long_op = LONG;
1807
1808 val= PyLong_AsLong(obj);
1809 if (val == -1 && PyErr_Occurred()) {
1810 /* out of range for int pickling */
1811 PyErr_Clear();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001812 }
Alexandre Vassalottided929b2013-11-24 22:41:13 -08001813 else if (self->bin &&
1814 (sizeof(long) <= 4 ||
Alexandre Vassalotti1048fb52013-11-25 11:35:46 -08001815 (val <= 0x7fffffffL && val >= (-0x7fffffffL - 1)))) {
Larry Hastings61272b72014-01-07 12:41:53 -08001816 /* result fits in a signed 4-byte integer.
Alexandre Vassalotti1048fb52013-11-25 11:35:46 -08001817
1818 Note: we can't use -0x80000000L in the above condition because some
1819 compilers (e.g., MSVC) will promote 0x80000000L to an unsigned type
1820 before applying the unary minus when sizeof(long) <= 4. The
1821 resulting value stays unsigned which is commonly not what we want,
1822 so MSVC happily warns us about it. However, that result would have
1823 been fine because we guard for sizeof(long) <= 4 which turns the
1824 condition true in that particular case. */
Alexandre Vassalottided929b2013-11-24 22:41:13 -08001825 char pdata[32];
1826 Py_ssize_t len = 0;
1827
1828 pdata[1] = (unsigned char)(val & 0xff);
1829 pdata[2] = (unsigned char)((val >> 8) & 0xff);
1830 pdata[3] = (unsigned char)((val >> 16) & 0xff);
1831 pdata[4] = (unsigned char)((val >> 24) & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001832
1833 if ((pdata[4] == 0) && (pdata[3] == 0)) {
1834 if (pdata[2] == 0) {
1835 pdata[0] = BININT1;
1836 len = 2;
1837 }
1838 else {
1839 pdata[0] = BININT2;
1840 len = 3;
1841 }
1842 }
1843 else {
1844 pdata[0] = BININT;
1845 len = 5;
1846 }
1847
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001848 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001849 return -1;
Alexandre Vassalottided929b2013-11-24 22:41:13 -08001850
1851 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001852 }
1853
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001854 if (self->proto >= 2) {
1855 /* Linear-time pickling. */
1856 size_t nbits;
1857 size_t nbytes;
1858 unsigned char *pdata;
1859 char header[5];
1860 int i;
1861 int sign = _PyLong_Sign(obj);
1862
1863 if (sign == 0) {
1864 header[0] = LONG1;
1865 header[1] = 0; /* It's 0 -- an empty bytestring. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001866 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001867 goto error;
1868 return 0;
1869 }
1870 nbits = _PyLong_NumBits(obj);
1871 if (nbits == (size_t)-1 && PyErr_Occurred())
1872 goto error;
1873 /* How many bytes do we need? There are nbits >> 3 full
1874 * bytes of data, and nbits & 7 leftover bits. If there
1875 * are any leftover bits, then we clearly need another
1876 * byte. Wnat's not so obvious is that we *probably*
1877 * need another byte even if there aren't any leftovers:
1878 * the most-significant bit of the most-significant byte
1879 * acts like a sign bit, and it's usually got a sense
Serhiy Storchaka95949422013-08-27 19:40:23 +03001880 * opposite of the one we need. The exception is ints
1881 * of the form -(2**(8*j-1)) for j > 0. Such an int is
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001882 * its own 256's-complement, so has the right sign bit
1883 * even without the extra byte. That's a pain to check
1884 * for in advance, though, so we always grab an extra
1885 * byte at the start, and cut it back later if possible.
1886 */
1887 nbytes = (nbits >> 3) + 1;
Antoine Pitroubf6ecf92012-11-24 20:40:21 +01001888 if (nbytes > 0x7fffffffL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001889 PyErr_SetString(PyExc_OverflowError,
Serhiy Storchaka95949422013-08-27 19:40:23 +03001890 "int too large to pickle");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001891 goto error;
1892 }
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001893 repr = PyBytes_FromStringAndSize(NULL, (Py_ssize_t)nbytes);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001894 if (repr == NULL)
1895 goto error;
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001896 pdata = (unsigned char *)PyBytes_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001897 i = _PyLong_AsByteArray((PyLongObject *)obj,
1898 pdata, nbytes,
1899 1 /* little endian */ , 1 /* signed */ );
1900 if (i < 0)
1901 goto error;
Serhiy Storchaka95949422013-08-27 19:40:23 +03001902 /* If the int is negative, this may be a byte more than
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001903 * needed. This is so iff the MSB is all redundant sign
1904 * bits.
1905 */
1906 if (sign < 0 &&
Victor Stinner121aab42011-09-29 23:40:53 +02001907 nbytes > 1 &&
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001908 pdata[nbytes - 1] == 0xff &&
1909 (pdata[nbytes - 2] & 0x80) != 0) {
1910 nbytes--;
1911 }
1912
1913 if (nbytes < 256) {
1914 header[0] = LONG1;
1915 header[1] = (unsigned char)nbytes;
1916 size = 2;
1917 }
1918 else {
1919 header[0] = LONG4;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001920 size = (Py_ssize_t) nbytes;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001921 for (i = 1; i < 5; i++) {
1922 header[i] = (unsigned char)(size & 0xff);
1923 size >>= 8;
1924 }
1925 size = 5;
1926 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001927 if (_Pickler_Write(self, header, size) < 0 ||
1928 _Pickler_Write(self, (char *)pdata, (int)nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001929 goto error;
1930 }
1931 else {
1932 char *string;
1933
Mark Dickinson8dd05142009-01-20 20:43:58 +00001934 /* proto < 2: write the repr and newline. This is quadratic-time (in
1935 the number of digits), in both directions. We add a trailing 'L'
1936 to the repr, for compatibility with Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001937
1938 repr = PyObject_Repr(obj);
1939 if (repr == NULL)
1940 goto error;
1941
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00001942 string = _PyUnicode_AsStringAndSize(repr, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001943 if (string == NULL)
1944 goto error;
1945
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001946 if (_Pickler_Write(self, &long_op, 1) < 0 ||
1947 _Pickler_Write(self, string, size) < 0 ||
1948 _Pickler_Write(self, "L\n", 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001949 goto error;
1950 }
1951
1952 if (0) {
1953 error:
1954 status = -1;
1955 }
1956 Py_XDECREF(repr);
1957
1958 return status;
1959}
1960
1961static int
1962save_float(PicklerObject *self, PyObject *obj)
1963{
1964 double x = PyFloat_AS_DOUBLE((PyFloatObject *)obj);
1965
1966 if (self->bin) {
1967 char pdata[9];
1968 pdata[0] = BINFLOAT;
1969 if (_PyFloat_Pack8(x, (unsigned char *)&pdata[1], 0) < 0)
1970 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001971 if (_Pickler_Write(self, pdata, 9) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001972 return -1;
Victor Stinner121aab42011-09-29 23:40:53 +02001973 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001974 else {
Eric Smith0923d1d2009-04-16 20:16:10 +00001975 int result = -1;
1976 char *buf = NULL;
1977 char op = FLOAT;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001978
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001979 if (_Pickler_Write(self, &op, 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001980 goto done;
1981
Mark Dickinson3e09f432009-04-17 08:41:23 +00001982 buf = PyOS_double_to_string(x, 'g', 17, 0, NULL);
Eric Smith0923d1d2009-04-16 20:16:10 +00001983 if (!buf) {
1984 PyErr_NoMemory();
1985 goto done;
1986 }
1987
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001988 if (_Pickler_Write(self, buf, strlen(buf)) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001989 goto done;
1990
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001991 if (_Pickler_Write(self, "\n", 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001992 goto done;
1993
1994 result = 0;
1995done:
1996 PyMem_Free(buf);
1997 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001998 }
1999
2000 return 0;
2001}
2002
2003static int
2004save_bytes(PicklerObject *self, PyObject *obj)
2005{
2006 if (self->proto < 3) {
2007 /* Older pickle protocols do not have an opcode for pickling bytes
2008 objects. Therefore, we need to fake the copy protocol (i.e.,
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05002009 the __reduce__ method) to permit bytes object unpickling.
2010
2011 Here we use a hack to be compatible with Python 2. Since in Python
2012 2 'bytes' is just an alias for 'str' (which has different
2013 parameters than the actual bytes object), we use codecs.encode
2014 to create the appropriate 'str' object when unpickled using
2015 Python 2 *and* the appropriate 'bytes' object when unpickled
2016 using Python 3. Again this is a hack and we don't need to do this
2017 with newer protocols. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002018 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002019 int status;
2020
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05002021 if (PyBytes_GET_SIZE(obj) == 0) {
2022 reduce_value = Py_BuildValue("(O())", (PyObject*)&PyBytes_Type);
2023 }
2024 else {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08002025 PickleState *st = _Pickle_GetGlobalState();
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05002026 PyObject *unicode_str =
2027 PyUnicode_DecodeLatin1(PyBytes_AS_STRING(obj),
2028 PyBytes_GET_SIZE(obj),
2029 "strict");
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002030 _Py_IDENTIFIER(latin1);
2031
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05002032 if (unicode_str == NULL)
2033 return -1;
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05002034 reduce_value = Py_BuildValue("(O(OO))",
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08002035 st->codecs_encode, unicode_str,
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002036 _PyUnicode_FromId(&PyId_latin1));
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05002037 Py_DECREF(unicode_str);
2038 }
2039
2040 if (reduce_value == NULL)
2041 return -1;
2042
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002043 /* save_reduce() will memoize the object automatically. */
2044 status = save_reduce(self, reduce_value, obj);
2045 Py_DECREF(reduce_value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002046 return status;
2047 }
2048 else {
2049 Py_ssize_t size;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002050 char header[9];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002051 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002052
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05002053 size = PyBytes_GET_SIZE(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002054 if (size < 0)
2055 return -1;
2056
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002057 if (size <= 0xff) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002058 header[0] = SHORT_BINBYTES;
2059 header[1] = (unsigned char)size;
2060 len = 2;
2061 }
Serhiy Storchaka67c719b2014-09-05 10:10:23 +03002062 else if ((size_t)size <= 0xffffffffUL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002063 header[0] = BINBYTES;
2064 header[1] = (unsigned char)(size & 0xff);
2065 header[2] = (unsigned char)((size >> 8) & 0xff);
2066 header[3] = (unsigned char)((size >> 16) & 0xff);
2067 header[4] = (unsigned char)((size >> 24) & 0xff);
2068 len = 5;
2069 }
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002070 else if (self->proto >= 4) {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002071 header[0] = BINBYTES8;
Alexandre Vassalotti1048fb52013-11-25 11:35:46 -08002072 _write_size64(header + 1, size);
Alexandre Vassalotti6e73ff12013-12-05 19:29:32 -08002073 len = 9;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002074 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002075 else {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002076 PyErr_SetString(PyExc_OverflowError,
Serhiy Storchakaf8def282013-02-16 17:29:56 +02002077 "cannot serialize a bytes object larger than 4 GiB");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002078 return -1; /* string too large */
2079 }
2080
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002081 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002082 return -1;
2083
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002084 if (_Pickler_Write(self, PyBytes_AS_STRING(obj), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002085 return -1;
2086
2087 if (memo_put(self, obj) < 0)
2088 return -1;
2089
2090 return 0;
2091 }
2092}
2093
2094/* A copy of PyUnicode_EncodeRawUnicodeEscape() that also translates
2095 backslash and newline characters to \uXXXX escapes. */
2096static PyObject *
Victor Stinnerc806fdc2011-09-29 23:50:23 +02002097raw_unicode_escape(PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002098{
Victor Stinner7270b7f2014-08-17 21:14:46 +02002099 PyObject *repr;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002100 char *p;
Victor Stinner049e5092014-08-17 22:20:00 +02002101 Py_ssize_t i, size;
2102 size_t expandsize;
Victor Stinnerc806fdc2011-09-29 23:50:23 +02002103 void *data;
2104 unsigned int kind;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002105
Victor Stinnerc806fdc2011-09-29 23:50:23 +02002106 if (PyUnicode_READY(obj))
2107 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002108
Victor Stinnerc806fdc2011-09-29 23:50:23 +02002109 size = PyUnicode_GET_LENGTH(obj);
2110 data = PyUnicode_DATA(obj);
2111 kind = PyUnicode_KIND(obj);
2112 if (kind == PyUnicode_4BYTE_KIND)
2113 expandsize = 10;
2114 else
2115 expandsize = 6;
Victor Stinner121aab42011-09-29 23:40:53 +02002116
Victor Stinner049e5092014-08-17 22:20:00 +02002117 if ((size_t)size > (size_t)PY_SSIZE_T_MAX / expandsize)
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00002118 return PyErr_NoMemory();
Victor Stinner7270b7f2014-08-17 21:14:46 +02002119 repr = PyBytes_FromStringAndSize(NULL, expandsize * size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002120 if (repr == NULL)
2121 return NULL;
2122 if (size == 0)
Victor Stinner7270b7f2014-08-17 21:14:46 +02002123 return repr;
2124 assert(Py_REFCNT(repr) == 1);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002125
Victor Stinner7270b7f2014-08-17 21:14:46 +02002126 p = PyBytes_AS_STRING(repr);
Victor Stinnerc806fdc2011-09-29 23:50:23 +02002127 for (i=0; i < size; i++) {
2128 Py_UCS4 ch = PyUnicode_READ(kind, data, i);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002129 /* Map 32-bit characters to '\Uxxxxxxxx' */
2130 if (ch >= 0x10000) {
2131 *p++ = '\\';
2132 *p++ = 'U';
Victor Stinnerf5cff562011-10-14 02:13:11 +02002133 *p++ = Py_hexdigits[(ch >> 28) & 0xf];
2134 *p++ = Py_hexdigits[(ch >> 24) & 0xf];
2135 *p++ = Py_hexdigits[(ch >> 20) & 0xf];
2136 *p++ = Py_hexdigits[(ch >> 16) & 0xf];
2137 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
2138 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
2139 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
2140 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002141 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002142 /* Map 16-bit characters to '\uxxxx' */
Victor Stinnerc806fdc2011-09-29 23:50:23 +02002143 else if (ch >= 256 || ch == '\\' || ch == '\n') {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002144 *p++ = '\\';
2145 *p++ = 'u';
Victor Stinnerf5cff562011-10-14 02:13:11 +02002146 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
2147 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
2148 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
2149 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002150 }
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00002151 /* Copy everything else as-is */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002152 else
2153 *p++ = (char) ch;
2154 }
Victor Stinner7270b7f2014-08-17 21:14:46 +02002155 size = p - PyBytes_AS_STRING(repr);
2156 if (_PyBytes_Resize(&repr, size) < 0)
2157 return NULL;
2158 return repr;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002159}
2160
2161static int
Antoine Pitrou299978d2013-04-07 17:38:11 +02002162write_utf8(PicklerObject *self, char *data, Py_ssize_t size)
2163{
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002164 char header[9];
2165 Py_ssize_t len;
Antoine Pitrou299978d2013-04-07 17:38:11 +02002166
Victor Stinnerf13c46c2014-08-17 21:05:55 +02002167 assert(size >= 0);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002168 if (size <= 0xff && self->proto >= 4) {
2169 header[0] = SHORT_BINUNICODE;
2170 header[1] = (unsigned char)(size & 0xff);
2171 len = 2;
2172 }
Victor Stinnerf13c46c2014-08-17 21:05:55 +02002173 else if ((size_t)size <= 0xffffffffUL) {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002174 header[0] = BINUNICODE;
2175 header[1] = (unsigned char)(size & 0xff);
2176 header[2] = (unsigned char)((size >> 8) & 0xff);
2177 header[3] = (unsigned char)((size >> 16) & 0xff);
2178 header[4] = (unsigned char)((size >> 24) & 0xff);
2179 len = 5;
2180 }
2181 else if (self->proto >= 4) {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002182 header[0] = BINUNICODE8;
Alexandre Vassalotti1048fb52013-11-25 11:35:46 -08002183 _write_size64(header + 1, size);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002184 len = 9;
2185 }
2186 else {
Antoine Pitrou299978d2013-04-07 17:38:11 +02002187 PyErr_SetString(PyExc_OverflowError,
Antoine Pitrou4b7b0f02013-04-07 23:46:52 +02002188 "cannot serialize a string larger than 4GiB");
Antoine Pitrou299978d2013-04-07 17:38:11 +02002189 return -1;
2190 }
Antoine Pitrou299978d2013-04-07 17:38:11 +02002191
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002192 if (_Pickler_Write(self, header, len) < 0)
Antoine Pitrou299978d2013-04-07 17:38:11 +02002193 return -1;
Antoine Pitrou299978d2013-04-07 17:38:11 +02002194 if (_Pickler_Write(self, data, size) < 0)
2195 return -1;
2196
2197 return 0;
2198}
2199
2200static int
2201write_unicode_binary(PicklerObject *self, PyObject *obj)
2202{
2203 PyObject *encoded = NULL;
2204 Py_ssize_t size;
2205 char *data;
2206 int r;
2207
2208 if (PyUnicode_READY(obj))
2209 return -1;
2210
2211 data = PyUnicode_AsUTF8AndSize(obj, &size);
2212 if (data != NULL)
2213 return write_utf8(self, data, size);
2214
2215 /* Issue #8383: for strings with lone surrogates, fallback on the
2216 "surrogatepass" error handler. */
2217 PyErr_Clear();
2218 encoded = PyUnicode_AsEncodedString(obj, "utf-8", "surrogatepass");
2219 if (encoded == NULL)
2220 return -1;
2221
2222 r = write_utf8(self, PyBytes_AS_STRING(encoded),
2223 PyBytes_GET_SIZE(encoded));
2224 Py_DECREF(encoded);
2225 return r;
2226}
2227
2228static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002229save_unicode(PicklerObject *self, PyObject *obj)
2230{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002231 if (self->bin) {
Antoine Pitrou299978d2013-04-07 17:38:11 +02002232 if (write_unicode_binary(self, obj) < 0)
2233 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002234 }
2235 else {
Antoine Pitrou299978d2013-04-07 17:38:11 +02002236 PyObject *encoded;
2237 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002238 const char unicode_op = UNICODE;
2239
Victor Stinnerc806fdc2011-09-29 23:50:23 +02002240 encoded = raw_unicode_escape(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002241 if (encoded == NULL)
Antoine Pitrou299978d2013-04-07 17:38:11 +02002242 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002243
Antoine Pitrou299978d2013-04-07 17:38:11 +02002244 if (_Pickler_Write(self, &unicode_op, 1) < 0) {
2245 Py_DECREF(encoded);
2246 return -1;
2247 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002248
2249 size = PyBytes_GET_SIZE(encoded);
Antoine Pitrou299978d2013-04-07 17:38:11 +02002250 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0) {
2251 Py_DECREF(encoded);
2252 return -1;
2253 }
2254 Py_DECREF(encoded);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002255
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002256 if (_Pickler_Write(self, "\n", 1) < 0)
Antoine Pitrou299978d2013-04-07 17:38:11 +02002257 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002258 }
2259 if (memo_put(self, obj) < 0)
Antoine Pitrou299978d2013-04-07 17:38:11 +02002260 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002261
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002262 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002263}
2264
2265/* A helper for save_tuple. Push the len elements in tuple t on the stack. */
2266static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002267store_tuple_elements(PicklerObject *self, PyObject *t, Py_ssize_t len)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002268{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002269 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002270
2271 assert(PyTuple_Size(t) == len);
2272
2273 for (i = 0; i < len; i++) {
2274 PyObject *element = PyTuple_GET_ITEM(t, i);
2275
2276 if (element == NULL)
2277 return -1;
2278 if (save(self, element, 0) < 0)
2279 return -1;
2280 }
2281
2282 return 0;
2283}
2284
2285/* Tuples are ubiquitous in the pickle protocols, so many techniques are
2286 * used across protocols to minimize the space needed to pickle them.
2287 * Tuples are also the only builtin immutable type that can be recursive
2288 * (a tuple can be reached from itself), and that requires some subtle
2289 * magic so that it works in all cases. IOW, this is a long routine.
2290 */
2291static int
2292save_tuple(PicklerObject *self, PyObject *obj)
2293{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002294 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002295
2296 const char mark_op = MARK;
2297 const char tuple_op = TUPLE;
2298 const char pop_op = POP;
2299 const char pop_mark_op = POP_MARK;
2300 const char len2opcode[] = {EMPTY_TUPLE, TUPLE1, TUPLE2, TUPLE3};
2301
2302 if ((len = PyTuple_Size(obj)) < 0)
2303 return -1;
2304
2305 if (len == 0) {
2306 char pdata[2];
2307
2308 if (self->proto) {
2309 pdata[0] = EMPTY_TUPLE;
2310 len = 1;
2311 }
2312 else {
2313 pdata[0] = MARK;
2314 pdata[1] = TUPLE;
2315 len = 2;
2316 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002317 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002318 return -1;
2319 return 0;
2320 }
2321
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002322 /* The tuple isn't in the memo now. If it shows up there after
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002323 * saving the tuple elements, the tuple must be recursive, in
2324 * which case we'll pop everything we put on the stack, and fetch
2325 * its value from the memo.
2326 */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002327 if (len <= 3 && self->proto >= 2) {
2328 /* Use TUPLE{1,2,3} opcodes. */
2329 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002330 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002331
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002332 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002333 /* pop the len elements */
2334 for (i = 0; i < len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002335 if (_Pickler_Write(self, &pop_op, 1) < 0)
2336 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002337 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002338 if (memo_get(self, obj) < 0)
2339 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002340
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002341 return 0;
2342 }
2343 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002344 if (_Pickler_Write(self, len2opcode + len, 1) < 0)
2345 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002346 }
2347 goto memoize;
2348 }
2349
2350 /* proto < 2 and len > 0, or proto >= 2 and len > 3.
2351 * Generate MARK e1 e2 ... TUPLE
2352 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002353 if (_Pickler_Write(self, &mark_op, 1) < 0)
2354 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002355
2356 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002357 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002358
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002359 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002360 /* pop the stack stuff we pushed */
2361 if (self->bin) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002362 if (_Pickler_Write(self, &pop_mark_op, 1) < 0)
2363 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002364 }
2365 else {
2366 /* Note that we pop one more than len, to remove
2367 * the MARK too.
2368 */
2369 for (i = 0; i <= len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002370 if (_Pickler_Write(self, &pop_op, 1) < 0)
2371 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002372 }
2373 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002374 if (memo_get(self, obj) < 0)
2375 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002376
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002377 return 0;
2378 }
2379 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002380 if (_Pickler_Write(self, &tuple_op, 1) < 0)
2381 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002382 }
2383
2384 memoize:
2385 if (memo_put(self, obj) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002386 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002387
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002388 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002389}
2390
2391/* iter is an iterator giving items, and we batch up chunks of
2392 * MARK item item ... item APPENDS
2393 * opcode sequences. Calling code should have arranged to first create an
2394 * empty list, or list-like object, for the APPENDS to operate on.
2395 * Returns 0 on success, <0 on error.
2396 */
2397static int
2398batch_list(PicklerObject *self, PyObject *iter)
2399{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002400 PyObject *obj = NULL;
2401 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002402 int i, n;
2403
2404 const char mark_op = MARK;
2405 const char append_op = APPEND;
2406 const char appends_op = APPENDS;
2407
2408 assert(iter != NULL);
2409
2410 /* XXX: I think this function could be made faster by avoiding the
2411 iterator interface and fetching objects directly from list using
2412 PyList_GET_ITEM.
2413 */
2414
2415 if (self->proto == 0) {
2416 /* APPENDS isn't available; do one at a time. */
2417 for (;;) {
2418 obj = PyIter_Next(iter);
2419 if (obj == NULL) {
2420 if (PyErr_Occurred())
2421 return -1;
2422 break;
2423 }
2424 i = save(self, obj, 0);
2425 Py_DECREF(obj);
2426 if (i < 0)
2427 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002428 if (_Pickler_Write(self, &append_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002429 return -1;
2430 }
2431 return 0;
2432 }
2433
2434 /* proto > 0: write in batches of BATCHSIZE. */
2435 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002436 /* Get first item */
2437 firstitem = PyIter_Next(iter);
2438 if (firstitem == NULL) {
2439 if (PyErr_Occurred())
2440 goto error;
2441
2442 /* nothing more to add */
2443 break;
2444 }
2445
2446 /* Try to get a second item */
2447 obj = PyIter_Next(iter);
2448 if (obj == NULL) {
2449 if (PyErr_Occurred())
2450 goto error;
2451
2452 /* Only one item to write */
2453 if (save(self, firstitem, 0) < 0)
2454 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002455 if (_Pickler_Write(self, &append_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002456 goto error;
2457 Py_CLEAR(firstitem);
2458 break;
2459 }
2460
2461 /* More than one item to write */
2462
2463 /* Pump out MARK, items, APPENDS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002464 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002465 goto error;
2466
2467 if (save(self, firstitem, 0) < 0)
2468 goto error;
2469 Py_CLEAR(firstitem);
2470 n = 1;
2471
2472 /* Fetch and save up to BATCHSIZE items */
2473 while (obj) {
2474 if (save(self, obj, 0) < 0)
2475 goto error;
2476 Py_CLEAR(obj);
2477 n += 1;
2478
2479 if (n == BATCHSIZE)
2480 break;
2481
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002482 obj = PyIter_Next(iter);
2483 if (obj == NULL) {
2484 if (PyErr_Occurred())
2485 goto error;
2486 break;
2487 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002488 }
2489
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002490 if (_Pickler_Write(self, &appends_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002491 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002492
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002493 } while (n == BATCHSIZE);
2494 return 0;
2495
2496 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002497 Py_XDECREF(firstitem);
2498 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002499 return -1;
2500}
2501
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002502/* This is a variant of batch_list() above, specialized for lists (with no
2503 * support for list subclasses). Like batch_list(), we batch up chunks of
2504 * MARK item item ... item APPENDS
2505 * opcode sequences. Calling code should have arranged to first create an
2506 * empty list, or list-like object, for the APPENDS to operate on.
2507 * Returns 0 on success, -1 on error.
2508 *
2509 * This version is considerably faster than batch_list(), if less general.
2510 *
2511 * Note that this only works for protocols > 0.
2512 */
2513static int
2514batch_list_exact(PicklerObject *self, PyObject *obj)
2515{
2516 PyObject *item = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002517 Py_ssize_t this_batch, total;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002518
2519 const char append_op = APPEND;
2520 const char appends_op = APPENDS;
2521 const char mark_op = MARK;
2522
2523 assert(obj != NULL);
2524 assert(self->proto > 0);
2525 assert(PyList_CheckExact(obj));
2526
2527 if (PyList_GET_SIZE(obj) == 1) {
2528 item = PyList_GET_ITEM(obj, 0);
2529 if (save(self, item, 0) < 0)
2530 return -1;
2531 if (_Pickler_Write(self, &append_op, 1) < 0)
2532 return -1;
2533 return 0;
2534 }
2535
2536 /* Write in batches of BATCHSIZE. */
2537 total = 0;
2538 do {
2539 this_batch = 0;
2540 if (_Pickler_Write(self, &mark_op, 1) < 0)
2541 return -1;
2542 while (total < PyList_GET_SIZE(obj)) {
2543 item = PyList_GET_ITEM(obj, total);
2544 if (save(self, item, 0) < 0)
2545 return -1;
2546 total++;
2547 if (++this_batch == BATCHSIZE)
2548 break;
2549 }
2550 if (_Pickler_Write(self, &appends_op, 1) < 0)
2551 return -1;
2552
2553 } while (total < PyList_GET_SIZE(obj));
2554
2555 return 0;
2556}
2557
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002558static int
2559save_list(PicklerObject *self, PyObject *obj)
2560{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002561 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002562 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002563 int status = 0;
2564
2565 if (self->fast && !fast_save_enter(self, obj))
2566 goto error;
2567
2568 /* Create an empty list. */
2569 if (self->bin) {
2570 header[0] = EMPTY_LIST;
2571 len = 1;
2572 }
2573 else {
2574 header[0] = MARK;
2575 header[1] = LIST;
2576 len = 2;
2577 }
2578
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002579 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002580 goto error;
2581
2582 /* Get list length, and bow out early if empty. */
2583 if ((len = PyList_Size(obj)) < 0)
2584 goto error;
2585
2586 if (memo_put(self, obj) < 0)
2587 goto error;
2588
2589 if (len != 0) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002590 /* Materialize the list elements. */
2591 if (PyList_CheckExact(obj) && self->proto > 0) {
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002592 if (Py_EnterRecursiveCall(" while pickling an object"))
2593 goto error;
2594 status = batch_list_exact(self, obj);
2595 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002596 } else {
2597 PyObject *iter = PyObject_GetIter(obj);
2598 if (iter == NULL)
2599 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002600
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002601 if (Py_EnterRecursiveCall(" while pickling an object")) {
2602 Py_DECREF(iter);
2603 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002604 }
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002605 status = batch_list(self, iter);
2606 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002607 Py_DECREF(iter);
2608 }
2609 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002610 if (0) {
2611 error:
2612 status = -1;
2613 }
2614
2615 if (self->fast && !fast_save_leave(self, obj))
2616 status = -1;
2617
2618 return status;
2619}
2620
2621/* iter is an iterator giving (key, value) pairs, and we batch up chunks of
2622 * MARK key value ... key value SETITEMS
2623 * opcode sequences. Calling code should have arranged to first create an
2624 * empty dict, or dict-like object, for the SETITEMS to operate on.
2625 * Returns 0 on success, <0 on error.
2626 *
2627 * This is very much like batch_list(). The difference between saving
2628 * elements directly, and picking apart two-tuples, is so long-winded at
2629 * the C level, though, that attempts to combine these routines were too
2630 * ugly to bear.
2631 */
2632static int
2633batch_dict(PicklerObject *self, PyObject *iter)
2634{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002635 PyObject *obj = NULL;
2636 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002637 int i, n;
2638
2639 const char mark_op = MARK;
2640 const char setitem_op = SETITEM;
2641 const char setitems_op = SETITEMS;
2642
2643 assert(iter != NULL);
2644
2645 if (self->proto == 0) {
2646 /* SETITEMS isn't available; do one at a time. */
2647 for (;;) {
2648 obj = PyIter_Next(iter);
2649 if (obj == NULL) {
2650 if (PyErr_Occurred())
2651 return -1;
2652 break;
2653 }
2654 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2655 PyErr_SetString(PyExc_TypeError, "dict items "
2656 "iterator must return 2-tuples");
2657 return -1;
2658 }
2659 i = save(self, PyTuple_GET_ITEM(obj, 0), 0);
2660 if (i >= 0)
2661 i = save(self, PyTuple_GET_ITEM(obj, 1), 0);
2662 Py_DECREF(obj);
2663 if (i < 0)
2664 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002665 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002666 return -1;
2667 }
2668 return 0;
2669 }
2670
2671 /* proto > 0: write in batches of BATCHSIZE. */
2672 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002673 /* Get first item */
2674 firstitem = PyIter_Next(iter);
2675 if (firstitem == NULL) {
2676 if (PyErr_Occurred())
2677 goto error;
2678
2679 /* nothing more to add */
2680 break;
2681 }
2682 if (!PyTuple_Check(firstitem) || PyTuple_Size(firstitem) != 2) {
2683 PyErr_SetString(PyExc_TypeError, "dict items "
2684 "iterator must return 2-tuples");
2685 goto error;
2686 }
2687
2688 /* Try to get a second item */
2689 obj = PyIter_Next(iter);
2690 if (obj == NULL) {
2691 if (PyErr_Occurred())
2692 goto error;
2693
2694 /* Only one item to write */
2695 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2696 goto error;
2697 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2698 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002699 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002700 goto error;
2701 Py_CLEAR(firstitem);
2702 break;
2703 }
2704
2705 /* More than one item to write */
2706
2707 /* Pump out MARK, items, SETITEMS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002708 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002709 goto error;
2710
2711 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2712 goto error;
2713 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2714 goto error;
2715 Py_CLEAR(firstitem);
2716 n = 1;
2717
2718 /* Fetch and save up to BATCHSIZE items */
2719 while (obj) {
2720 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2721 PyErr_SetString(PyExc_TypeError, "dict items "
2722 "iterator must return 2-tuples");
2723 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002724 }
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002725 if (save(self, PyTuple_GET_ITEM(obj, 0), 0) < 0 ||
2726 save(self, PyTuple_GET_ITEM(obj, 1), 0) < 0)
2727 goto error;
2728 Py_CLEAR(obj);
2729 n += 1;
2730
2731 if (n == BATCHSIZE)
2732 break;
2733
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002734 obj = PyIter_Next(iter);
2735 if (obj == NULL) {
2736 if (PyErr_Occurred())
2737 goto error;
2738 break;
2739 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002740 }
2741
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002742 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002743 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002744
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002745 } while (n == BATCHSIZE);
2746 return 0;
2747
2748 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002749 Py_XDECREF(firstitem);
2750 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002751 return -1;
2752}
2753
Collin Winter5c9b02d2009-05-25 05:43:30 +00002754/* This is a variant of batch_dict() above that specializes for dicts, with no
2755 * support for dict subclasses. Like batch_dict(), we batch up chunks of
2756 * MARK key value ... key value SETITEMS
2757 * opcode sequences. Calling code should have arranged to first create an
2758 * empty dict, or dict-like object, for the SETITEMS to operate on.
2759 * Returns 0 on success, -1 on error.
2760 *
2761 * Note that this currently doesn't work for protocol 0.
2762 */
2763static int
2764batch_dict_exact(PicklerObject *self, PyObject *obj)
2765{
2766 PyObject *key = NULL, *value = NULL;
2767 int i;
2768 Py_ssize_t dict_size, ppos = 0;
2769
Alexandre Vassalottif70b1292009-05-25 18:00:52 +00002770 const char mark_op = MARK;
2771 const char setitem_op = SETITEM;
2772 const char setitems_op = SETITEMS;
Collin Winter5c9b02d2009-05-25 05:43:30 +00002773
2774 assert(obj != NULL);
2775 assert(self->proto > 0);
2776
2777 dict_size = PyDict_Size(obj);
2778
2779 /* Special-case len(d) == 1 to save space. */
2780 if (dict_size == 1) {
2781 PyDict_Next(obj, &ppos, &key, &value);
2782 if (save(self, key, 0) < 0)
2783 return -1;
2784 if (save(self, value, 0) < 0)
2785 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002786 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002787 return -1;
2788 return 0;
2789 }
2790
2791 /* Write in batches of BATCHSIZE. */
2792 do {
2793 i = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002794 if (_Pickler_Write(self, &mark_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002795 return -1;
2796 while (PyDict_Next(obj, &ppos, &key, &value)) {
2797 if (save(self, key, 0) < 0)
2798 return -1;
2799 if (save(self, value, 0) < 0)
2800 return -1;
2801 if (++i == BATCHSIZE)
2802 break;
2803 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002804 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002805 return -1;
2806 if (PyDict_Size(obj) != dict_size) {
2807 PyErr_Format(
2808 PyExc_RuntimeError,
2809 "dictionary changed size during iteration");
2810 return -1;
2811 }
2812
2813 } while (i == BATCHSIZE);
2814 return 0;
2815}
2816
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002817static int
2818save_dict(PicklerObject *self, PyObject *obj)
2819{
2820 PyObject *items, *iter;
2821 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002822 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002823 int status = 0;
2824
2825 if (self->fast && !fast_save_enter(self, obj))
2826 goto error;
2827
2828 /* Create an empty dict. */
2829 if (self->bin) {
2830 header[0] = EMPTY_DICT;
2831 len = 1;
2832 }
2833 else {
2834 header[0] = MARK;
2835 header[1] = DICT;
2836 len = 2;
2837 }
2838
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002839 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002840 goto error;
2841
2842 /* Get dict size, and bow out early if empty. */
2843 if ((len = PyDict_Size(obj)) < 0)
2844 goto error;
2845
2846 if (memo_put(self, obj) < 0)
2847 goto error;
2848
2849 if (len != 0) {
2850 /* Save the dict items. */
Collin Winter5c9b02d2009-05-25 05:43:30 +00002851 if (PyDict_CheckExact(obj) && self->proto > 0) {
2852 /* We can take certain shortcuts if we know this is a dict and
2853 not a dict subclass. */
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002854 if (Py_EnterRecursiveCall(" while pickling an object"))
2855 goto error;
2856 status = batch_dict_exact(self, obj);
2857 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002858 } else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02002859 _Py_IDENTIFIER(items);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02002860
2861 items = _PyObject_CallMethodId(obj, &PyId_items, "()");
Collin Winter5c9b02d2009-05-25 05:43:30 +00002862 if (items == NULL)
2863 goto error;
2864 iter = PyObject_GetIter(items);
2865 Py_DECREF(items);
2866 if (iter == NULL)
2867 goto error;
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002868 if (Py_EnterRecursiveCall(" while pickling an object")) {
2869 Py_DECREF(iter);
2870 goto error;
2871 }
Collin Winter5c9b02d2009-05-25 05:43:30 +00002872 status = batch_dict(self, iter);
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002873 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002874 Py_DECREF(iter);
2875 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002876 }
2877
2878 if (0) {
2879 error:
2880 status = -1;
2881 }
2882
2883 if (self->fast && !fast_save_leave(self, obj))
2884 status = -1;
2885
2886 return status;
2887}
2888
2889static int
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002890save_set(PicklerObject *self, PyObject *obj)
2891{
2892 PyObject *item;
2893 int i;
2894 Py_ssize_t set_size, ppos = 0;
2895 Py_hash_t hash;
2896
2897 const char empty_set_op = EMPTY_SET;
2898 const char mark_op = MARK;
2899 const char additems_op = ADDITEMS;
2900
2901 if (self->proto < 4) {
2902 PyObject *items;
2903 PyObject *reduce_value;
2904 int status;
2905
2906 items = PySequence_List(obj);
2907 if (items == NULL) {
2908 return -1;
2909 }
2910 reduce_value = Py_BuildValue("(O(O))", (PyObject*)&PySet_Type, items);
2911 Py_DECREF(items);
2912 if (reduce_value == NULL) {
2913 return -1;
2914 }
2915 /* save_reduce() will memoize the object automatically. */
2916 status = save_reduce(self, reduce_value, obj);
2917 Py_DECREF(reduce_value);
2918 return status;
2919 }
2920
2921 if (_Pickler_Write(self, &empty_set_op, 1) < 0)
2922 return -1;
2923
2924 if (memo_put(self, obj) < 0)
2925 return -1;
2926
2927 set_size = PySet_GET_SIZE(obj);
2928 if (set_size == 0)
2929 return 0; /* nothing to do */
2930
2931 /* Write in batches of BATCHSIZE. */
2932 do {
2933 i = 0;
2934 if (_Pickler_Write(self, &mark_op, 1) < 0)
2935 return -1;
2936 while (_PySet_NextEntry(obj, &ppos, &item, &hash)) {
2937 if (save(self, item, 0) < 0)
2938 return -1;
2939 if (++i == BATCHSIZE)
2940 break;
2941 }
2942 if (_Pickler_Write(self, &additems_op, 1) < 0)
2943 return -1;
2944 if (PySet_GET_SIZE(obj) != set_size) {
2945 PyErr_Format(
2946 PyExc_RuntimeError,
2947 "set changed size during iteration");
2948 return -1;
2949 }
2950 } while (i == BATCHSIZE);
2951
2952 return 0;
2953}
2954
2955static int
2956save_frozenset(PicklerObject *self, PyObject *obj)
2957{
2958 PyObject *iter;
2959
2960 const char mark_op = MARK;
2961 const char frozenset_op = FROZENSET;
2962
2963 if (self->fast && !fast_save_enter(self, obj))
2964 return -1;
2965
2966 if (self->proto < 4) {
2967 PyObject *items;
2968 PyObject *reduce_value;
2969 int status;
2970
2971 items = PySequence_List(obj);
2972 if (items == NULL) {
2973 return -1;
2974 }
2975 reduce_value = Py_BuildValue("(O(O))", (PyObject*)&PyFrozenSet_Type,
2976 items);
2977 Py_DECREF(items);
2978 if (reduce_value == NULL) {
2979 return -1;
2980 }
2981 /* save_reduce() will memoize the object automatically. */
2982 status = save_reduce(self, reduce_value, obj);
2983 Py_DECREF(reduce_value);
2984 return status;
2985 }
2986
2987 if (_Pickler_Write(self, &mark_op, 1) < 0)
2988 return -1;
2989
2990 iter = PyObject_GetIter(obj);
Christian Heimesb3d3ee42013-11-23 21:01:40 +01002991 if (iter == NULL) {
Christian Heimes74d8d632013-11-23 21:05:31 +01002992 return -1;
Christian Heimesb3d3ee42013-11-23 21:01:40 +01002993 }
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01002994 for (;;) {
2995 PyObject *item;
2996
2997 item = PyIter_Next(iter);
2998 if (item == NULL) {
2999 if (PyErr_Occurred()) {
3000 Py_DECREF(iter);
3001 return -1;
3002 }
3003 break;
3004 }
3005 if (save(self, item, 0) < 0) {
3006 Py_DECREF(item);
3007 Py_DECREF(iter);
3008 return -1;
3009 }
3010 Py_DECREF(item);
3011 }
3012 Py_DECREF(iter);
3013
3014 /* If the object is already in the memo, this means it is
3015 recursive. In this case, throw away everything we put on the
3016 stack, and fetch the object back from the memo. */
3017 if (PyMemoTable_Get(self->memo, obj)) {
3018 const char pop_mark_op = POP_MARK;
3019
3020 if (_Pickler_Write(self, &pop_mark_op, 1) < 0)
3021 return -1;
3022 if (memo_get(self, obj) < 0)
3023 return -1;
3024 return 0;
3025 }
3026
3027 if (_Pickler_Write(self, &frozenset_op, 1) < 0)
3028 return -1;
3029 if (memo_put(self, obj) < 0)
3030 return -1;
3031
3032 return 0;
3033}
3034
3035static int
3036fix_imports(PyObject **module_name, PyObject **global_name)
3037{
3038 PyObject *key;
3039 PyObject *item;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003040 PickleState *st = _Pickle_GetGlobalState();
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003041
3042 key = PyTuple_Pack(2, *module_name, *global_name);
3043 if (key == NULL)
3044 return -1;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003045 item = PyDict_GetItemWithError(st->name_mapping_3to2, key);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003046 Py_DECREF(key);
3047 if (item) {
3048 PyObject *fixed_module_name;
3049 PyObject *fixed_global_name;
3050
3051 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
3052 PyErr_Format(PyExc_RuntimeError,
3053 "_compat_pickle.REVERSE_NAME_MAPPING values "
3054 "should be 2-tuples, not %.200s",
3055 Py_TYPE(item)->tp_name);
3056 return -1;
3057 }
3058 fixed_module_name = PyTuple_GET_ITEM(item, 0);
3059 fixed_global_name = PyTuple_GET_ITEM(item, 1);
3060 if (!PyUnicode_Check(fixed_module_name) ||
3061 !PyUnicode_Check(fixed_global_name)) {
3062 PyErr_Format(PyExc_RuntimeError,
3063 "_compat_pickle.REVERSE_NAME_MAPPING values "
3064 "should be pairs of str, not (%.200s, %.200s)",
3065 Py_TYPE(fixed_module_name)->tp_name,
3066 Py_TYPE(fixed_global_name)->tp_name);
3067 return -1;
3068 }
3069
3070 Py_CLEAR(*module_name);
3071 Py_CLEAR(*global_name);
3072 Py_INCREF(fixed_module_name);
3073 Py_INCREF(fixed_global_name);
3074 *module_name = fixed_module_name;
3075 *global_name = fixed_global_name;
3076 }
3077 else if (PyErr_Occurred()) {
3078 return -1;
3079 }
3080
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003081 item = PyDict_GetItemWithError(st->import_mapping_3to2, *module_name);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003082 if (item) {
3083 if (!PyUnicode_Check(item)) {
3084 PyErr_Format(PyExc_RuntimeError,
3085 "_compat_pickle.REVERSE_IMPORT_MAPPING values "
3086 "should be strings, not %.200s",
3087 Py_TYPE(item)->tp_name);
3088 return -1;
3089 }
3090 Py_CLEAR(*module_name);
3091 Py_INCREF(item);
3092 *module_name = item;
3093 }
3094 else if (PyErr_Occurred()) {
3095 return -1;
3096 }
3097
3098 return 0;
3099}
3100
3101static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003102save_global(PicklerObject *self, PyObject *obj, PyObject *name)
3103{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003104 PyObject *global_name = NULL;
3105 PyObject *module_name = NULL;
3106 PyObject *module = NULL;
3107 PyObject *cls;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003108 PickleState *st = _Pickle_GetGlobalState();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003109 int status = 0;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003110 _Py_IDENTIFIER(__name__);
3111 _Py_IDENTIFIER(__qualname__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003112
3113 const char global_op = GLOBAL;
3114
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003115 if (name) {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003116 Py_INCREF(name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003117 global_name = name;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003118 }
3119 else {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003120 if (self->proto >= 4) {
3121 global_name = _PyObject_GetAttrId(obj, &PyId___qualname__);
3122 if (global_name == NULL) {
3123 if (!PyErr_ExceptionMatches(PyExc_AttributeError))
3124 goto error;
3125 PyErr_Clear();
3126 }
3127 }
3128 if (global_name == NULL) {
3129 global_name = _PyObject_GetAttrId(obj, &PyId___name__);
3130 if (global_name == NULL)
3131 goto error;
3132 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003133 }
3134
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003135 module_name = whichmodule(obj, global_name, self->proto >= 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003136 if (module_name == NULL)
3137 goto error;
3138
3139 /* XXX: Change to use the import C API directly with level=0 to disallow
3140 relative imports.
3141
3142 XXX: PyImport_ImportModuleLevel could be used. However, this bypasses
3143 builtins.__import__. Therefore, _pickle, unlike pickle.py, will ignore
3144 custom import functions (IMHO, this would be a nice security
3145 feature). The import C API would need to be extended to support the
3146 extra parameters of __import__ to fix that. */
3147 module = PyImport_Import(module_name);
3148 if (module == NULL) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003149 PyErr_Format(st->PicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003150 "Can't pickle %R: import of module %R failed",
3151 obj, module_name);
3152 goto error;
3153 }
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003154 cls = getattribute(module, global_name, self->proto >= 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003155 if (cls == NULL) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003156 PyErr_Format(st->PicklingError,
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003157 "Can't pickle %R: attribute lookup %S on %S failed",
3158 obj, global_name, module_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003159 goto error;
3160 }
3161 if (cls != obj) {
3162 Py_DECREF(cls);
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003163 PyErr_Format(st->PicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003164 "Can't pickle %R: it's not the same object as %S.%S",
3165 obj, module_name, global_name);
3166 goto error;
3167 }
3168 Py_DECREF(cls);
3169
3170 if (self->proto >= 2) {
3171 /* See whether this is in the extension registry, and if
3172 * so generate an EXT opcode.
3173 */
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003174 PyObject *extension_key;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003175 PyObject *code_obj; /* extension code as Python object */
3176 long code; /* extension code as C value */
3177 char pdata[5];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003178 Py_ssize_t n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003179
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003180 extension_key = PyTuple_Pack(2, module_name, global_name);
3181 if (extension_key == NULL) {
3182 goto error;
3183 }
Alexandre Vassalotti567eba12013-11-28 17:09:16 -08003184 code_obj = PyDict_GetItemWithError(st->extension_registry,
3185 extension_key);
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003186 Py_DECREF(extension_key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003187 /* The object is not registered in the extension registry.
3188 This is the most likely code path. */
Alexandre Vassalotti567eba12013-11-28 17:09:16 -08003189 if (code_obj == NULL) {
3190 if (PyErr_Occurred()) {
3191 goto error;
3192 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003193 goto gen_global;
Alexandre Vassalotti567eba12013-11-28 17:09:16 -08003194 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003195
3196 /* XXX: pickle.py doesn't check neither the type, nor the range
3197 of the value returned by the extension_registry. It should for
3198 consistency. */
3199
3200 /* Verify code_obj has the right type and value. */
3201 if (!PyLong_Check(code_obj)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003202 PyErr_Format(st->PicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003203 "Can't pickle %R: extension code %R isn't an integer",
3204 obj, code_obj);
3205 goto error;
3206 }
3207 code = PyLong_AS_LONG(code_obj);
3208 if (code <= 0 || code > 0x7fffffffL) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003209 if (!PyErr_Occurred())
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003210 PyErr_Format(st->PicklingError, "Can't pickle %R: extension "
3211 "code %ld is out of range", obj, code);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003212 goto error;
3213 }
3214
3215 /* Generate an EXT opcode. */
3216 if (code <= 0xff) {
3217 pdata[0] = EXT1;
3218 pdata[1] = (unsigned char)code;
3219 n = 2;
3220 }
3221 else if (code <= 0xffff) {
3222 pdata[0] = EXT2;
3223 pdata[1] = (unsigned char)(code & 0xff);
3224 pdata[2] = (unsigned char)((code >> 8) & 0xff);
3225 n = 3;
3226 }
3227 else {
3228 pdata[0] = EXT4;
3229 pdata[1] = (unsigned char)(code & 0xff);
3230 pdata[2] = (unsigned char)((code >> 8) & 0xff);
3231 pdata[3] = (unsigned char)((code >> 16) & 0xff);
3232 pdata[4] = (unsigned char)((code >> 24) & 0xff);
3233 n = 5;
3234 }
3235
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003236 if (_Pickler_Write(self, pdata, n) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003237 goto error;
3238 }
3239 else {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003240 gen_global:
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003241 if (self->proto >= 4) {
3242 const char stack_global_op = STACK_GLOBAL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003243
Christian Heimese8b1ba12013-11-23 21:13:39 +01003244 if (save(self, module_name, 0) < 0)
3245 goto error;
3246 if (save(self, global_name, 0) < 0)
3247 goto error;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003248
3249 if (_Pickler_Write(self, &stack_global_op, 1) < 0)
3250 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003251 }
3252 else {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003253 /* Generate a normal global opcode if we are using a pickle
3254 protocol < 4, or if the object is not registered in the
3255 extension registry. */
3256 PyObject *encoded;
3257 PyObject *(*unicode_encoder)(PyObject *);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003258
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003259 if (_Pickler_Write(self, &global_op, 1) < 0)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003260 goto error;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003261
3262 /* For protocol < 3 and if the user didn't request against doing
3263 so, we convert module names to the old 2.x module names. */
3264 if (self->proto < 3 && self->fix_imports) {
3265 if (fix_imports(&module_name, &global_name) < 0) {
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003266 goto error;
3267 }
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003268 }
3269
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003270 /* Since Python 3.0 now supports non-ASCII identifiers, we encode
3271 both the module name and the global name using UTF-8. We do so
3272 only when we are using the pickle protocol newer than version
3273 3. This is to ensure compatibility with older Unpickler running
3274 on Python 2.x. */
3275 if (self->proto == 3) {
3276 unicode_encoder = PyUnicode_AsUTF8String;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003277 }
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003278 else {
3279 unicode_encoder = PyUnicode_AsASCIIString;
3280 }
3281 encoded = unicode_encoder(module_name);
3282 if (encoded == NULL) {
3283 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003284 PyErr_Format(st->PicklingError,
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003285 "can't pickle module identifier '%S' using "
3286 "pickle protocol %i",
3287 module_name, self->proto);
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003288 goto error;
3289 }
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003290 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
3291 PyBytes_GET_SIZE(encoded)) < 0) {
3292 Py_DECREF(encoded);
3293 goto error;
3294 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003295 Py_DECREF(encoded);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003296 if(_Pickler_Write(self, "\n", 1) < 0)
3297 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003298
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003299 /* Save the name of the module. */
3300 encoded = unicode_encoder(global_name);
3301 if (encoded == NULL) {
3302 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003303 PyErr_Format(st->PicklingError,
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003304 "can't pickle global identifier '%S' using "
3305 "pickle protocol %i",
3306 global_name, self->proto);
3307 goto error;
3308 }
3309 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
3310 PyBytes_GET_SIZE(encoded)) < 0) {
3311 Py_DECREF(encoded);
3312 goto error;
3313 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003314 Py_DECREF(encoded);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003315 if (_Pickler_Write(self, "\n", 1) < 0)
3316 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003317 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003318 /* Memoize the object. */
3319 if (memo_put(self, obj) < 0)
3320 goto error;
3321 }
3322
3323 if (0) {
3324 error:
3325 status = -1;
3326 }
3327 Py_XDECREF(module_name);
3328 Py_XDECREF(global_name);
3329 Py_XDECREF(module);
3330
3331 return status;
3332}
3333
3334static int
Alexandre Vassalotti19b6fa62013-11-30 16:06:39 -08003335save_singleton_type(PicklerObject *self, PyObject *obj, PyObject *singleton)
3336{
3337 PyObject *reduce_value;
3338 int status;
3339
3340 reduce_value = Py_BuildValue("O(O)", &PyType_Type, singleton);
3341 if (reduce_value == NULL) {
3342 return -1;
3343 }
3344 status = save_reduce(self, reduce_value, obj);
3345 Py_DECREF(reduce_value);
3346 return status;
3347}
3348
3349static int
3350save_type(PicklerObject *self, PyObject *obj)
3351{
Alexandre Vassalotti65846c62013-11-30 17:55:48 -08003352 if (obj == (PyObject *)&_PyNone_Type) {
Alexandre Vassalotti19b6fa62013-11-30 16:06:39 -08003353 return save_singleton_type(self, obj, Py_None);
3354 }
3355 else if (obj == (PyObject *)&PyEllipsis_Type) {
3356 return save_singleton_type(self, obj, Py_Ellipsis);
3357 }
Alexandre Vassalotti65846c62013-11-30 17:55:48 -08003358 else if (obj == (PyObject *)&_PyNotImplemented_Type) {
Alexandre Vassalotti19b6fa62013-11-30 16:06:39 -08003359 return save_singleton_type(self, obj, Py_NotImplemented);
3360 }
3361 return save_global(self, obj, NULL);
3362}
3363
3364static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003365save_pers(PicklerObject *self, PyObject *obj, PyObject *func)
3366{
3367 PyObject *pid = NULL;
3368 int status = 0;
3369
3370 const char persid_op = PERSID;
3371 const char binpersid_op = BINPERSID;
3372
3373 Py_INCREF(obj);
Alexandre Vassalotti20c28c12013-11-27 02:26:54 -08003374 pid = _Pickle_FastCall(func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003375 if (pid == NULL)
3376 return -1;
3377
3378 if (pid != Py_None) {
3379 if (self->bin) {
3380 if (save(self, pid, 1) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003381 _Pickler_Write(self, &binpersid_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003382 goto error;
3383 }
3384 else {
3385 PyObject *pid_str = NULL;
3386 char *pid_ascii_bytes;
3387 Py_ssize_t size;
3388
3389 pid_str = PyObject_Str(pid);
3390 if (pid_str == NULL)
3391 goto error;
3392
3393 /* XXX: Should it check whether the persistent id only contains
3394 ASCII characters? And what if the pid contains embedded
3395 newlines? */
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00003396 pid_ascii_bytes = _PyUnicode_AsStringAndSize(pid_str, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003397 Py_DECREF(pid_str);
3398 if (pid_ascii_bytes == NULL)
3399 goto error;
3400
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003401 if (_Pickler_Write(self, &persid_op, 1) < 0 ||
3402 _Pickler_Write(self, pid_ascii_bytes, size) < 0 ||
3403 _Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003404 goto error;
3405 }
3406 status = 1;
3407 }
3408
3409 if (0) {
3410 error:
3411 status = -1;
3412 }
3413 Py_XDECREF(pid);
3414
3415 return status;
3416}
3417
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003418static PyObject *
3419get_class(PyObject *obj)
3420{
3421 PyObject *cls;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003422 _Py_IDENTIFIER(__class__);
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003423
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003424 cls = _PyObject_GetAttrId(obj, &PyId___class__);
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003425 if (cls == NULL) {
3426 if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
3427 PyErr_Clear();
3428 cls = (PyObject *) Py_TYPE(obj);
3429 Py_INCREF(cls);
3430 }
3431 }
3432 return cls;
3433}
3434
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003435/* We're saving obj, and args is the 2-thru-5 tuple returned by the
3436 * appropriate __reduce__ method for obj.
3437 */
3438static int
3439save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
3440{
3441 PyObject *callable;
3442 PyObject *argtup;
3443 PyObject *state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003444 PyObject *listitems = Py_None;
3445 PyObject *dictitems = Py_None;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003446 PickleState *st = _Pickle_GetGlobalState();
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00003447 Py_ssize_t size;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003448 int use_newobj = 0, use_newobj_ex = 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003449
3450 const char reduce_op = REDUCE;
3451 const char build_op = BUILD;
3452 const char newobj_op = NEWOBJ;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003453 const char newobj_ex_op = NEWOBJ_EX;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003454
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00003455 size = PyTuple_Size(args);
3456 if (size < 2 || size > 5) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003457 PyErr_SetString(st->PicklingError, "tuple returned by "
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00003458 "__reduce__ must contain 2 through 5 elements");
3459 return -1;
3460 }
3461
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003462 if (!PyArg_UnpackTuple(args, "save_reduce", 2, 5,
3463 &callable, &argtup, &state, &listitems, &dictitems))
3464 return -1;
3465
3466 if (!PyCallable_Check(callable)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003467 PyErr_SetString(st->PicklingError, "first item of the tuple "
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003468 "returned by __reduce__ must be callable");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003469 return -1;
3470 }
3471 if (!PyTuple_Check(argtup)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003472 PyErr_SetString(st->PicklingError, "second item of the tuple "
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003473 "returned by __reduce__ must be a tuple");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003474 return -1;
3475 }
3476
3477 if (state == Py_None)
3478 state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003479
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003480 if (listitems == Py_None)
3481 listitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003482 else if (!PyIter_Check(listitems)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003483 PyErr_Format(st->PicklingError, "fourth element of the tuple "
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003484 "returned by __reduce__ must be an iterator, not %s",
3485 Py_TYPE(listitems)->tp_name);
3486 return -1;
3487 }
3488
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003489 if (dictitems == Py_None)
3490 dictitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003491 else if (!PyIter_Check(dictitems)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003492 PyErr_Format(st->PicklingError, "fifth element of the tuple "
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003493 "returned by __reduce__ must be an iterator, not %s",
3494 Py_TYPE(dictitems)->tp_name);
3495 return -1;
3496 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003497
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003498 if (self->proto >= 2) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003499 PyObject *name;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003500 _Py_IDENTIFIER(__name__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003501
Victor Stinner804e05e2013-11-14 01:26:17 +01003502 name = _PyObject_GetAttrId(callable, &PyId___name__);
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003503 if (name == NULL) {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003504 if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003505 return -1;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003506 }
3507 PyErr_Clear();
3508 }
Serhiy Storchaka707b5cc2014-12-16 19:43:46 +02003509 else if (PyUnicode_Check(name)) {
3510 if (self->proto >= 4) {
3511 _Py_IDENTIFIER(__newobj_ex__);
3512 use_newobj_ex = PyUnicode_Compare(
3513 name, _PyUnicode_FromId(&PyId___newobj_ex__)) == 0;
3514 }
3515 if (!use_newobj_ex) {
3516 _Py_IDENTIFIER(__newobj__);
3517 use_newobj = PyUnicode_Compare(
3518 name, _PyUnicode_FromId(&PyId___newobj__)) == 0;
3519 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003520 }
Serhiy Storchaka707b5cc2014-12-16 19:43:46 +02003521 Py_XDECREF(name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003522 }
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003523
3524 if (use_newobj_ex) {
3525 PyObject *cls;
3526 PyObject *args;
3527 PyObject *kwargs;
3528
3529 if (Py_SIZE(argtup) != 3) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003530 PyErr_Format(st->PicklingError,
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003531 "length of the NEWOBJ_EX argument tuple must be "
3532 "exactly 3, not %zd", Py_SIZE(argtup));
3533 return -1;
3534 }
3535
3536 cls = PyTuple_GET_ITEM(argtup, 0);
3537 if (!PyType_Check(cls)) {
Larry Hastings61272b72014-01-07 12:41:53 -08003538 PyErr_Format(st->PicklingError,
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003539 "first item from NEWOBJ_EX argument tuple must "
3540 "be a class, not %.200s", Py_TYPE(cls)->tp_name);
3541 return -1;
3542 }
3543 args = PyTuple_GET_ITEM(argtup, 1);
3544 if (!PyTuple_Check(args)) {
Larry Hastings61272b72014-01-07 12:41:53 -08003545 PyErr_Format(st->PicklingError,
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003546 "second item from NEWOBJ_EX argument tuple must "
3547 "be a tuple, not %.200s", Py_TYPE(args)->tp_name);
3548 return -1;
3549 }
3550 kwargs = PyTuple_GET_ITEM(argtup, 2);
3551 if (!PyDict_Check(kwargs)) {
Larry Hastings61272b72014-01-07 12:41:53 -08003552 PyErr_Format(st->PicklingError,
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003553 "third item from NEWOBJ_EX argument tuple must "
3554 "be a dict, not %.200s", Py_TYPE(kwargs)->tp_name);
3555 return -1;
3556 }
3557
3558 if (save(self, cls, 0) < 0 ||
3559 save(self, args, 0) < 0 ||
3560 save(self, kwargs, 0) < 0 ||
3561 _Pickler_Write(self, &newobj_ex_op, 1) < 0) {
3562 return -1;
3563 }
3564 }
3565 else if (use_newobj) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003566 PyObject *cls;
3567 PyObject *newargtup;
3568 PyObject *obj_class;
3569 int p;
3570
3571 /* Sanity checks. */
3572 if (Py_SIZE(argtup) < 1) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003573 PyErr_SetString(st->PicklingError, "__newobj__ arglist is empty");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003574 return -1;
3575 }
3576
3577 cls = PyTuple_GET_ITEM(argtup, 0);
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003578 if (!PyType_Check(cls)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003579 PyErr_SetString(st->PicklingError, "args[0] from "
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003580 "__newobj__ args is not a type");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003581 return -1;
3582 }
3583
3584 if (obj != NULL) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003585 obj_class = get_class(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003586 p = obj_class != cls; /* true iff a problem */
3587 Py_DECREF(obj_class);
3588 if (p) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003589 PyErr_SetString(st->PicklingError, "args[0] from "
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003590 "__newobj__ args has the wrong class");
3591 return -1;
3592 }
3593 }
3594 /* XXX: These calls save() are prone to infinite recursion. Imagine
3595 what happen if the value returned by the __reduce__() method of
3596 some extension type contains another object of the same type. Ouch!
3597
3598 Here is a quick example, that I ran into, to illustrate what I
3599 mean:
3600
3601 >>> import pickle, copyreg
3602 >>> copyreg.dispatch_table.pop(complex)
3603 >>> pickle.dumps(1+2j)
3604 Traceback (most recent call last):
3605 ...
3606 RuntimeError: maximum recursion depth exceeded
3607
3608 Removing the complex class from copyreg.dispatch_table made the
3609 __reduce_ex__() method emit another complex object:
3610
3611 >>> (1+1j).__reduce_ex__(2)
3612 (<function __newobj__ at 0xb7b71c3c>,
3613 (<class 'complex'>, (1+1j)), None, None, None)
3614
3615 Thus when save() was called on newargstup (the 2nd item) recursion
3616 ensued. Of course, the bug was in the complex class which had a
3617 broken __getnewargs__() that emitted another complex object. But,
3618 the point, here, is it is quite easy to end up with a broken reduce
3619 function. */
3620
3621 /* Save the class and its __new__ arguments. */
3622 if (save(self, cls, 0) < 0)
3623 return -1;
3624
3625 newargtup = PyTuple_GetSlice(argtup, 1, Py_SIZE(argtup));
3626 if (newargtup == NULL)
3627 return -1;
3628
3629 p = save(self, newargtup, 0);
3630 Py_DECREF(newargtup);
3631 if (p < 0)
3632 return -1;
3633
3634 /* Add NEWOBJ opcode. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003635 if (_Pickler_Write(self, &newobj_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003636 return -1;
3637 }
3638 else { /* Not using NEWOBJ. */
3639 if (save(self, callable, 0) < 0 ||
3640 save(self, argtup, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003641 _Pickler_Write(self, &reduce_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003642 return -1;
3643 }
3644
3645 /* obj can be NULL when save_reduce() is used directly. A NULL obj means
3646 the caller do not want to memoize the object. Not particularly useful,
3647 but that is to mimic the behavior save_reduce() in pickle.py when
3648 obj is None. */
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003649 if (obj != NULL) {
3650 /* If the object is already in the memo, this means it is
3651 recursive. In this case, throw away everything we put on the
3652 stack, and fetch the object back from the memo. */
3653 if (PyMemoTable_Get(self->memo, obj)) {
3654 const char pop_op = POP;
3655
3656 if (_Pickler_Write(self, &pop_op, 1) < 0)
3657 return -1;
3658 if (memo_get(self, obj) < 0)
3659 return -1;
3660
3661 return 0;
3662 }
3663 else if (memo_put(self, obj) < 0)
3664 return -1;
3665 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003666
3667 if (listitems && batch_list(self, listitems) < 0)
3668 return -1;
3669
3670 if (dictitems && batch_dict(self, dictitems) < 0)
3671 return -1;
3672
3673 if (state) {
Victor Stinner121aab42011-09-29 23:40:53 +02003674 if (save(self, state, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003675 _Pickler_Write(self, &build_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003676 return -1;
3677 }
3678
3679 return 0;
3680}
3681
3682static int
3683save(PicklerObject *self, PyObject *obj, int pers_save)
3684{
3685 PyTypeObject *type;
3686 PyObject *reduce_func = NULL;
3687 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003688 int status = 0;
3689
Alexandre Vassalottib6a2f2a2013-11-23 20:30:03 -08003690 if (_Pickler_OpcodeBoundary(self) < 0)
3691 return -1;
3692
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00003693 if (Py_EnterRecursiveCall(" while pickling an object"))
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003694 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003695
3696 /* The extra pers_save argument is necessary to avoid calling save_pers()
3697 on its returned object. */
3698 if (!pers_save && self->pers_func) {
3699 /* save_pers() returns:
3700 -1 to signal an error;
3701 0 if it did nothing successfully;
3702 1 if a persistent id was saved.
3703 */
3704 if ((status = save_pers(self, obj, self->pers_func)) != 0)
3705 goto done;
3706 }
3707
3708 type = Py_TYPE(obj);
3709
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003710 /* The old cPickle had an optimization that used switch-case statement
3711 dispatching on the first letter of the type name. This has was removed
3712 since benchmarks shown that this optimization was actually slowing
3713 things down. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003714
3715 /* Atom types; these aren't memoized, so don't check the memo. */
3716
3717 if (obj == Py_None) {
3718 status = save_none(self, obj);
3719 goto done;
3720 }
3721 else if (obj == Py_False || obj == Py_True) {
3722 status = save_bool(self, obj);
3723 goto done;
3724 }
3725 else if (type == &PyLong_Type) {
3726 status = save_long(self, obj);
3727 goto done;
3728 }
3729 else if (type == &PyFloat_Type) {
3730 status = save_float(self, obj);
3731 goto done;
3732 }
3733
3734 /* Check the memo to see if it has the object. If so, generate
3735 a GET (or BINGET) opcode, instead of pickling the object
3736 once again. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003737 if (PyMemoTable_Get(self->memo, obj)) {
3738 if (memo_get(self, obj) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003739 goto error;
3740 goto done;
3741 }
3742
3743 if (type == &PyBytes_Type) {
3744 status = save_bytes(self, obj);
3745 goto done;
3746 }
3747 else if (type == &PyUnicode_Type) {
3748 status = save_unicode(self, obj);
3749 goto done;
3750 }
3751 else if (type == &PyDict_Type) {
3752 status = save_dict(self, obj);
3753 goto done;
3754 }
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003755 else if (type == &PySet_Type) {
3756 status = save_set(self, obj);
3757 goto done;
3758 }
3759 else if (type == &PyFrozenSet_Type) {
3760 status = save_frozenset(self, obj);
3761 goto done;
3762 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003763 else if (type == &PyList_Type) {
3764 status = save_list(self, obj);
3765 goto done;
3766 }
3767 else if (type == &PyTuple_Type) {
3768 status = save_tuple(self, obj);
3769 goto done;
3770 }
3771 else if (type == &PyType_Type) {
Alexandre Vassalotti19b6fa62013-11-30 16:06:39 -08003772 status = save_type(self, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003773 goto done;
3774 }
3775 else if (type == &PyFunction_Type) {
3776 status = save_global(self, obj, NULL);
Alexandre Vassalottifc912852013-11-24 03:07:35 -08003777 goto done;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003778 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003779
3780 /* XXX: This part needs some unit tests. */
3781
3782 /* Get a reduction callable, and call it. This may come from
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003783 * self.dispatch_table, copyreg.dispatch_table, the object's
3784 * __reduce_ex__ method, or the object's __reduce__ method.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003785 */
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003786 if (self->dispatch_table == NULL) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003787 PickleState *st = _Pickle_GetGlobalState();
Alexandre Vassalotti567eba12013-11-28 17:09:16 -08003788 reduce_func = PyDict_GetItemWithError(st->dispatch_table,
3789 (PyObject *)type);
3790 if (reduce_func == NULL) {
3791 if (PyErr_Occurred()) {
3792 goto error;
3793 }
3794 } else {
3795 /* PyDict_GetItemWithError() returns a borrowed reference.
3796 Increase the reference count to be consistent with
3797 PyObject_GetItem and _PyObject_GetAttrId used below. */
3798 Py_INCREF(reduce_func);
3799 }
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003800 } else {
Alexandre Vassalotti567eba12013-11-28 17:09:16 -08003801 reduce_func = PyObject_GetItem(self->dispatch_table,
3802 (PyObject *)type);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003803 if (reduce_func == NULL) {
3804 if (PyErr_ExceptionMatches(PyExc_KeyError))
3805 PyErr_Clear();
3806 else
3807 goto error;
3808 }
3809 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003810 if (reduce_func != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003811 Py_INCREF(obj);
Alexandre Vassalotti20c28c12013-11-27 02:26:54 -08003812 reduce_value = _Pickle_FastCall(reduce_func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003813 }
Antoine Pitrouffd41d92011-10-04 09:23:04 +02003814 else if (PyType_IsSubtype(type, &PyType_Type)) {
3815 status = save_global(self, obj, NULL);
3816 goto done;
3817 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003818 else {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003819 _Py_IDENTIFIER(__reduce__);
3820 _Py_IDENTIFIER(__reduce_ex__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003821
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003822
3823 /* XXX: If the __reduce__ method is defined, __reduce_ex__ is
3824 automatically defined as __reduce__. While this is convenient, this
3825 make it impossible to know which method was actually called. Of
3826 course, this is not a big deal. But still, it would be nice to let
3827 the user know which method was called when something go
3828 wrong. Incidentally, this means if __reduce_ex__ is not defined, we
3829 don't actually have to check for a __reduce__ method. */
3830
3831 /* Check for a __reduce_ex__ method. */
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003832 reduce_func = _PyObject_GetAttrId(obj, &PyId___reduce_ex__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003833 if (reduce_func != NULL) {
3834 PyObject *proto;
3835 proto = PyLong_FromLong(self->proto);
3836 if (proto != NULL) {
Alexandre Vassalotti20c28c12013-11-27 02:26:54 -08003837 reduce_value = _Pickle_FastCall(reduce_func, proto);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003838 }
3839 }
3840 else {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003841 PickleState *st = _Pickle_GetGlobalState();
3842
3843 if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003844 PyErr_Clear();
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003845 }
3846 else {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003847 goto error;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003848 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003849 /* Check for a __reduce__ method. */
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003850 reduce_func = _PyObject_GetAttrId(obj, &PyId___reduce__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003851 if (reduce_func != NULL) {
Alexandre Vassalotti6bf41e52013-11-28 15:17:29 -08003852 PyObject *empty_tuple = PyTuple_New(0);
3853 reduce_value = PyObject_Call(reduce_func, empty_tuple,
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003854 NULL);
Alexandre Vassalotti6bf41e52013-11-28 15:17:29 -08003855 Py_DECREF(empty_tuple);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003856 }
3857 else {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003858 PyErr_Format(st->PicklingError,
3859 "can't pickle '%.200s' object: %R",
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003860 type->tp_name, obj);
3861 goto error;
3862 }
3863 }
3864 }
3865
3866 if (reduce_value == NULL)
3867 goto error;
3868
3869 if (PyUnicode_Check(reduce_value)) {
3870 status = save_global(self, obj, reduce_value);
3871 goto done;
3872 }
3873
3874 if (!PyTuple_Check(reduce_value)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003875 PickleState *st = _Pickle_GetGlobalState();
3876 PyErr_SetString(st->PicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003877 "__reduce__ must return a string or tuple");
3878 goto error;
3879 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003880
3881 status = save_reduce(self, reduce_value, obj);
3882
3883 if (0) {
3884 error:
3885 status = -1;
3886 }
3887 done:
Alexandre Vassalottib6a2f2a2013-11-23 20:30:03 -08003888
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003889 Py_LeaveRecursiveCall();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003890 Py_XDECREF(reduce_func);
3891 Py_XDECREF(reduce_value);
3892
3893 return status;
3894}
3895
3896static int
3897dump(PicklerObject *self, PyObject *obj)
3898{
3899 const char stop_op = STOP;
3900
3901 if (self->proto >= 2) {
3902 char header[2];
3903
3904 header[0] = PROTO;
3905 assert(self->proto >= 0 && self->proto < 256);
3906 header[1] = (unsigned char)self->proto;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003907 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003908 return -1;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01003909 if (self->proto >= 4)
3910 self->framing = 1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003911 }
3912
3913 if (save(self, obj, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003914 _Pickler_Write(self, &stop_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003915 return -1;
3916
3917 return 0;
3918}
3919
Larry Hastings61272b72014-01-07 12:41:53 -08003920/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08003921
3922_pickle.Pickler.clear_memo
3923
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08003924Clears the pickler's "memo".
3925
3926The memo is the data structure that remembers which objects the
3927pickler has already seen, so that shared or recursive objects are
3928pickled by reference and not by value. This method is useful when
3929re-using picklers.
Larry Hastings61272b72014-01-07 12:41:53 -08003930[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08003931
Larry Hastings3cceb382014-01-04 11:09:09 -08003932static PyObject *
3933_pickle_Pickler_clear_memo_impl(PicklerObject *self)
Larry Hastings581ee362014-01-28 05:00:08 -08003934/*[clinic end generated code: output=8665c8658aaa094b input=01bdad52f3d93e56]*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003935{
3936 if (self->memo)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003937 PyMemoTable_Clear(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003938
3939 Py_RETURN_NONE;
3940}
3941
Larry Hastings61272b72014-01-07 12:41:53 -08003942/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08003943
3944_pickle.Pickler.dump
3945
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08003946 obj: object
3947 /
3948
3949Write a pickled representation of the given object to the open file.
Larry Hastings61272b72014-01-07 12:41:53 -08003950[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08003951
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003952static PyObject *
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08003953_pickle_Pickler_dump(PicklerObject *self, PyObject *obj)
Larry Hastings581ee362014-01-28 05:00:08 -08003954/*[clinic end generated code: output=87ecad1261e02ac7 input=552eb1c0f52260d9]*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003955{
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003956 /* Check whether the Pickler was initialized correctly (issue3664).
3957 Developers often forget to call __init__() in their subclasses, which
3958 would trigger a segfault without this check. */
3959 if (self->write == NULL) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08003960 PickleState *st = _Pickle_GetGlobalState();
3961 PyErr_Format(st->PicklingError,
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003962 "Pickler.__init__() was not called by %s.__init__()",
3963 Py_TYPE(self)->tp_name);
3964 return NULL;
3965 }
3966
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003967 if (_Pickler_ClearBuffer(self) < 0)
3968 return NULL;
3969
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003970 if (dump(self, obj) < 0)
3971 return NULL;
3972
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003973 if (_Pickler_FlushToFile(self) < 0)
3974 return NULL;
3975
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003976 Py_RETURN_NONE;
3977}
3978
Serhiy Storchaka5bbd2312014-12-16 19:39:08 +02003979/*[clinic input]
3980
3981_pickle.Pickler.__sizeof__ -> Py_ssize_t
3982
3983Returns size in memory, in bytes.
3984[clinic start generated code]*/
3985
3986static Py_ssize_t
3987_pickle_Pickler___sizeof___impl(PicklerObject *self)
3988/*[clinic end generated code: output=106edb3123f332e1 input=8cbbec9bd5540d42]*/
3989{
3990 Py_ssize_t res, s;
3991
3992 res = sizeof(PicklerObject);
3993 if (self->memo != NULL) {
3994 res += sizeof(PyMemoTable);
3995 res += self->memo->mt_allocated * sizeof(PyMemoEntry);
3996 }
3997 if (self->output_buffer != NULL) {
3998 s = _PySys_GetSizeOf(self->output_buffer);
3999 if (s == -1)
4000 return -1;
4001 res += s;
4002 }
4003 return res;
4004}
4005
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004006static struct PyMethodDef Pickler_methods[] = {
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004007 _PICKLE_PICKLER_DUMP_METHODDEF
4008 _PICKLE_PICKLER_CLEAR_MEMO_METHODDEF
Serhiy Storchaka5bbd2312014-12-16 19:39:08 +02004009 _PICKLE_PICKLER___SIZEOF___METHODDEF
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004010 {NULL, NULL} /* sentinel */
4011};
4012
4013static void
4014Pickler_dealloc(PicklerObject *self)
4015{
4016 PyObject_GC_UnTrack(self);
4017
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004018 Py_XDECREF(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004019 Py_XDECREF(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004020 Py_XDECREF(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01004021 Py_XDECREF(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004022 Py_XDECREF(self->fast_memo);
4023
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004024 PyMemoTable_Del(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004025
4026 Py_TYPE(self)->tp_free((PyObject *)self);
4027}
4028
4029static int
4030Pickler_traverse(PicklerObject *self, visitproc visit, void *arg)
4031{
4032 Py_VISIT(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004033 Py_VISIT(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01004034 Py_VISIT(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004035 Py_VISIT(self->fast_memo);
4036 return 0;
4037}
4038
4039static int
4040Pickler_clear(PicklerObject *self)
4041{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004042 Py_CLEAR(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004043 Py_CLEAR(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004044 Py_CLEAR(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01004045 Py_CLEAR(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004046 Py_CLEAR(self->fast_memo);
4047
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004048 if (self->memo != NULL) {
4049 PyMemoTable *memo = self->memo;
4050 self->memo = NULL;
4051 PyMemoTable_Del(memo);
4052 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004053 return 0;
4054}
4055
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004056
Larry Hastings61272b72014-01-07 12:41:53 -08004057/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004058
4059_pickle.Pickler.__init__
4060
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004061 file: object
4062 protocol: object = NULL
4063 fix_imports: bool = True
4064
4065This takes a binary file for writing a pickle data stream.
4066
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08004067The optional *protocol* argument tells the pickler to use the given
4068protocol; supported protocols are 0, 1, 2, 3 and 4. The default
4069protocol is 3; a backward-incompatible protocol designed for Python 3.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004070
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08004071Specifying a negative protocol version selects the highest protocol
4072version supported. The higher the protocol used, the more recent the
4073version of Python needed to read the pickle produced.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004074
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08004075The *file* argument must have a write() method that accepts a single
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004076bytes argument. It can thus be a file object opened for binary
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08004077writing, a io.BytesIO instance, or any other custom object that meets
4078this interface.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004079
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08004080If *fix_imports* is True and protocol is less than 3, pickle will try
4081to map the new Python 3 names to the old module names used in Python
40822, so that the pickle data stream is readable with Python 2.
Larry Hastings61272b72014-01-07 12:41:53 -08004083[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004084
Larry Hastingsb7ccb202014-01-18 23:50:21 -08004085static int
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004086_pickle_Pickler___init___impl(PicklerObject *self, PyObject *file, PyObject *protocol, int fix_imports)
Larry Hastings581ee362014-01-28 05:00:08 -08004087/*[clinic end generated code: output=56e229f3b1f4332f input=b8cdeb7e3f5ee674]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004088{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02004089 _Py_IDENTIFIER(persistent_id);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01004090 _Py_IDENTIFIER(dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004091
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004092 /* In case of multiple __init__() calls, clear previous content. */
4093 if (self->write != NULL)
4094 (void)Pickler_clear(self);
4095
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004096 if (_Pickler_SetProtocol(self, protocol, fix_imports) < 0)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08004097 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004098
4099 if (_Pickler_SetOutputStream(self, file) < 0)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08004100 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004101
4102 /* memo and output_buffer may have already been created in _Pickler_New */
4103 if (self->memo == NULL) {
4104 self->memo = PyMemoTable_New();
4105 if (self->memo == NULL)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08004106 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004107 }
4108 self->output_len = 0;
4109 if (self->output_buffer == NULL) {
4110 self->max_output_len = WRITE_BUF_SIZE;
4111 self->output_buffer = PyBytes_FromStringAndSize(NULL,
4112 self->max_output_len);
4113 if (self->output_buffer == NULL)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08004114 return -1;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00004115 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004116
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00004117 self->fast = 0;
4118 self->fast_nesting = 0;
4119 self->fast_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004120 self->pers_func = NULL;
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02004121 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_id)) {
4122 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
4123 &PyId_persistent_id);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004124 if (self->pers_func == NULL)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08004125 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004126 }
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01004127 self->dispatch_table = NULL;
4128 if (_PyObject_HasAttrId((PyObject *)self, &PyId_dispatch_table)) {
4129 self->dispatch_table = _PyObject_GetAttrId((PyObject *)self,
4130 &PyId_dispatch_table);
4131 if (self->dispatch_table == NULL)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08004132 return -1;
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004133 }
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08004134
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004135 return 0;
4136}
4137
Larry Hastingsb7ccb202014-01-18 23:50:21 -08004138
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004139/* Define a proxy object for the Pickler's internal memo object. This is to
4140 * avoid breaking code like:
4141 * pickler.memo.clear()
4142 * and
4143 * pickler.memo = saved_memo
4144 * Is this a good idea? Not really, but we don't want to break code that uses
4145 * it. Note that we don't implement the entire mapping API here. This is
4146 * intentional, as these should be treated as black-box implementation details.
4147 */
4148
Larry Hastings61272b72014-01-07 12:41:53 -08004149/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004150_pickle.PicklerMemoProxy.clear
4151
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004152Remove all items from memo.
Larry Hastings61272b72014-01-07 12:41:53 -08004153[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004154
Larry Hastings3cceb382014-01-04 11:09:09 -08004155static PyObject *
4156_pickle_PicklerMemoProxy_clear_impl(PicklerMemoProxyObject *self)
Larry Hastings581ee362014-01-28 05:00:08 -08004157/*[clinic end generated code: output=5fb9370d48ae8b05 input=ccc186dacd0f1405]*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004158{
4159 if (self->pickler->memo)
4160 PyMemoTable_Clear(self->pickler->memo);
4161 Py_RETURN_NONE;
4162}
4163
Larry Hastings61272b72014-01-07 12:41:53 -08004164/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004165_pickle.PicklerMemoProxy.copy
4166
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004167Copy the memo to a new object.
Larry Hastings61272b72014-01-07 12:41:53 -08004168[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004169
Larry Hastings3cceb382014-01-04 11:09:09 -08004170static PyObject *
4171_pickle_PicklerMemoProxy_copy_impl(PicklerMemoProxyObject *self)
Larry Hastings581ee362014-01-28 05:00:08 -08004172/*[clinic end generated code: output=bb83a919d29225ef input=b73043485ac30b36]*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004173{
4174 Py_ssize_t i;
4175 PyMemoTable *memo;
4176 PyObject *new_memo = PyDict_New();
4177 if (new_memo == NULL)
4178 return NULL;
4179
4180 memo = self->pickler->memo;
4181 for (i = 0; i < memo->mt_allocated; ++i) {
4182 PyMemoEntry entry = memo->mt_table[i];
4183 if (entry.me_key != NULL) {
4184 int status;
4185 PyObject *key, *value;
4186
4187 key = PyLong_FromVoidPtr(entry.me_key);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004188 value = Py_BuildValue("nO", entry.me_value, entry.me_key);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004189
4190 if (key == NULL || value == NULL) {
4191 Py_XDECREF(key);
4192 Py_XDECREF(value);
4193 goto error;
4194 }
4195 status = PyDict_SetItem(new_memo, key, value);
4196 Py_DECREF(key);
4197 Py_DECREF(value);
4198 if (status < 0)
4199 goto error;
4200 }
4201 }
4202 return new_memo;
4203
4204 error:
4205 Py_XDECREF(new_memo);
4206 return NULL;
4207}
4208
Larry Hastings61272b72014-01-07 12:41:53 -08004209/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004210_pickle.PicklerMemoProxy.__reduce__
4211
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004212Implement pickle support.
Larry Hastings61272b72014-01-07 12:41:53 -08004213[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004214
Larry Hastings3cceb382014-01-04 11:09:09 -08004215static PyObject *
4216_pickle_PicklerMemoProxy___reduce___impl(PicklerMemoProxyObject *self)
Larry Hastings581ee362014-01-28 05:00:08 -08004217/*[clinic end generated code: output=bebba1168863ab1d input=2f7c540e24b7aae4]*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004218{
4219 PyObject *reduce_value, *dict_args;
Larry Hastings3cceb382014-01-04 11:09:09 -08004220 PyObject *contents = _pickle_PicklerMemoProxy_copy_impl(self);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004221 if (contents == NULL)
4222 return NULL;
4223
4224 reduce_value = PyTuple_New(2);
4225 if (reduce_value == NULL) {
4226 Py_DECREF(contents);
4227 return NULL;
4228 }
4229 dict_args = PyTuple_New(1);
4230 if (dict_args == NULL) {
4231 Py_DECREF(contents);
4232 Py_DECREF(reduce_value);
4233 return NULL;
4234 }
4235 PyTuple_SET_ITEM(dict_args, 0, contents);
4236 Py_INCREF((PyObject *)&PyDict_Type);
4237 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
4238 PyTuple_SET_ITEM(reduce_value, 1, dict_args);
4239 return reduce_value;
4240}
4241
4242static PyMethodDef picklerproxy_methods[] = {
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004243 _PICKLE_PICKLERMEMOPROXY_CLEAR_METHODDEF
4244 _PICKLE_PICKLERMEMOPROXY_COPY_METHODDEF
4245 _PICKLE_PICKLERMEMOPROXY___REDUCE___METHODDEF
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004246 {NULL, NULL} /* sentinel */
4247};
4248
4249static void
4250PicklerMemoProxy_dealloc(PicklerMemoProxyObject *self)
4251{
4252 PyObject_GC_UnTrack(self);
4253 Py_XDECREF(self->pickler);
4254 PyObject_GC_Del((PyObject *)self);
4255}
4256
4257static int
4258PicklerMemoProxy_traverse(PicklerMemoProxyObject *self,
4259 visitproc visit, void *arg)
4260{
4261 Py_VISIT(self->pickler);
4262 return 0;
4263}
4264
4265static int
4266PicklerMemoProxy_clear(PicklerMemoProxyObject *self)
4267{
4268 Py_CLEAR(self->pickler);
4269 return 0;
4270}
4271
4272static PyTypeObject PicklerMemoProxyType = {
4273 PyVarObject_HEAD_INIT(NULL, 0)
4274 "_pickle.PicklerMemoProxy", /*tp_name*/
4275 sizeof(PicklerMemoProxyObject), /*tp_basicsize*/
4276 0,
4277 (destructor)PicklerMemoProxy_dealloc, /* tp_dealloc */
4278 0, /* tp_print */
4279 0, /* tp_getattr */
4280 0, /* tp_setattr */
4281 0, /* tp_compare */
4282 0, /* tp_repr */
4283 0, /* tp_as_number */
4284 0, /* tp_as_sequence */
4285 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00004286 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004287 0, /* tp_call */
4288 0, /* tp_str */
4289 PyObject_GenericGetAttr, /* tp_getattro */
4290 PyObject_GenericSetAttr, /* tp_setattro */
4291 0, /* tp_as_buffer */
4292 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
4293 0, /* tp_doc */
4294 (traverseproc)PicklerMemoProxy_traverse, /* tp_traverse */
4295 (inquiry)PicklerMemoProxy_clear, /* tp_clear */
4296 0, /* tp_richcompare */
4297 0, /* tp_weaklistoffset */
4298 0, /* tp_iter */
4299 0, /* tp_iternext */
4300 picklerproxy_methods, /* tp_methods */
4301};
4302
4303static PyObject *
4304PicklerMemoProxy_New(PicklerObject *pickler)
4305{
4306 PicklerMemoProxyObject *self;
4307
4308 self = PyObject_GC_New(PicklerMemoProxyObject, &PicklerMemoProxyType);
4309 if (self == NULL)
4310 return NULL;
4311 Py_INCREF(pickler);
4312 self->pickler = pickler;
4313 PyObject_GC_Track(self);
4314 return (PyObject *)self;
4315}
4316
4317/*****************************************************************************/
4318
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004319static PyObject *
4320Pickler_get_memo(PicklerObject *self)
4321{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004322 return PicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004323}
4324
4325static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004326Pickler_set_memo(PicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004327{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004328 PyMemoTable *new_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004329
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004330 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004331 PyErr_SetString(PyExc_TypeError,
4332 "attribute deletion is not supported");
4333 return -1;
4334 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004335
4336 if (Py_TYPE(obj) == &PicklerMemoProxyType) {
4337 PicklerObject *pickler =
4338 ((PicklerMemoProxyObject *)obj)->pickler;
4339
4340 new_memo = PyMemoTable_Copy(pickler->memo);
4341 if (new_memo == NULL)
4342 return -1;
4343 }
4344 else if (PyDict_Check(obj)) {
4345 Py_ssize_t i = 0;
4346 PyObject *key, *value;
4347
4348 new_memo = PyMemoTable_New();
4349 if (new_memo == NULL)
4350 return -1;
4351
4352 while (PyDict_Next(obj, &i, &key, &value)) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004353 Py_ssize_t memo_id;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004354 PyObject *memo_obj;
4355
4356 if (!PyTuple_Check(value) || Py_SIZE(value) != 2) {
4357 PyErr_SetString(PyExc_TypeError,
4358 "'memo' values must be 2-item tuples");
4359 goto error;
4360 }
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004361 memo_id = PyLong_AsSsize_t(PyTuple_GET_ITEM(value, 0));
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004362 if (memo_id == -1 && PyErr_Occurred())
4363 goto error;
4364 memo_obj = PyTuple_GET_ITEM(value, 1);
4365 if (PyMemoTable_Set(new_memo, memo_obj, memo_id) < 0)
4366 goto error;
4367 }
4368 }
4369 else {
4370 PyErr_Format(PyExc_TypeError,
4371 "'memo' attribute must be an PicklerMemoProxy object"
4372 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004373 return -1;
4374 }
4375
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004376 PyMemoTable_Del(self->memo);
4377 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004378
4379 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004380
4381 error:
4382 if (new_memo)
4383 PyMemoTable_Del(new_memo);
4384 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004385}
4386
4387static PyObject *
4388Pickler_get_persid(PicklerObject *self)
4389{
4390 if (self->pers_func == NULL)
4391 PyErr_SetString(PyExc_AttributeError, "persistent_id");
4392 else
4393 Py_INCREF(self->pers_func);
4394 return self->pers_func;
4395}
4396
4397static int
4398Pickler_set_persid(PicklerObject *self, PyObject *value)
4399{
4400 PyObject *tmp;
4401
4402 if (value == NULL) {
4403 PyErr_SetString(PyExc_TypeError,
4404 "attribute deletion is not supported");
4405 return -1;
4406 }
4407 if (!PyCallable_Check(value)) {
4408 PyErr_SetString(PyExc_TypeError,
4409 "persistent_id must be a callable taking one argument");
4410 return -1;
4411 }
4412
4413 tmp = self->pers_func;
4414 Py_INCREF(value);
4415 self->pers_func = value;
4416 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
4417
4418 return 0;
4419}
4420
4421static PyMemberDef Pickler_members[] = {
4422 {"bin", T_INT, offsetof(PicklerObject, bin)},
4423 {"fast", T_INT, offsetof(PicklerObject, fast)},
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01004424 {"dispatch_table", T_OBJECT_EX, offsetof(PicklerObject, dispatch_table)},
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004425 {NULL}
4426};
4427
4428static PyGetSetDef Pickler_getsets[] = {
4429 {"memo", (getter)Pickler_get_memo,
4430 (setter)Pickler_set_memo},
4431 {"persistent_id", (getter)Pickler_get_persid,
4432 (setter)Pickler_set_persid},
4433 {NULL}
4434};
4435
4436static PyTypeObject Pickler_Type = {
4437 PyVarObject_HEAD_INIT(NULL, 0)
4438 "_pickle.Pickler" , /*tp_name*/
4439 sizeof(PicklerObject), /*tp_basicsize*/
4440 0, /*tp_itemsize*/
4441 (destructor)Pickler_dealloc, /*tp_dealloc*/
4442 0, /*tp_print*/
4443 0, /*tp_getattr*/
4444 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00004445 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004446 0, /*tp_repr*/
4447 0, /*tp_as_number*/
4448 0, /*tp_as_sequence*/
4449 0, /*tp_as_mapping*/
4450 0, /*tp_hash*/
4451 0, /*tp_call*/
4452 0, /*tp_str*/
4453 0, /*tp_getattro*/
4454 0, /*tp_setattro*/
4455 0, /*tp_as_buffer*/
4456 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08004457 _pickle_Pickler___init____doc__, /*tp_doc*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004458 (traverseproc)Pickler_traverse, /*tp_traverse*/
4459 (inquiry)Pickler_clear, /*tp_clear*/
4460 0, /*tp_richcompare*/
4461 0, /*tp_weaklistoffset*/
4462 0, /*tp_iter*/
4463 0, /*tp_iternext*/
4464 Pickler_methods, /*tp_methods*/
4465 Pickler_members, /*tp_members*/
4466 Pickler_getsets, /*tp_getset*/
4467 0, /*tp_base*/
4468 0, /*tp_dict*/
4469 0, /*tp_descr_get*/
4470 0, /*tp_descr_set*/
4471 0, /*tp_dictoffset*/
Larry Hastingsb7ccb202014-01-18 23:50:21 -08004472 _pickle_Pickler___init__, /*tp_init*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004473 PyType_GenericAlloc, /*tp_alloc*/
4474 PyType_GenericNew, /*tp_new*/
4475 PyObject_GC_Del, /*tp_free*/
4476 0, /*tp_is_gc*/
4477};
4478
Victor Stinner121aab42011-09-29 23:40:53 +02004479/* Temporary helper for calling self.find_class().
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004480
4481 XXX: It would be nice to able to avoid Python function call overhead, by
4482 using directly the C version of find_class(), when find_class() is not
4483 overridden by a subclass. Although, this could become rather hackish. A
4484 simpler optimization would be to call the C function when self is not a
4485 subclass instance. */
4486static PyObject *
4487find_class(UnpicklerObject *self, PyObject *module_name, PyObject *global_name)
4488{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02004489 _Py_IDENTIFIER(find_class);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02004490
4491 return _PyObject_CallMethodId((PyObject *)self, &PyId_find_class, "OO",
4492 module_name, global_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004493}
4494
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004495static Py_ssize_t
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004496marker(UnpicklerObject *self)
4497{
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08004498 PickleState *st = _Pickle_GetGlobalState();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004499 if (self->num_marks < 1) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08004500 PyErr_SetString(st->UnpicklingError, "could not find MARK");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004501 return -1;
4502 }
4503
4504 return self->marks[--self->num_marks];
4505}
4506
4507static int
4508load_none(UnpicklerObject *self)
4509{
4510 PDATA_APPEND(self->stack, Py_None, -1);
4511 return 0;
4512}
4513
4514static int
4515bad_readline(void)
4516{
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08004517 PickleState *st = _Pickle_GetGlobalState();
4518 PyErr_SetString(st->UnpicklingError, "pickle data was truncated");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004519 return -1;
4520}
4521
4522static int
4523load_int(UnpicklerObject *self)
4524{
4525 PyObject *value;
4526 char *endptr, *s;
4527 Py_ssize_t len;
4528 long x;
4529
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004530 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004531 return -1;
4532 if (len < 2)
4533 return bad_readline();
4534
4535 errno = 0;
Victor Stinner121aab42011-09-29 23:40:53 +02004536 /* XXX: Should the base argument of strtol() be explicitly set to 10?
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004537 XXX(avassalotti): Should this uses PyOS_strtol()? */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004538 x = strtol(s, &endptr, 0);
4539
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004540 if (errno || (*endptr != '\n' && *endptr != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004541 /* Hm, maybe we've got something long. Let's try reading
Serhiy Storchaka95949422013-08-27 19:40:23 +03004542 * it as a Python int object. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004543 errno = 0;
4544 /* XXX: Same thing about the base here. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004545 value = PyLong_FromString(s, NULL, 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004546 if (value == NULL) {
4547 PyErr_SetString(PyExc_ValueError,
4548 "could not convert string to int");
4549 return -1;
4550 }
4551 }
4552 else {
4553 if (len == 3 && (x == 0 || x == 1)) {
4554 if ((value = PyBool_FromLong(x)) == NULL)
4555 return -1;
4556 }
4557 else {
4558 if ((value = PyLong_FromLong(x)) == NULL)
4559 return -1;
4560 }
4561 }
4562
4563 PDATA_PUSH(self->stack, value, -1);
4564 return 0;
4565}
4566
4567static int
4568load_bool(UnpicklerObject *self, PyObject *boolean)
4569{
4570 assert(boolean == Py_True || boolean == Py_False);
4571 PDATA_APPEND(self->stack, boolean, -1);
4572 return 0;
4573}
4574
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004575/* s contains x bytes of an unsigned little-endian integer. Return its value
4576 * as a C Py_ssize_t, or -1 if it's higher than PY_SSIZE_T_MAX.
4577 */
4578static Py_ssize_t
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004579calc_binsize(char *bytes, int nbytes)
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004580{
4581 unsigned char *s = (unsigned char *)bytes;
Victor Stinnerf13c46c2014-08-17 21:05:55 +02004582 Py_ssize_t i;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004583 size_t x = 0;
4584
Victor Stinnerf13c46c2014-08-17 21:05:55 +02004585 for (i = 0; i < nbytes && (size_t)i < sizeof(size_t); i++) {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004586 x |= (size_t) s[i] << (8 * i);
4587 }
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004588
4589 if (x > PY_SSIZE_T_MAX)
4590 return -1;
4591 else
4592 return (Py_ssize_t) x;
4593}
4594
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004595/* s contains x bytes of a little-endian integer. Return its value as a
4596 * C int. Obscure: when x is 1 or 2, this is an unsigned little-endian
4597 * int, but when x is 4 it's a signed one. This is an historical source
4598 * of x-platform bugs.
4599 */
4600static long
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004601calc_binint(char *bytes, int nbytes)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004602{
4603 unsigned char *s = (unsigned char *)bytes;
Victor Stinnerf13c46c2014-08-17 21:05:55 +02004604 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004605 long x = 0;
4606
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004607 for (i = 0; i < nbytes; i++) {
4608 x |= (long)s[i] << (8 * i);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004609 }
4610
4611 /* Unlike BININT1 and BININT2, BININT (more accurately BININT4)
4612 * is signed, so on a box with longs bigger than 4 bytes we need
4613 * to extend a BININT's sign bit to the full width.
4614 */
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004615 if (SIZEOF_LONG > 4 && nbytes == 4) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004616 x |= -(x & (1L << 31));
4617 }
4618
4619 return x;
4620}
4621
4622static int
4623load_binintx(UnpicklerObject *self, char *s, int size)
4624{
4625 PyObject *value;
4626 long x;
4627
4628 x = calc_binint(s, size);
4629
4630 if ((value = PyLong_FromLong(x)) == NULL)
4631 return -1;
4632
4633 PDATA_PUSH(self->stack, value, -1);
4634 return 0;
4635}
4636
4637static int
4638load_binint(UnpicklerObject *self)
4639{
4640 char *s;
4641
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004642 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004643 return -1;
4644
4645 return load_binintx(self, s, 4);
4646}
4647
4648static int
4649load_binint1(UnpicklerObject *self)
4650{
4651 char *s;
4652
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004653 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004654 return -1;
4655
4656 return load_binintx(self, s, 1);
4657}
4658
4659static int
4660load_binint2(UnpicklerObject *self)
4661{
4662 char *s;
4663
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004664 if (_Unpickler_Read(self, &s, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004665 return -1;
4666
4667 return load_binintx(self, s, 2);
4668}
4669
4670static int
4671load_long(UnpicklerObject *self)
4672{
4673 PyObject *value;
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004674 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004675 Py_ssize_t len;
4676
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004677 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004678 return -1;
4679 if (len < 2)
4680 return bad_readline();
4681
Mark Dickinson8dd05142009-01-20 20:43:58 +00004682 /* s[len-2] will usually be 'L' (and s[len-1] is '\n'); we need to remove
4683 the 'L' before calling PyLong_FromString. In order to maintain
4684 compatibility with Python 3.0.0, we don't actually *require*
4685 the 'L' to be present. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004686 if (s[len-2] == 'L')
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004687 s[len-2] = '\0';
Alexandre Vassalottie4bccb72009-01-24 01:47:57 +00004688 /* XXX: Should the base argument explicitly set to 10? */
4689 value = PyLong_FromString(s, NULL, 0);
Mark Dickinson8dd05142009-01-20 20:43:58 +00004690 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004691 return -1;
4692
4693 PDATA_PUSH(self->stack, value, -1);
4694 return 0;
4695}
4696
4697/* 'size' bytes contain the # of bytes of little-endian 256's-complement
4698 * data following.
4699 */
4700static int
4701load_counted_long(UnpicklerObject *self, int size)
4702{
4703 PyObject *value;
4704 char *nbytes;
4705 char *pdata;
4706
4707 assert(size == 1 || size == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004708 if (_Unpickler_Read(self, &nbytes, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004709 return -1;
4710
4711 size = calc_binint(nbytes, size);
4712 if (size < 0) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08004713 PickleState *st = _Pickle_GetGlobalState();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004714 /* Corrupt or hostile pickle -- we never write one like this */
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08004715 PyErr_SetString(st->UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004716 "LONG pickle has negative byte count");
4717 return -1;
4718 }
4719
4720 if (size == 0)
4721 value = PyLong_FromLong(0L);
4722 else {
4723 /* Read the raw little-endian bytes and convert. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004724 if (_Unpickler_Read(self, &pdata, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004725 return -1;
4726 value = _PyLong_FromByteArray((unsigned char *)pdata, (size_t)size,
4727 1 /* little endian */ , 1 /* signed */ );
4728 }
4729 if (value == NULL)
4730 return -1;
4731 PDATA_PUSH(self->stack, value, -1);
4732 return 0;
4733}
4734
4735static int
4736load_float(UnpicklerObject *self)
4737{
4738 PyObject *value;
4739 char *endptr, *s;
4740 Py_ssize_t len;
4741 double d;
4742
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004743 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004744 return -1;
4745 if (len < 2)
4746 return bad_readline();
4747
4748 errno = 0;
Mark Dickinson725bfd82009-05-03 20:33:40 +00004749 d = PyOS_string_to_double(s, &endptr, PyExc_OverflowError);
4750 if (d == -1.0 && PyErr_Occurred())
4751 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004752 if ((endptr[0] != '\n') && (endptr[0] != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004753 PyErr_SetString(PyExc_ValueError, "could not convert string to float");
4754 return -1;
4755 }
Mark Dickinson725bfd82009-05-03 20:33:40 +00004756 value = PyFloat_FromDouble(d);
4757 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004758 return -1;
4759
4760 PDATA_PUSH(self->stack, value, -1);
4761 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004762}
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004763
4764static int
4765load_binfloat(UnpicklerObject *self)
4766{
4767 PyObject *value;
4768 double x;
4769 char *s;
4770
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004771 if (_Unpickler_Read(self, &s, 8) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004772 return -1;
4773
4774 x = _PyFloat_Unpack8((unsigned char *)s, 0);
4775 if (x == -1.0 && PyErr_Occurred())
4776 return -1;
4777
4778 if ((value = PyFloat_FromDouble(x)) == NULL)
4779 return -1;
4780
4781 PDATA_PUSH(self->stack, value, -1);
4782 return 0;
4783}
4784
4785static int
4786load_string(UnpicklerObject *self)
4787{
4788 PyObject *bytes;
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08004789 PyObject *obj;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004790 Py_ssize_t len;
4791 char *s, *p;
4792
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004793 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004794 return -1;
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004795 /* Strip the newline */
4796 len--;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004797 /* Strip outermost quotes */
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004798 if (len >= 2 && s[0] == s[len - 1] && (s[0] == '\'' || s[0] == '"')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004799 p = s + 1;
4800 len -= 2;
4801 }
4802 else {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08004803 PickleState *st = _Pickle_GetGlobalState();
4804 PyErr_SetString(st->UnpicklingError,
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004805 "the STRING opcode argument must be quoted");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004806 return -1;
4807 }
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004808 assert(len >= 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004809
4810 /* Use the PyBytes API to decode the string, since that is what is used
4811 to encode, and then coerce the result to Unicode. */
4812 bytes = PyBytes_DecodeEscape(p, len, NULL, 0, NULL);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004813 if (bytes == NULL)
4814 return -1;
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08004815
4816 /* Leave the Python 2.x strings as bytes if the *encoding* given to the
4817 Unpickler was 'bytes'. Otherwise, convert them to unicode. */
4818 if (strcmp(self->encoding, "bytes") == 0) {
4819 obj = bytes;
4820 }
4821 else {
4822 obj = PyUnicode_FromEncodedObject(bytes, self->encoding, self->errors);
4823 Py_DECREF(bytes);
4824 if (obj == NULL) {
4825 return -1;
4826 }
4827 }
4828
4829 PDATA_PUSH(self->stack, obj, -1);
4830 return 0;
4831}
4832
4833static int
4834load_counted_binstring(UnpicklerObject *self, int nbytes)
4835{
4836 PyObject *obj;
4837 Py_ssize_t size;
4838 char *s;
4839
4840 if (_Unpickler_Read(self, &s, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004841 return -1;
4842
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08004843 size = calc_binsize(s, nbytes);
4844 if (size < 0) {
4845 PickleState *st = _Pickle_GetGlobalState();
4846 PyErr_Format(st->UnpicklingError,
4847 "BINSTRING exceeds system's maximum size of %zd bytes",
4848 PY_SSIZE_T_MAX);
4849 return -1;
4850 }
4851
4852 if (_Unpickler_Read(self, &s, size) < 0)
4853 return -1;
4854
4855 /* Convert Python 2.x strings to bytes if the *encoding* given to the
4856 Unpickler was 'bytes'. Otherwise, convert them to unicode. */
4857 if (strcmp(self->encoding, "bytes") == 0) {
4858 obj = PyBytes_FromStringAndSize(s, size);
4859 }
4860 else {
4861 obj = PyUnicode_Decode(s, size, self->encoding, self->errors);
4862 }
4863 if (obj == NULL) {
4864 return -1;
4865 }
4866
4867 PDATA_PUSH(self->stack, obj, -1);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004868 return 0;
4869}
4870
4871static int
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004872load_counted_binbytes(UnpicklerObject *self, int nbytes)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004873{
4874 PyObject *bytes;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004875 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004876 char *s;
4877
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004878 if (_Unpickler_Read(self, &s, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004879 return -1;
4880
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004881 size = calc_binsize(s, nbytes);
4882 if (size < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004883 PyErr_Format(PyExc_OverflowError,
4884 "BINBYTES exceeds system's maximum size of %zd bytes",
Alexandre Vassalotticc757172013-04-14 02:25:10 -07004885 PY_SSIZE_T_MAX);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004886 return -1;
4887 }
4888
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004889 if (_Unpickler_Read(self, &s, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004890 return -1;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004891
4892 bytes = PyBytes_FromStringAndSize(s, size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004893 if (bytes == NULL)
4894 return -1;
4895
4896 PDATA_PUSH(self->stack, bytes, -1);
4897 return 0;
4898}
4899
4900static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004901load_unicode(UnpicklerObject *self)
4902{
4903 PyObject *str;
4904 Py_ssize_t len;
4905 char *s;
4906
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004907 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004908 return -1;
4909 if (len < 1)
4910 return bad_readline();
4911
4912 str = PyUnicode_DecodeRawUnicodeEscape(s, len - 1, NULL);
4913 if (str == NULL)
4914 return -1;
4915
4916 PDATA_PUSH(self->stack, str, -1);
4917 return 0;
4918}
4919
4920static int
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004921load_counted_binunicode(UnpicklerObject *self, int nbytes)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004922{
4923 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004924 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004925 char *s;
4926
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004927 if (_Unpickler_Read(self, &s, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004928 return -1;
4929
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01004930 size = calc_binsize(s, nbytes);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004931 if (size < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004932 PyErr_Format(PyExc_OverflowError,
4933 "BINUNICODE exceeds system's maximum size of %zd bytes",
Alexandre Vassalotticc757172013-04-14 02:25:10 -07004934 PY_SSIZE_T_MAX);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004935 return -1;
4936 }
4937
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004938 if (_Unpickler_Read(self, &s, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004939 return -1;
4940
Victor Stinner485fb562010-04-13 11:07:24 +00004941 str = PyUnicode_DecodeUTF8(s, size, "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004942 if (str == NULL)
4943 return -1;
4944
4945 PDATA_PUSH(self->stack, str, -1);
4946 return 0;
4947}
4948
4949static int
4950load_tuple(UnpicklerObject *self)
4951{
4952 PyObject *tuple;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004953 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004954
4955 if ((i = marker(self)) < 0)
4956 return -1;
4957
4958 tuple = Pdata_poptuple(self->stack, i);
4959 if (tuple == NULL)
4960 return -1;
4961 PDATA_PUSH(self->stack, tuple, -1);
4962 return 0;
4963}
4964
4965static int
4966load_counted_tuple(UnpicklerObject *self, int len)
4967{
4968 PyObject *tuple;
4969
4970 tuple = PyTuple_New(len);
4971 if (tuple == NULL)
4972 return -1;
4973
4974 while (--len >= 0) {
4975 PyObject *item;
4976
4977 PDATA_POP(self->stack, item);
4978 if (item == NULL)
4979 return -1;
4980 PyTuple_SET_ITEM(tuple, len, item);
4981 }
4982 PDATA_PUSH(self->stack, tuple, -1);
4983 return 0;
4984}
4985
4986static int
4987load_empty_list(UnpicklerObject *self)
4988{
4989 PyObject *list;
4990
4991 if ((list = PyList_New(0)) == NULL)
4992 return -1;
4993 PDATA_PUSH(self->stack, list, -1);
4994 return 0;
4995}
4996
4997static int
4998load_empty_dict(UnpicklerObject *self)
4999{
5000 PyObject *dict;
5001
5002 if ((dict = PyDict_New()) == NULL)
5003 return -1;
5004 PDATA_PUSH(self->stack, dict, -1);
5005 return 0;
5006}
5007
5008static int
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005009load_empty_set(UnpicklerObject *self)
5010{
5011 PyObject *set;
5012
5013 if ((set = PySet_New(NULL)) == NULL)
5014 return -1;
5015 PDATA_PUSH(self->stack, set, -1);
5016 return 0;
5017}
5018
5019static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005020load_list(UnpicklerObject *self)
5021{
5022 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005023 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005024
5025 if ((i = marker(self)) < 0)
5026 return -1;
5027
5028 list = Pdata_poplist(self->stack, i);
5029 if (list == NULL)
5030 return -1;
5031 PDATA_PUSH(self->stack, list, -1);
5032 return 0;
5033}
5034
5035static int
5036load_dict(UnpicklerObject *self)
5037{
5038 PyObject *dict, *key, *value;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005039 Py_ssize_t i, j, k;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005040
5041 if ((i = marker(self)) < 0)
5042 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005043 j = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005044
5045 if ((dict = PyDict_New()) == NULL)
5046 return -1;
5047
5048 for (k = i + 1; k < j; k += 2) {
5049 key = self->stack->data[k - 1];
5050 value = self->stack->data[k];
5051 if (PyDict_SetItem(dict, key, value) < 0) {
5052 Py_DECREF(dict);
5053 return -1;
5054 }
5055 }
5056 Pdata_clear(self->stack, i);
5057 PDATA_PUSH(self->stack, dict, -1);
5058 return 0;
5059}
5060
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005061static int
5062load_frozenset(UnpicklerObject *self)
5063{
5064 PyObject *items;
5065 PyObject *frozenset;
5066 Py_ssize_t i;
5067
5068 if ((i = marker(self)) < 0)
5069 return -1;
5070
5071 items = Pdata_poptuple(self->stack, i);
5072 if (items == NULL)
5073 return -1;
5074
5075 frozenset = PyFrozenSet_New(items);
5076 Py_DECREF(items);
5077 if (frozenset == NULL)
5078 return -1;
5079
5080 PDATA_PUSH(self->stack, frozenset, -1);
5081 return 0;
5082}
5083
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005084static PyObject *
5085instantiate(PyObject *cls, PyObject *args)
5086{
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005087 PyObject *result = NULL;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005088 _Py_IDENTIFIER(__getinitargs__);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005089 /* Caller must assure args are a tuple. Normally, args come from
5090 Pdata_poptuple which packs objects from the top of the stack
5091 into a newly created tuple. */
5092 assert(PyTuple_Check(args));
5093 if (Py_SIZE(args) > 0 || !PyType_Check(cls) ||
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005094 _PyObject_HasAttrId(cls, &PyId___getinitargs__)) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005095 result = PyObject_CallObject(cls, args);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005096 }
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005097 else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005098 _Py_IDENTIFIER(__new__);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02005099
5100 result = _PyObject_CallMethodId(cls, &PyId___new__, "O", cls);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005101 }
5102 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005103}
5104
5105static int
5106load_obj(UnpicklerObject *self)
5107{
5108 PyObject *cls, *args, *obj = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005109 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005110
5111 if ((i = marker(self)) < 0)
5112 return -1;
5113
5114 args = Pdata_poptuple(self->stack, i + 1);
5115 if (args == NULL)
5116 return -1;
5117
5118 PDATA_POP(self->stack, cls);
5119 if (cls) {
5120 obj = instantiate(cls, args);
5121 Py_DECREF(cls);
5122 }
5123 Py_DECREF(args);
5124 if (obj == NULL)
5125 return -1;
5126
5127 PDATA_PUSH(self->stack, obj, -1);
5128 return 0;
5129}
5130
5131static int
5132load_inst(UnpicklerObject *self)
5133{
5134 PyObject *cls = NULL;
5135 PyObject *args = NULL;
5136 PyObject *obj = NULL;
5137 PyObject *module_name;
5138 PyObject *class_name;
5139 Py_ssize_t len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005140 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005141 char *s;
5142
5143 if ((i = marker(self)) < 0)
5144 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005145 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005146 return -1;
5147 if (len < 2)
5148 return bad_readline();
5149
5150 /* Here it is safe to use PyUnicode_DecodeASCII(), even though non-ASCII
5151 identifiers are permitted in Python 3.0, since the INST opcode is only
5152 supported by older protocols on Python 2.x. */
5153 module_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
5154 if (module_name == NULL)
5155 return -1;
5156
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005157 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005158 if (len < 2)
5159 return bad_readline();
5160 class_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005161 if (class_name != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005162 cls = find_class(self, module_name, class_name);
5163 Py_DECREF(class_name);
5164 }
5165 }
5166 Py_DECREF(module_name);
5167
5168 if (cls == NULL)
5169 return -1;
5170
5171 if ((args = Pdata_poptuple(self->stack, i)) != NULL) {
5172 obj = instantiate(cls, args);
5173 Py_DECREF(args);
5174 }
5175 Py_DECREF(cls);
5176
5177 if (obj == NULL)
5178 return -1;
5179
5180 PDATA_PUSH(self->stack, obj, -1);
5181 return 0;
5182}
5183
5184static int
5185load_newobj(UnpicklerObject *self)
5186{
5187 PyObject *args = NULL;
5188 PyObject *clsraw = NULL;
5189 PyTypeObject *cls; /* clsraw cast to its true type */
5190 PyObject *obj;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005191 PickleState *st = _Pickle_GetGlobalState();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005192
5193 /* Stack is ... cls argtuple, and we want to call
5194 * cls.__new__(cls, *argtuple).
5195 */
5196 PDATA_POP(self->stack, args);
5197 if (args == NULL)
5198 goto error;
5199 if (!PyTuple_Check(args)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005200 PyErr_SetString(st->UnpicklingError,
5201 "NEWOBJ expected an arg " "tuple.");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005202 goto error;
5203 }
5204
5205 PDATA_POP(self->stack, clsraw);
5206 cls = (PyTypeObject *)clsraw;
5207 if (cls == NULL)
5208 goto error;
5209 if (!PyType_Check(cls)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005210 PyErr_SetString(st->UnpicklingError, "NEWOBJ class argument "
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005211 "isn't a type object");
5212 goto error;
5213 }
5214 if (cls->tp_new == NULL) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005215 PyErr_SetString(st->UnpicklingError, "NEWOBJ class argument "
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005216 "has NULL tp_new");
5217 goto error;
5218 }
5219
5220 /* Call __new__. */
5221 obj = cls->tp_new(cls, args, NULL);
5222 if (obj == NULL)
5223 goto error;
5224
5225 Py_DECREF(args);
5226 Py_DECREF(clsraw);
5227 PDATA_PUSH(self->stack, obj, -1);
5228 return 0;
5229
5230 error:
5231 Py_XDECREF(args);
5232 Py_XDECREF(clsraw);
5233 return -1;
5234}
5235
5236static int
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005237load_newobj_ex(UnpicklerObject *self)
5238{
5239 PyObject *cls, *args, *kwargs;
5240 PyObject *obj;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005241 PickleState *st = _Pickle_GetGlobalState();
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005242
5243 PDATA_POP(self->stack, kwargs);
5244 if (kwargs == NULL) {
5245 return -1;
5246 }
5247 PDATA_POP(self->stack, args);
5248 if (args == NULL) {
5249 Py_DECREF(kwargs);
5250 return -1;
5251 }
5252 PDATA_POP(self->stack, cls);
5253 if (cls == NULL) {
5254 Py_DECREF(kwargs);
5255 Py_DECREF(args);
5256 return -1;
5257 }
Larry Hastings61272b72014-01-07 12:41:53 -08005258
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005259 if (!PyType_Check(cls)) {
5260 Py_DECREF(kwargs);
5261 Py_DECREF(args);
5262 Py_DECREF(cls);
Larry Hastings61272b72014-01-07 12:41:53 -08005263 PyErr_Format(st->UnpicklingError,
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005264 "NEWOBJ_EX class argument must be a type, not %.200s",
5265 Py_TYPE(cls)->tp_name);
5266 return -1;
5267 }
5268
5269 if (((PyTypeObject *)cls)->tp_new == NULL) {
5270 Py_DECREF(kwargs);
5271 Py_DECREF(args);
5272 Py_DECREF(cls);
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005273 PyErr_SetString(st->UnpicklingError,
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005274 "NEWOBJ_EX class argument doesn't have __new__");
5275 return -1;
5276 }
5277 obj = ((PyTypeObject *)cls)->tp_new((PyTypeObject *)cls, args, kwargs);
5278 Py_DECREF(kwargs);
5279 Py_DECREF(args);
5280 Py_DECREF(cls);
5281 if (obj == NULL) {
5282 return -1;
5283 }
5284 PDATA_PUSH(self->stack, obj, -1);
5285 return 0;
5286}
5287
5288static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005289load_global(UnpicklerObject *self)
5290{
5291 PyObject *global = NULL;
5292 PyObject *module_name;
5293 PyObject *global_name;
5294 Py_ssize_t len;
5295 char *s;
5296
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005297 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005298 return -1;
5299 if (len < 2)
5300 return bad_readline();
5301 module_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
5302 if (!module_name)
5303 return -1;
5304
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005305 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005306 if (len < 2) {
5307 Py_DECREF(module_name);
5308 return bad_readline();
5309 }
5310 global_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
5311 if (global_name) {
5312 global = find_class(self, module_name, global_name);
5313 Py_DECREF(global_name);
5314 }
5315 }
5316 Py_DECREF(module_name);
5317
5318 if (global == NULL)
5319 return -1;
5320 PDATA_PUSH(self->stack, global, -1);
5321 return 0;
5322}
5323
5324static int
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005325load_stack_global(UnpicklerObject *self)
5326{
5327 PyObject *global;
5328 PyObject *module_name;
5329 PyObject *global_name;
5330
5331 PDATA_POP(self->stack, global_name);
5332 PDATA_POP(self->stack, module_name);
5333 if (module_name == NULL || !PyUnicode_CheckExact(module_name) ||
5334 global_name == NULL || !PyUnicode_CheckExact(global_name)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005335 PickleState *st = _Pickle_GetGlobalState();
5336 PyErr_SetString(st->UnpicklingError, "STACK_GLOBAL requires str");
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005337 Py_XDECREF(global_name);
5338 Py_XDECREF(module_name);
5339 return -1;
5340 }
5341 global = find_class(self, module_name, global_name);
5342 Py_DECREF(global_name);
5343 Py_DECREF(module_name);
5344 if (global == NULL)
5345 return -1;
5346 PDATA_PUSH(self->stack, global, -1);
5347 return 0;
5348}
5349
5350static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005351load_persid(UnpicklerObject *self)
5352{
5353 PyObject *pid;
5354 Py_ssize_t len;
5355 char *s;
5356
5357 if (self->pers_func) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005358 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005359 return -1;
Alexandre Vassalotti896414f2013-11-30 13:52:35 -08005360 if (len < 1)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005361 return bad_readline();
5362
5363 pid = PyBytes_FromStringAndSize(s, len - 1);
5364 if (pid == NULL)
5365 return -1;
5366
Alexandre Vassalotti20c28c12013-11-27 02:26:54 -08005367 /* This does not leak since _Pickle_FastCall() steals the reference
5368 to pid first. */
5369 pid = _Pickle_FastCall(self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005370 if (pid == NULL)
5371 return -1;
5372
5373 PDATA_PUSH(self->stack, pid, -1);
5374 return 0;
5375 }
5376 else {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005377 PickleState *st = _Pickle_GetGlobalState();
5378 PyErr_SetString(st->UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005379 "A load persistent id instruction was encountered,\n"
5380 "but no persistent_load function was specified.");
5381 return -1;
5382 }
5383}
5384
5385static int
5386load_binpersid(UnpicklerObject *self)
5387{
5388 PyObject *pid;
5389
5390 if (self->pers_func) {
5391 PDATA_POP(self->stack, pid);
5392 if (pid == NULL)
5393 return -1;
5394
Alexandre Vassalotti20c28c12013-11-27 02:26:54 -08005395 /* This does not leak since _Pickle_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005396 reference to pid first. */
Alexandre Vassalotti20c28c12013-11-27 02:26:54 -08005397 pid = _Pickle_FastCall(self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005398 if (pid == NULL)
5399 return -1;
5400
5401 PDATA_PUSH(self->stack, pid, -1);
5402 return 0;
5403 }
5404 else {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005405 PickleState *st = _Pickle_GetGlobalState();
5406 PyErr_SetString(st->UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005407 "A load persistent id instruction was encountered,\n"
5408 "but no persistent_load function was specified.");
5409 return -1;
5410 }
5411}
5412
5413static int
5414load_pop(UnpicklerObject *self)
5415{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005416 Py_ssize_t len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005417
5418 /* Note that we split the (pickle.py) stack into two stacks,
5419 * an object stack and a mark stack. We have to be clever and
5420 * pop the right one. We do this by looking at the top of the
Collin Winter8ca69de2009-05-26 16:53:41 +00005421 * mark stack first, and only signalling a stack underflow if
5422 * the object stack is empty and the mark stack doesn't match
5423 * our expectations.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005424 */
Collin Winter8ca69de2009-05-26 16:53:41 +00005425 if (self->num_marks > 0 && self->marks[self->num_marks - 1] == len) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005426 self->num_marks--;
Antoine Pitrou01a15ea2010-01-07 17:57:31 +00005427 } else if (len > 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005428 len--;
5429 Py_DECREF(self->stack->data[len]);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005430 Py_SIZE(self->stack) = len;
Collin Winter8ca69de2009-05-26 16:53:41 +00005431 } else {
5432 return stack_underflow();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005433 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005434 return 0;
5435}
5436
5437static int
5438load_pop_mark(UnpicklerObject *self)
5439{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005440 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005441
5442 if ((i = marker(self)) < 0)
5443 return -1;
5444
5445 Pdata_clear(self->stack, i);
5446
5447 return 0;
5448}
5449
5450static int
5451load_dup(UnpicklerObject *self)
5452{
5453 PyObject *last;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005454 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005455
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005456 if ((len = Py_SIZE(self->stack)) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005457 return stack_underflow();
5458 last = self->stack->data[len - 1];
5459 PDATA_APPEND(self->stack, last, -1);
5460 return 0;
5461}
5462
5463static int
5464load_get(UnpicklerObject *self)
5465{
5466 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005467 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005468 Py_ssize_t len;
5469 char *s;
5470
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005471 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005472 return -1;
5473 if (len < 2)
5474 return bad_readline();
5475
5476 key = PyLong_FromString(s, NULL, 10);
5477 if (key == NULL)
5478 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005479 idx = PyLong_AsSsize_t(key);
5480 if (idx == -1 && PyErr_Occurred()) {
5481 Py_DECREF(key);
5482 return -1;
5483 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005484
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005485 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005486 if (value == NULL) {
5487 if (!PyErr_Occurred())
5488 PyErr_SetObject(PyExc_KeyError, key);
5489 Py_DECREF(key);
5490 return -1;
5491 }
5492 Py_DECREF(key);
5493
5494 PDATA_APPEND(self->stack, value, -1);
5495 return 0;
5496}
5497
5498static int
5499load_binget(UnpicklerObject *self)
5500{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005501 PyObject *value;
5502 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005503 char *s;
5504
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005505 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005506 return -1;
5507
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005508 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005509
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005510 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005511 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005512 PyObject *key = PyLong_FromSsize_t(idx);
Christian Heimes9ee5c372013-07-26 22:45:00 +02005513 if (key != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005514 PyErr_SetObject(PyExc_KeyError, key);
Christian Heimes9ee5c372013-07-26 22:45:00 +02005515 Py_DECREF(key);
5516 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005517 return -1;
5518 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005519
5520 PDATA_APPEND(self->stack, value, -1);
5521 return 0;
5522}
5523
5524static int
5525load_long_binget(UnpicklerObject *self)
5526{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005527 PyObject *value;
5528 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005529 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005530
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005531 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005532 return -1;
5533
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005534 idx = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005535
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005536 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005537 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005538 PyObject *key = PyLong_FromSsize_t(idx);
Christian Heimes9ee5c372013-07-26 22:45:00 +02005539 if (key != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005540 PyErr_SetObject(PyExc_KeyError, key);
Christian Heimes9ee5c372013-07-26 22:45:00 +02005541 Py_DECREF(key);
5542 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005543 return -1;
5544 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005545
5546 PDATA_APPEND(self->stack, value, -1);
5547 return 0;
5548}
5549
5550/* Push an object from the extension registry (EXT[124]). nbytes is
5551 * the number of bytes following the opcode, holding the index (code) value.
5552 */
5553static int
5554load_extension(UnpicklerObject *self, int nbytes)
5555{
5556 char *codebytes; /* the nbytes bytes after the opcode */
5557 long code; /* calc_binint returns long */
5558 PyObject *py_code; /* code as a Python int */
5559 PyObject *obj; /* the object to push */
5560 PyObject *pair; /* (module_name, class_name) */
5561 PyObject *module_name, *class_name;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005562 PickleState *st = _Pickle_GetGlobalState();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005563
5564 assert(nbytes == 1 || nbytes == 2 || nbytes == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005565 if (_Unpickler_Read(self, &codebytes, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005566 return -1;
5567 code = calc_binint(codebytes, nbytes);
5568 if (code <= 0) { /* note that 0 is forbidden */
5569 /* Corrupt or hostile pickle. */
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005570 PyErr_SetString(st->UnpicklingError, "EXT specifies code <= 0");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005571 return -1;
5572 }
5573
5574 /* Look for the code in the cache. */
5575 py_code = PyLong_FromLong(code);
5576 if (py_code == NULL)
5577 return -1;
Alexandre Vassalotti567eba12013-11-28 17:09:16 -08005578 obj = PyDict_GetItemWithError(st->extension_cache, py_code);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005579 if (obj != NULL) {
5580 /* Bingo. */
5581 Py_DECREF(py_code);
5582 PDATA_APPEND(self->stack, obj, -1);
5583 return 0;
5584 }
Alexandre Vassalotti567eba12013-11-28 17:09:16 -08005585 if (PyErr_Occurred()) {
5586 Py_DECREF(py_code);
5587 return -1;
5588 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005589
5590 /* Look up the (module_name, class_name) pair. */
Alexandre Vassalotti567eba12013-11-28 17:09:16 -08005591 pair = PyDict_GetItemWithError(st->inverted_registry, py_code);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005592 if (pair == NULL) {
5593 Py_DECREF(py_code);
Alexandre Vassalotti567eba12013-11-28 17:09:16 -08005594 if (!PyErr_Occurred()) {
5595 PyErr_Format(PyExc_ValueError, "unregistered extension "
5596 "code %ld", code);
5597 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005598 return -1;
5599 }
5600 /* Since the extension registry is manipulable via Python code,
5601 * confirm that pair is really a 2-tuple of strings.
5602 */
5603 if (!PyTuple_Check(pair) || PyTuple_Size(pair) != 2 ||
5604 !PyUnicode_Check(module_name = PyTuple_GET_ITEM(pair, 0)) ||
5605 !PyUnicode_Check(class_name = PyTuple_GET_ITEM(pair, 1))) {
5606 Py_DECREF(py_code);
5607 PyErr_Format(PyExc_ValueError, "_inverted_registry[%ld] "
5608 "isn't a 2-tuple of strings", code);
5609 return -1;
5610 }
5611 /* Load the object. */
5612 obj = find_class(self, module_name, class_name);
5613 if (obj == NULL) {
5614 Py_DECREF(py_code);
5615 return -1;
5616 }
5617 /* Cache code -> obj. */
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005618 code = PyDict_SetItem(st->extension_cache, py_code, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005619 Py_DECREF(py_code);
5620 if (code < 0) {
5621 Py_DECREF(obj);
5622 return -1;
5623 }
5624 PDATA_PUSH(self->stack, obj, -1);
5625 return 0;
5626}
5627
5628static int
5629load_put(UnpicklerObject *self)
5630{
5631 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005632 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005633 Py_ssize_t len;
5634 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005635
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005636 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005637 return -1;
5638 if (len < 2)
5639 return bad_readline();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005640 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005641 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005642 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005643
5644 key = PyLong_FromString(s, NULL, 10);
5645 if (key == NULL)
5646 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005647 idx = PyLong_AsSsize_t(key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005648 Py_DECREF(key);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02005649 if (idx < 0) {
5650 if (!PyErr_Occurred())
5651 PyErr_SetString(PyExc_ValueError,
5652 "negative PUT argument");
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005653 return -1;
Antoine Pitrou55549ec2011-08-30 00:27:10 +02005654 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005655
5656 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005657}
5658
5659static int
5660load_binput(UnpicklerObject *self)
5661{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005662 PyObject *value;
5663 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005664 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005665
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005666 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005667 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005668
5669 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005670 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005671 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005672
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005673 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005674
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005675 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005676}
5677
5678static int
5679load_long_binput(UnpicklerObject *self)
5680{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005681 PyObject *value;
5682 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005683 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005684
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005685 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005686 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005687
5688 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005689 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005690 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005691
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005692 idx = calc_binsize(s, 4);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02005693 if (idx < 0) {
5694 PyErr_SetString(PyExc_ValueError,
5695 "negative LONG_BINPUT argument");
5696 return -1;
5697 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005698
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005699 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005700}
5701
5702static int
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005703load_memoize(UnpicklerObject *self)
5704{
5705 PyObject *value;
5706
5707 if (Py_SIZE(self->stack) <= 0)
5708 return stack_underflow();
5709 value = self->stack->data[Py_SIZE(self->stack) - 1];
5710
5711 return _Unpickler_MemoPut(self, self->memo_len, value);
5712}
5713
5714static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005715do_append(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005716{
5717 PyObject *value;
5718 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005719 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005720
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005721 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005722 if (x > len || x <= 0)
5723 return stack_underflow();
5724 if (len == x) /* nothing to do */
5725 return 0;
5726
5727 list = self->stack->data[x - 1];
5728
5729 if (PyList_Check(list)) {
5730 PyObject *slice;
5731 Py_ssize_t list_len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005732 int ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005733
5734 slice = Pdata_poplist(self->stack, x);
5735 if (!slice)
5736 return -1;
5737 list_len = PyList_GET_SIZE(list);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005738 ret = PyList_SetSlice(list, list_len, list_len, slice);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005739 Py_DECREF(slice);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005740 return ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005741 }
5742 else {
5743 PyObject *append_func;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005744 _Py_IDENTIFIER(append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005745
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005746 append_func = _PyObject_GetAttrId(list, &PyId_append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005747 if (append_func == NULL)
5748 return -1;
5749 for (i = x; i < len; i++) {
5750 PyObject *result;
5751
5752 value = self->stack->data[i];
Alexandre Vassalotti20c28c12013-11-27 02:26:54 -08005753 result = _Pickle_FastCall(append_func, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005754 if (result == NULL) {
5755 Pdata_clear(self->stack, i + 1);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005756 Py_SIZE(self->stack) = x;
Alexandre Vassalotti637c7c42013-04-20 21:28:21 -07005757 Py_DECREF(append_func);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005758 return -1;
5759 }
5760 Py_DECREF(result);
5761 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005762 Py_SIZE(self->stack) = x;
Alexandre Vassalotti637c7c42013-04-20 21:28:21 -07005763 Py_DECREF(append_func);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005764 }
5765
5766 return 0;
5767}
5768
5769static int
5770load_append(UnpicklerObject *self)
5771{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005772 return do_append(self, Py_SIZE(self->stack) - 1);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005773}
5774
5775static int
5776load_appends(UnpicklerObject *self)
5777{
5778 return do_append(self, marker(self));
5779}
5780
5781static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005782do_setitems(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005783{
5784 PyObject *value, *key;
5785 PyObject *dict;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005786 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005787 int status = 0;
5788
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005789 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005790 if (x > len || x <= 0)
5791 return stack_underflow();
5792 if (len == x) /* nothing to do */
5793 return 0;
Victor Stinner121aab42011-09-29 23:40:53 +02005794 if ((len - x) % 2 != 0) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005795 PickleState *st = _Pickle_GetGlobalState();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005796 /* Currupt or hostile pickle -- we never write one like this. */
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005797 PyErr_SetString(st->UnpicklingError,
5798 "odd number of items for SETITEMS");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005799 return -1;
5800 }
5801
5802 /* Here, dict does not actually need to be a PyDict; it could be anything
5803 that supports the __setitem__ attribute. */
5804 dict = self->stack->data[x - 1];
5805
5806 for (i = x + 1; i < len; i += 2) {
5807 key = self->stack->data[i - 1];
5808 value = self->stack->data[i];
5809 if (PyObject_SetItem(dict, key, value) < 0) {
5810 status = -1;
5811 break;
5812 }
5813 }
5814
5815 Pdata_clear(self->stack, x);
5816 return status;
5817}
5818
5819static int
5820load_setitem(UnpicklerObject *self)
5821{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005822 return do_setitems(self, Py_SIZE(self->stack) - 2);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005823}
5824
5825static int
5826load_setitems(UnpicklerObject *self)
5827{
5828 return do_setitems(self, marker(self));
5829}
5830
5831static int
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005832load_additems(UnpicklerObject *self)
5833{
5834 PyObject *set;
5835 Py_ssize_t mark, len, i;
5836
5837 mark = marker(self);
5838 len = Py_SIZE(self->stack);
5839 if (mark > len || mark <= 0)
5840 return stack_underflow();
5841 if (len == mark) /* nothing to do */
5842 return 0;
5843
5844 set = self->stack->data[mark - 1];
5845
5846 if (PySet_Check(set)) {
5847 PyObject *items;
5848 int status;
5849
5850 items = Pdata_poptuple(self->stack, mark);
5851 if (items == NULL)
5852 return -1;
5853
5854 status = _PySet_Update(set, items);
5855 Py_DECREF(items);
5856 return status;
5857 }
5858 else {
5859 PyObject *add_func;
5860 _Py_IDENTIFIER(add);
5861
5862 add_func = _PyObject_GetAttrId(set, &PyId_add);
5863 if (add_func == NULL)
5864 return -1;
5865 for (i = mark; i < len; i++) {
5866 PyObject *result;
5867 PyObject *item;
5868
5869 item = self->stack->data[i];
Alexandre Vassalotti20c28c12013-11-27 02:26:54 -08005870 result = _Pickle_FastCall(add_func, item);
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01005871 if (result == NULL) {
5872 Pdata_clear(self->stack, i + 1);
5873 Py_SIZE(self->stack) = mark;
5874 return -1;
5875 }
5876 Py_DECREF(result);
5877 }
5878 Py_SIZE(self->stack) = mark;
5879 }
5880
5881 return 0;
5882}
5883
5884static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005885load_build(UnpicklerObject *self)
5886{
5887 PyObject *state, *inst, *slotstate;
5888 PyObject *setstate;
5889 int status = 0;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005890 _Py_IDENTIFIER(__setstate__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005891
5892 /* Stack is ... instance, state. We want to leave instance at
5893 * the stack top, possibly mutated via instance.__setstate__(state).
5894 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005895 if (Py_SIZE(self->stack) < 2)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005896 return stack_underflow();
5897
5898 PDATA_POP(self->stack, state);
5899 if (state == NULL)
5900 return -1;
5901
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005902 inst = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005903
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005904 setstate = _PyObject_GetAttrId(inst, &PyId___setstate__);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005905 if (setstate == NULL) {
5906 if (PyErr_ExceptionMatches(PyExc_AttributeError))
5907 PyErr_Clear();
Antoine Pitroud79dc622008-09-05 00:03:33 +00005908 else {
5909 Py_DECREF(state);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005910 return -1;
Antoine Pitroud79dc622008-09-05 00:03:33 +00005911 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005912 }
5913 else {
5914 PyObject *result;
5915
5916 /* The explicit __setstate__ is responsible for everything. */
Alexandre Vassalotti20c28c12013-11-27 02:26:54 -08005917 result = _Pickle_FastCall(setstate, state);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005918 Py_DECREF(setstate);
5919 if (result == NULL)
5920 return -1;
5921 Py_DECREF(result);
5922 return 0;
5923 }
5924
5925 /* A default __setstate__. First see whether state embeds a
5926 * slot state dict too (a proto 2 addition).
5927 */
5928 if (PyTuple_Check(state) && Py_SIZE(state) == 2) {
5929 PyObject *tmp = state;
5930
5931 state = PyTuple_GET_ITEM(tmp, 0);
5932 slotstate = PyTuple_GET_ITEM(tmp, 1);
5933 Py_INCREF(state);
5934 Py_INCREF(slotstate);
5935 Py_DECREF(tmp);
5936 }
5937 else
5938 slotstate = NULL;
5939
5940 /* Set inst.__dict__ from the state dict (if any). */
5941 if (state != Py_None) {
5942 PyObject *dict;
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005943 PyObject *d_key, *d_value;
5944 Py_ssize_t i;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005945 _Py_IDENTIFIER(__dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005946
5947 if (!PyDict_Check(state)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005948 PickleState *st = _Pickle_GetGlobalState();
5949 PyErr_SetString(st->UnpicklingError, "state is not a dictionary");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005950 goto error;
5951 }
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005952 dict = _PyObject_GetAttrId(inst, &PyId___dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005953 if (dict == NULL)
5954 goto error;
5955
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005956 i = 0;
5957 while (PyDict_Next(state, &i, &d_key, &d_value)) {
5958 /* normally the keys for instance attributes are
5959 interned. we should try to do that here. */
5960 Py_INCREF(d_key);
5961 if (PyUnicode_CheckExact(d_key))
5962 PyUnicode_InternInPlace(&d_key);
5963 if (PyObject_SetItem(dict, d_key, d_value) < 0) {
5964 Py_DECREF(d_key);
5965 goto error;
5966 }
5967 Py_DECREF(d_key);
5968 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005969 Py_DECREF(dict);
5970 }
5971
5972 /* Also set instance attributes from the slotstate dict (if any). */
5973 if (slotstate != NULL) {
5974 PyObject *d_key, *d_value;
5975 Py_ssize_t i;
5976
5977 if (!PyDict_Check(slotstate)) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08005978 PickleState *st = _Pickle_GetGlobalState();
5979 PyErr_SetString(st->UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005980 "slot state is not a dictionary");
5981 goto error;
5982 }
5983 i = 0;
5984 while (PyDict_Next(slotstate, &i, &d_key, &d_value)) {
5985 if (PyObject_SetAttr(inst, d_key, d_value) < 0)
5986 goto error;
5987 }
5988 }
5989
5990 if (0) {
5991 error:
5992 status = -1;
5993 }
5994
5995 Py_DECREF(state);
5996 Py_XDECREF(slotstate);
5997 return status;
5998}
5999
6000static int
6001load_mark(UnpicklerObject *self)
6002{
6003
6004 /* Note that we split the (pickle.py) stack into two stacks, an
6005 * object stack and a mark stack. Here we push a mark onto the
6006 * mark stack.
6007 */
6008
6009 if ((self->num_marks + 1) >= self->marks_size) {
6010 size_t alloc;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02006011 Py_ssize_t *marks;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006012
6013 /* Use the size_t type to check for overflow. */
6014 alloc = ((size_t)self->num_marks << 1) + 20;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02006015 if (alloc > (PY_SSIZE_T_MAX / sizeof(Py_ssize_t)) ||
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00006016 alloc <= ((size_t)self->num_marks + 1)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006017 PyErr_NoMemory();
6018 return -1;
6019 }
6020
6021 if (self->marks == NULL)
Antoine Pitrou82be19f2011-08-29 23:09:33 +02006022 marks = (Py_ssize_t *) PyMem_Malloc(alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006023 else
Antoine Pitrou82be19f2011-08-29 23:09:33 +02006024 marks = (Py_ssize_t *) PyMem_Realloc(self->marks,
6025 alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006026 if (marks == NULL) {
6027 PyErr_NoMemory();
6028 return -1;
6029 }
6030 self->marks = marks;
6031 self->marks_size = (Py_ssize_t)alloc;
6032 }
6033
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006034 self->marks[self->num_marks++] = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006035
6036 return 0;
6037}
6038
6039static int
6040load_reduce(UnpicklerObject *self)
6041{
6042 PyObject *callable = NULL;
6043 PyObject *argtup = NULL;
6044 PyObject *obj = NULL;
6045
6046 PDATA_POP(self->stack, argtup);
6047 if (argtup == NULL)
6048 return -1;
6049 PDATA_POP(self->stack, callable);
6050 if (callable) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00006051 obj = PyObject_CallObject(callable, argtup);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006052 Py_DECREF(callable);
6053 }
6054 Py_DECREF(argtup);
6055
6056 if (obj == NULL)
6057 return -1;
6058
6059 PDATA_PUSH(self->stack, obj, -1);
6060 return 0;
6061}
6062
6063/* Just raises an error if we don't know the protocol specified. PROTO
6064 * is the first opcode for protocols >= 2.
6065 */
6066static int
6067load_proto(UnpicklerObject *self)
6068{
6069 char *s;
6070 int i;
6071
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006072 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006073 return -1;
6074
6075 i = (unsigned char)s[0];
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006076 if (i <= HIGHEST_PROTOCOL) {
6077 self->proto = i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006078 return 0;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006079 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006080
6081 PyErr_Format(PyExc_ValueError, "unsupported pickle protocol: %d", i);
6082 return -1;
6083}
6084
Alexandre Vassalottib6a2f2a2013-11-23 20:30:03 -08006085static int
6086load_frame(UnpicklerObject *self)
6087{
6088 char *s;
6089 Py_ssize_t frame_len;
6090
6091 if (_Unpickler_Read(self, &s, 8) < 0)
6092 return -1;
6093
6094 frame_len = calc_binsize(s, 8);
6095 if (frame_len < 0) {
6096 PyErr_Format(PyExc_OverflowError,
6097 "FRAME length exceeds system's maximum of %zd bytes",
6098 PY_SSIZE_T_MAX);
6099 return -1;
6100 }
6101
6102 if (_Unpickler_Read(self, &s, frame_len) < 0)
6103 return -1;
6104
6105 /* Rewind to start of frame */
6106 self->next_read_idx -= frame_len;
6107 return 0;
6108}
6109
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006110static PyObject *
6111load(UnpicklerObject *self)
6112{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006113 PyObject *value = NULL;
Christian Heimes27ea78b2014-01-27 01:03:53 +01006114 char *s = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006115
6116 self->num_marks = 0;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01006117 self->proto = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006118 if (Py_SIZE(self->stack))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006119 Pdata_clear(self->stack, 0);
6120
6121 /* Convenient macros for the dispatch while-switch loop just below. */
6122#define OP(opcode, load_func) \
6123 case opcode: if (load_func(self) < 0) break; continue;
6124
6125#define OP_ARG(opcode, load_func, arg) \
6126 case opcode: if (load_func(self, (arg)) < 0) break; continue;
6127
6128 while (1) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006129 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006130 break;
6131
6132 switch ((enum opcode)s[0]) {
6133 OP(NONE, load_none)
6134 OP(BININT, load_binint)
6135 OP(BININT1, load_binint1)
6136 OP(BININT2, load_binint2)
6137 OP(INT, load_int)
6138 OP(LONG, load_long)
6139 OP_ARG(LONG1, load_counted_long, 1)
6140 OP_ARG(LONG4, load_counted_long, 4)
6141 OP(FLOAT, load_float)
6142 OP(BINFLOAT, load_binfloat)
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01006143 OP_ARG(SHORT_BINBYTES, load_counted_binbytes, 1)
6144 OP_ARG(BINBYTES, load_counted_binbytes, 4)
6145 OP_ARG(BINBYTES8, load_counted_binbytes, 8)
6146 OP_ARG(SHORT_BINSTRING, load_counted_binstring, 1)
6147 OP_ARG(BINSTRING, load_counted_binstring, 4)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006148 OP(STRING, load_string)
6149 OP(UNICODE, load_unicode)
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01006150 OP_ARG(SHORT_BINUNICODE, load_counted_binunicode, 1)
6151 OP_ARG(BINUNICODE, load_counted_binunicode, 4)
6152 OP_ARG(BINUNICODE8, load_counted_binunicode, 8)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006153 OP_ARG(EMPTY_TUPLE, load_counted_tuple, 0)
6154 OP_ARG(TUPLE1, load_counted_tuple, 1)
6155 OP_ARG(TUPLE2, load_counted_tuple, 2)
6156 OP_ARG(TUPLE3, load_counted_tuple, 3)
6157 OP(TUPLE, load_tuple)
6158 OP(EMPTY_LIST, load_empty_list)
6159 OP(LIST, load_list)
6160 OP(EMPTY_DICT, load_empty_dict)
6161 OP(DICT, load_dict)
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01006162 OP(EMPTY_SET, load_empty_set)
6163 OP(ADDITEMS, load_additems)
6164 OP(FROZENSET, load_frozenset)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006165 OP(OBJ, load_obj)
6166 OP(INST, load_inst)
6167 OP(NEWOBJ, load_newobj)
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01006168 OP(NEWOBJ_EX, load_newobj_ex)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006169 OP(GLOBAL, load_global)
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01006170 OP(STACK_GLOBAL, load_stack_global)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006171 OP(APPEND, load_append)
6172 OP(APPENDS, load_appends)
6173 OP(BUILD, load_build)
6174 OP(DUP, load_dup)
6175 OP(BINGET, load_binget)
6176 OP(LONG_BINGET, load_long_binget)
6177 OP(GET, load_get)
6178 OP(MARK, load_mark)
6179 OP(BINPUT, load_binput)
6180 OP(LONG_BINPUT, load_long_binput)
6181 OP(PUT, load_put)
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01006182 OP(MEMOIZE, load_memoize)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006183 OP(POP, load_pop)
6184 OP(POP_MARK, load_pop_mark)
6185 OP(SETITEM, load_setitem)
6186 OP(SETITEMS, load_setitems)
6187 OP(PERSID, load_persid)
6188 OP(BINPERSID, load_binpersid)
6189 OP(REDUCE, load_reduce)
6190 OP(PROTO, load_proto)
Alexandre Vassalottib6a2f2a2013-11-23 20:30:03 -08006191 OP(FRAME, load_frame)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006192 OP_ARG(EXT1, load_extension, 1)
6193 OP_ARG(EXT2, load_extension, 2)
6194 OP_ARG(EXT4, load_extension, 4)
6195 OP_ARG(NEWTRUE, load_bool, Py_True)
6196 OP_ARG(NEWFALSE, load_bool, Py_False)
6197
6198 case STOP:
6199 break;
6200
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006201 default:
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08006202 if (s[0] == '\0') {
Benjamin Petersonadde86d2011-09-23 13:41:41 -04006203 PyErr_SetNone(PyExc_EOFError);
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08006204 }
6205 else {
6206 PickleState *st = _Pickle_GetGlobalState();
6207 PyErr_Format(st->UnpicklingError,
Benjamin Petersonadde86d2011-09-23 13:41:41 -04006208 "invalid load key, '%c'.", s[0]);
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08006209 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006210 return NULL;
6211 }
6212
6213 break; /* and we are done! */
6214 }
6215
Alexandre Vassalottib6a2f2a2013-11-23 20:30:03 -08006216 if (PyErr_Occurred()) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006217 return NULL;
6218 }
6219
Victor Stinner2ae57e32013-10-31 13:39:23 +01006220 if (_Unpickler_SkipConsumed(self) < 0)
6221 return NULL;
6222
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006223 PDATA_POP(self->stack, value);
6224 return value;
6225}
6226
Larry Hastings61272b72014-01-07 12:41:53 -08006227/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006228
6229_pickle.Unpickler.load
6230
6231Load a pickle.
6232
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006233Read a pickled object representation from the open file object given
6234in the constructor, and return the reconstituted object hierarchy
6235specified therein.
Larry Hastings61272b72014-01-07 12:41:53 -08006236[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006237
Larry Hastings3cceb382014-01-04 11:09:09 -08006238static PyObject *
Larry Hastingsc2047262014-01-25 20:43:29 -08006239_pickle_Unpickler_load_impl(UnpicklerObject *self)
Larry Hastings581ee362014-01-28 05:00:08 -08006240/*[clinic end generated code: output=fdcc488aad675b14 input=acbb91a42fa9b7b9]*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006241{
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006242 UnpicklerObject *unpickler = (UnpicklerObject*)self;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08006243
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006244 /* Check whether the Unpickler was initialized correctly. This prevents
6245 segfaulting if a subclass overridden __init__ with a function that does
6246 not call Unpickler.__init__(). Here, we simply ensure that self->read
6247 is not NULL. */
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006248 if (unpickler->read == NULL) {
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08006249 PickleState *st = _Pickle_GetGlobalState();
6250 PyErr_Format(st->UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006251 "Unpickler.__init__() was not called by %s.__init__()",
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006252 Py_TYPE(unpickler)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006253 return NULL;
6254 }
6255
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006256 return load(unpickler);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006257}
6258
6259/* The name of find_class() is misleading. In newer pickle protocols, this
6260 function is used for loading any global (i.e., functions), not just
6261 classes. The name is kept only for backward compatibility. */
6262
Larry Hastings61272b72014-01-07 12:41:53 -08006263/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006264
6265_pickle.Unpickler.find_class
6266
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006267 module_name: object
6268 global_name: object
6269 /
6270
6271Return an object from a specified module.
6272
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006273If necessary, the module will be imported. Subclasses may override
6274this method (e.g. to restrict unpickling of arbitrary classes and
6275functions).
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006276
6277This method is called whenever a class or a function object is
6278needed. Both arguments passed are str objects.
Larry Hastings61272b72014-01-07 12:41:53 -08006279[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006280
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006281static PyObject *
6282_pickle_Unpickler_find_class_impl(UnpicklerObject *self, PyObject *module_name, PyObject *global_name)
Larry Hastings581ee362014-01-28 05:00:08 -08006283/*[clinic end generated code: output=64c77437e088e188 input=e2e6a865de093ef4]*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006284{
6285 PyObject *global;
6286 PyObject *modules_dict;
6287 PyObject *module;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01006288 _Py_IDENTIFIER(modules);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006289
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006290 /* Try to map the old names used in Python 2.x to the new ones used in
6291 Python 3.x. We do this only with old pickle protocols and when the
6292 user has not disabled the feature. */
6293 if (self->proto < 3 && self->fix_imports) {
6294 PyObject *key;
6295 PyObject *item;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08006296 PickleState *st = _Pickle_GetGlobalState();
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006297
6298 /* Check if the global (i.e., a function or a class) was renamed
6299 or moved to another module. */
6300 key = PyTuple_Pack(2, module_name, global_name);
6301 if (key == NULL)
6302 return NULL;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08006303 item = PyDict_GetItemWithError(st->name_mapping_2to3, key);
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006304 Py_DECREF(key);
6305 if (item) {
6306 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
6307 PyErr_Format(PyExc_RuntimeError,
6308 "_compat_pickle.NAME_MAPPING values should be "
6309 "2-tuples, not %.200s", Py_TYPE(item)->tp_name);
6310 return NULL;
6311 }
6312 module_name = PyTuple_GET_ITEM(item, 0);
6313 global_name = PyTuple_GET_ITEM(item, 1);
6314 if (!PyUnicode_Check(module_name) ||
6315 !PyUnicode_Check(global_name)) {
6316 PyErr_Format(PyExc_RuntimeError,
6317 "_compat_pickle.NAME_MAPPING values should be "
6318 "pairs of str, not (%.200s, %.200s)",
6319 Py_TYPE(module_name)->tp_name,
6320 Py_TYPE(global_name)->tp_name);
6321 return NULL;
6322 }
6323 }
6324 else if (PyErr_Occurred()) {
6325 return NULL;
6326 }
6327
6328 /* Check if the module was renamed. */
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08006329 item = PyDict_GetItemWithError(st->import_mapping_2to3, module_name);
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006330 if (item) {
6331 if (!PyUnicode_Check(item)) {
6332 PyErr_Format(PyExc_RuntimeError,
6333 "_compat_pickle.IMPORT_MAPPING values should be "
6334 "strings, not %.200s", Py_TYPE(item)->tp_name);
6335 return NULL;
6336 }
6337 module_name = item;
6338 }
6339 else if (PyErr_Occurred()) {
6340 return NULL;
6341 }
6342 }
6343
Victor Stinnerbb520202013-11-06 22:40:41 +01006344 modules_dict = _PySys_GetObjectId(&PyId_modules);
Victor Stinner1e53bba2013-07-16 22:26:05 +02006345 if (modules_dict == NULL) {
6346 PyErr_SetString(PyExc_RuntimeError, "unable to get sys.modules");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006347 return NULL;
Victor Stinner1e53bba2013-07-16 22:26:05 +02006348 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006349
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006350 module = PyDict_GetItemWithError(modules_dict, module_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006351 if (module == NULL) {
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006352 if (PyErr_Occurred())
6353 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006354 module = PyImport_Import(module_name);
6355 if (module == NULL)
6356 return NULL;
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01006357 global = getattribute(module, global_name, self->proto >= 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006358 Py_DECREF(module);
6359 }
Victor Stinner121aab42011-09-29 23:40:53 +02006360 else {
Antoine Pitrouc9dc4a22013-11-23 18:59:12 +01006361 global = getattribute(module, global_name, self->proto >= 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006362 }
6363 return global;
6364}
6365
Serhiy Storchaka5bbd2312014-12-16 19:39:08 +02006366/*[clinic input]
6367
6368_pickle.Unpickler.__sizeof__ -> Py_ssize_t
6369
6370Returns size in memory, in bytes.
6371[clinic start generated code]*/
6372
6373static Py_ssize_t
6374_pickle_Unpickler___sizeof___impl(UnpicklerObject *self)
6375/*[clinic end generated code: output=119d9d03ad4c7651 input=13333471fdeedf5e]*/
6376{
6377 Py_ssize_t res;
6378
6379 res = sizeof(UnpicklerObject);
6380 if (self->memo != NULL)
6381 res += self->memo_size * sizeof(PyObject *);
6382 if (self->marks != NULL)
6383 res += self->marks_size * sizeof(Py_ssize_t);
6384 if (self->input_line != NULL)
6385 res += strlen(self->input_line) + 1;
6386 if (self->encoding != NULL)
6387 res += strlen(self->encoding) + 1;
6388 if (self->errors != NULL)
6389 res += strlen(self->errors) + 1;
6390 return res;
6391}
6392
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006393static struct PyMethodDef Unpickler_methods[] = {
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006394 _PICKLE_UNPICKLER_LOAD_METHODDEF
6395 _PICKLE_UNPICKLER_FIND_CLASS_METHODDEF
Serhiy Storchaka5bbd2312014-12-16 19:39:08 +02006396 _PICKLE_UNPICKLER___SIZEOF___METHODDEF
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006397 {NULL, NULL} /* sentinel */
6398};
6399
6400static void
6401Unpickler_dealloc(UnpicklerObject *self)
6402{
6403 PyObject_GC_UnTrack((PyObject *)self);
6404 Py_XDECREF(self->readline);
6405 Py_XDECREF(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00006406 Py_XDECREF(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006407 Py_XDECREF(self->stack);
6408 Py_XDECREF(self->pers_func);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006409 if (self->buffer.buf != NULL) {
6410 PyBuffer_Release(&self->buffer);
6411 self->buffer.buf = NULL;
6412 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006413
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006414 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006415 PyMem_Free(self->marks);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006416 PyMem_Free(self->input_line);
Victor Stinner49fc8ec2013-07-07 23:30:24 +02006417 PyMem_Free(self->encoding);
6418 PyMem_Free(self->errors);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006419
6420 Py_TYPE(self)->tp_free((PyObject *)self);
6421}
6422
6423static int
6424Unpickler_traverse(UnpicklerObject *self, visitproc visit, void *arg)
6425{
6426 Py_VISIT(self->readline);
6427 Py_VISIT(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00006428 Py_VISIT(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006429 Py_VISIT(self->stack);
6430 Py_VISIT(self->pers_func);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006431 return 0;
6432}
6433
6434static int
6435Unpickler_clear(UnpicklerObject *self)
6436{
6437 Py_CLEAR(self->readline);
6438 Py_CLEAR(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00006439 Py_CLEAR(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006440 Py_CLEAR(self->stack);
6441 Py_CLEAR(self->pers_func);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006442 if (self->buffer.buf != NULL) {
6443 PyBuffer_Release(&self->buffer);
6444 self->buffer.buf = NULL;
6445 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006446
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006447 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006448 PyMem_Free(self->marks);
6449 self->marks = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006450 PyMem_Free(self->input_line);
6451 self->input_line = NULL;
Victor Stinner49fc8ec2013-07-07 23:30:24 +02006452 PyMem_Free(self->encoding);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006453 self->encoding = NULL;
Victor Stinner49fc8ec2013-07-07 23:30:24 +02006454 PyMem_Free(self->errors);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006455 self->errors = NULL;
6456
6457 return 0;
6458}
6459
Larry Hastings61272b72014-01-07 12:41:53 -08006460/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006461
6462_pickle.Unpickler.__init__
6463
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006464 file: object
6465 *
6466 fix_imports: bool = True
6467 encoding: str = 'ASCII'
6468 errors: str = 'strict'
6469
6470This takes a binary file for reading a pickle data stream.
6471
6472The protocol version of the pickle is detected automatically, so no
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006473protocol argument is needed. Bytes past the pickled object's
6474representation are ignored.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006475
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006476The argument *file* must have two methods, a read() method that takes
6477an integer argument, and a readline() method that requires no
6478arguments. Both methods should return bytes. Thus *file* can be a
6479binary file object opened for reading, a io.BytesIO object, or any
6480other custom object that meets this interface.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006481
6482Optional keyword arguments are *fix_imports*, *encoding* and *errors*,
6483which are used to control compatiblity support for pickle stream
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006484generated by Python 2. If *fix_imports* is True, pickle will try to
6485map the old Python 2 names to the new names used in Python 3. The
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006486*encoding* and *errors* tell pickle how to decode 8-bit string
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006487instances pickled by Python 2; these default to 'ASCII' and 'strict',
6488respectively. The *encoding* can be 'bytes' to read these 8-bit
6489string instances as bytes objects.
Larry Hastings61272b72014-01-07 12:41:53 -08006490[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006491
Larry Hastingsb7ccb202014-01-18 23:50:21 -08006492static int
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006493_pickle_Unpickler___init___impl(UnpicklerObject *self, PyObject *file, int fix_imports, const char *encoding, const char *errors)
Larry Hastings581ee362014-01-28 05:00:08 -08006494/*[clinic end generated code: output=b9ed1d84d315f3b5 input=30b4dc9e976b890c]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006495{
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02006496 _Py_IDENTIFIER(persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006497
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006498 /* In case of multiple __init__() calls, clear previous content. */
6499 if (self->read != NULL)
6500 (void)Unpickler_clear(self);
6501
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006502 if (_Unpickler_SetInputStream(self, file) < 0)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08006503 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006504
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006505 if (_Unpickler_SetInputEncoding(self, encoding, errors) < 0)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08006506 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006507
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006508 self->fix_imports = fix_imports;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006509 if (self->fix_imports == -1)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08006510 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006511
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02006512 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_load)) {
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02006513 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
6514 &PyId_persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006515 if (self->pers_func == NULL)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08006516 return 1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006517 }
6518 else {
6519 self->pers_func = NULL;
6520 }
6521
6522 self->stack = (Pdata *)Pdata_New();
6523 if (self->stack == NULL)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08006524 return 1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006525
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006526 self->memo_size = 32;
6527 self->memo = _Unpickler_NewMemo(self->memo_size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006528 if (self->memo == NULL)
Larry Hastingsb7ccb202014-01-18 23:50:21 -08006529 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006530
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006531 self->proto = 0;
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00006532
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006533 return 0;
6534}
6535
Larry Hastingsb7ccb202014-01-18 23:50:21 -08006536
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006537/* Define a proxy object for the Unpickler's internal memo object. This is to
6538 * avoid breaking code like:
6539 * unpickler.memo.clear()
6540 * and
6541 * unpickler.memo = saved_memo
6542 * Is this a good idea? Not really, but we don't want to break code that uses
6543 * it. Note that we don't implement the entire mapping API here. This is
6544 * intentional, as these should be treated as black-box implementation details.
6545 *
6546 * We do, however, have to implement pickling/unpickling support because of
Victor Stinner121aab42011-09-29 23:40:53 +02006547 * real-world code like cvs2svn.
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006548 */
6549
Larry Hastings61272b72014-01-07 12:41:53 -08006550/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006551_pickle.UnpicklerMemoProxy.clear
6552
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006553Remove all items from memo.
Larry Hastings61272b72014-01-07 12:41:53 -08006554[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006555
Larry Hastings3cceb382014-01-04 11:09:09 -08006556static PyObject *
6557_pickle_UnpicklerMemoProxy_clear_impl(UnpicklerMemoProxyObject *self)
Larry Hastings581ee362014-01-28 05:00:08 -08006558/*[clinic end generated code: output=d20cd43f4ba1fb1f input=b1df7c52e7afd9bd]*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006559{
6560 _Unpickler_MemoCleanup(self->unpickler);
6561 self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size);
6562 if (self->unpickler->memo == NULL)
6563 return NULL;
6564 Py_RETURN_NONE;
6565}
6566
Larry Hastings61272b72014-01-07 12:41:53 -08006567/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006568_pickle.UnpicklerMemoProxy.copy
6569
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006570Copy the memo to a new object.
Larry Hastings61272b72014-01-07 12:41:53 -08006571[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006572
Larry Hastings3cceb382014-01-04 11:09:09 -08006573static PyObject *
6574_pickle_UnpicklerMemoProxy_copy_impl(UnpicklerMemoProxyObject *self)
Larry Hastings581ee362014-01-28 05:00:08 -08006575/*[clinic end generated code: output=e12af7e9bc1e4c77 input=97769247ce032c1d]*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006576{
6577 Py_ssize_t i;
6578 PyObject *new_memo = PyDict_New();
6579 if (new_memo == NULL)
6580 return NULL;
6581
6582 for (i = 0; i < self->unpickler->memo_size; i++) {
6583 int status;
6584 PyObject *key, *value;
6585
6586 value = self->unpickler->memo[i];
6587 if (value == NULL)
6588 continue;
6589
6590 key = PyLong_FromSsize_t(i);
6591 if (key == NULL)
6592 goto error;
6593 status = PyDict_SetItem(new_memo, key, value);
6594 Py_DECREF(key);
6595 if (status < 0)
6596 goto error;
6597 }
6598 return new_memo;
6599
6600error:
6601 Py_DECREF(new_memo);
6602 return NULL;
6603}
6604
Larry Hastings61272b72014-01-07 12:41:53 -08006605/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006606_pickle.UnpicklerMemoProxy.__reduce__
6607
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006608Implement pickling support.
Larry Hastings61272b72014-01-07 12:41:53 -08006609[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006610
Larry Hastings3cceb382014-01-04 11:09:09 -08006611static PyObject *
6612_pickle_UnpicklerMemoProxy___reduce___impl(UnpicklerMemoProxyObject *self)
Larry Hastings581ee362014-01-28 05:00:08 -08006613/*[clinic end generated code: output=6da34ac048d94cca input=6920862413407199]*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006614{
6615 PyObject *reduce_value;
6616 PyObject *constructor_args;
Larry Hastings3cceb382014-01-04 11:09:09 -08006617 PyObject *contents = _pickle_UnpicklerMemoProxy_copy_impl(self);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006618 if (contents == NULL)
6619 return NULL;
6620
6621 reduce_value = PyTuple_New(2);
6622 if (reduce_value == NULL) {
6623 Py_DECREF(contents);
6624 return NULL;
6625 }
6626 constructor_args = PyTuple_New(1);
6627 if (constructor_args == NULL) {
6628 Py_DECREF(contents);
6629 Py_DECREF(reduce_value);
6630 return NULL;
6631 }
6632 PyTuple_SET_ITEM(constructor_args, 0, contents);
6633 Py_INCREF((PyObject *)&PyDict_Type);
6634 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
6635 PyTuple_SET_ITEM(reduce_value, 1, constructor_args);
6636 return reduce_value;
6637}
6638
6639static PyMethodDef unpicklerproxy_methods[] = {
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006640 _PICKLE_UNPICKLERMEMOPROXY_CLEAR_METHODDEF
6641 _PICKLE_UNPICKLERMEMOPROXY_COPY_METHODDEF
6642 _PICKLE_UNPICKLERMEMOPROXY___REDUCE___METHODDEF
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006643 {NULL, NULL} /* sentinel */
6644};
6645
6646static void
6647UnpicklerMemoProxy_dealloc(UnpicklerMemoProxyObject *self)
6648{
6649 PyObject_GC_UnTrack(self);
6650 Py_XDECREF(self->unpickler);
6651 PyObject_GC_Del((PyObject *)self);
6652}
6653
6654static int
6655UnpicklerMemoProxy_traverse(UnpicklerMemoProxyObject *self,
6656 visitproc visit, void *arg)
6657{
6658 Py_VISIT(self->unpickler);
6659 return 0;
6660}
6661
6662static int
6663UnpicklerMemoProxy_clear(UnpicklerMemoProxyObject *self)
6664{
6665 Py_CLEAR(self->unpickler);
6666 return 0;
6667}
6668
6669static PyTypeObject UnpicklerMemoProxyType = {
6670 PyVarObject_HEAD_INIT(NULL, 0)
6671 "_pickle.UnpicklerMemoProxy", /*tp_name*/
6672 sizeof(UnpicklerMemoProxyObject), /*tp_basicsize*/
6673 0,
6674 (destructor)UnpicklerMemoProxy_dealloc, /* tp_dealloc */
6675 0, /* tp_print */
6676 0, /* tp_getattr */
6677 0, /* tp_setattr */
6678 0, /* tp_compare */
6679 0, /* tp_repr */
6680 0, /* tp_as_number */
6681 0, /* tp_as_sequence */
6682 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00006683 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006684 0, /* tp_call */
6685 0, /* tp_str */
6686 PyObject_GenericGetAttr, /* tp_getattro */
6687 PyObject_GenericSetAttr, /* tp_setattro */
6688 0, /* tp_as_buffer */
6689 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
6690 0, /* tp_doc */
6691 (traverseproc)UnpicklerMemoProxy_traverse, /* tp_traverse */
6692 (inquiry)UnpicklerMemoProxy_clear, /* tp_clear */
6693 0, /* tp_richcompare */
6694 0, /* tp_weaklistoffset */
6695 0, /* tp_iter */
6696 0, /* tp_iternext */
6697 unpicklerproxy_methods, /* tp_methods */
6698};
6699
6700static PyObject *
6701UnpicklerMemoProxy_New(UnpicklerObject *unpickler)
6702{
6703 UnpicklerMemoProxyObject *self;
6704
6705 self = PyObject_GC_New(UnpicklerMemoProxyObject,
6706 &UnpicklerMemoProxyType);
6707 if (self == NULL)
6708 return NULL;
6709 Py_INCREF(unpickler);
6710 self->unpickler = unpickler;
6711 PyObject_GC_Track(self);
6712 return (PyObject *)self;
6713}
6714
6715/*****************************************************************************/
6716
6717
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006718static PyObject *
6719Unpickler_get_memo(UnpicklerObject *self)
6720{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006721 return UnpicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006722}
6723
6724static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006725Unpickler_set_memo(UnpicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006726{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006727 PyObject **new_memo;
6728 Py_ssize_t new_memo_size = 0;
6729 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006730
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006731 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006732 PyErr_SetString(PyExc_TypeError,
6733 "attribute deletion is not supported");
6734 return -1;
6735 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006736
6737 if (Py_TYPE(obj) == &UnpicklerMemoProxyType) {
6738 UnpicklerObject *unpickler =
6739 ((UnpicklerMemoProxyObject *)obj)->unpickler;
6740
6741 new_memo_size = unpickler->memo_size;
6742 new_memo = _Unpickler_NewMemo(new_memo_size);
6743 if (new_memo == NULL)
6744 return -1;
6745
6746 for (i = 0; i < new_memo_size; i++) {
6747 Py_XINCREF(unpickler->memo[i]);
6748 new_memo[i] = unpickler->memo[i];
6749 }
6750 }
6751 else if (PyDict_Check(obj)) {
6752 Py_ssize_t i = 0;
6753 PyObject *key, *value;
6754
6755 new_memo_size = PyDict_Size(obj);
6756 new_memo = _Unpickler_NewMemo(new_memo_size);
6757 if (new_memo == NULL)
6758 return -1;
6759
6760 while (PyDict_Next(obj, &i, &key, &value)) {
6761 Py_ssize_t idx;
6762 if (!PyLong_Check(key)) {
6763 PyErr_SetString(PyExc_TypeError,
6764 "memo key must be integers");
6765 goto error;
6766 }
6767 idx = PyLong_AsSsize_t(key);
6768 if (idx == -1 && PyErr_Occurred())
6769 goto error;
Christian Heimesa24b4d22013-07-01 15:17:45 +02006770 if (idx < 0) {
6771 PyErr_SetString(PyExc_ValueError,
Christian Heimes80878792013-07-01 15:23:39 +02006772 "memo key must be positive integers.");
Christian Heimesa24b4d22013-07-01 15:17:45 +02006773 goto error;
6774 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006775 if (_Unpickler_MemoPut(self, idx, value) < 0)
6776 goto error;
6777 }
6778 }
6779 else {
6780 PyErr_Format(PyExc_TypeError,
6781 "'memo' attribute must be an UnpicklerMemoProxy object"
6782 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006783 return -1;
6784 }
6785
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006786 _Unpickler_MemoCleanup(self);
6787 self->memo_size = new_memo_size;
6788 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006789
6790 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006791
6792 error:
6793 if (new_memo_size) {
6794 i = new_memo_size;
6795 while (--i >= 0) {
6796 Py_XDECREF(new_memo[i]);
6797 }
6798 PyMem_FREE(new_memo);
6799 }
6800 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006801}
6802
6803static PyObject *
6804Unpickler_get_persload(UnpicklerObject *self)
6805{
6806 if (self->pers_func == NULL)
6807 PyErr_SetString(PyExc_AttributeError, "persistent_load");
6808 else
6809 Py_INCREF(self->pers_func);
6810 return self->pers_func;
6811}
6812
6813static int
6814Unpickler_set_persload(UnpicklerObject *self, PyObject *value)
6815{
6816 PyObject *tmp;
6817
6818 if (value == NULL) {
6819 PyErr_SetString(PyExc_TypeError,
6820 "attribute deletion is not supported");
6821 return -1;
6822 }
6823 if (!PyCallable_Check(value)) {
6824 PyErr_SetString(PyExc_TypeError,
6825 "persistent_load must be a callable taking "
6826 "one argument");
6827 return -1;
6828 }
6829
6830 tmp = self->pers_func;
6831 Py_INCREF(value);
6832 self->pers_func = value;
6833 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
6834
6835 return 0;
6836}
6837
6838static PyGetSetDef Unpickler_getsets[] = {
6839 {"memo", (getter)Unpickler_get_memo, (setter)Unpickler_set_memo},
6840 {"persistent_load", (getter)Unpickler_get_persload,
6841 (setter)Unpickler_set_persload},
6842 {NULL}
6843};
6844
6845static PyTypeObject Unpickler_Type = {
6846 PyVarObject_HEAD_INIT(NULL, 0)
6847 "_pickle.Unpickler", /*tp_name*/
6848 sizeof(UnpicklerObject), /*tp_basicsize*/
6849 0, /*tp_itemsize*/
6850 (destructor)Unpickler_dealloc, /*tp_dealloc*/
6851 0, /*tp_print*/
6852 0, /*tp_getattr*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006853 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00006854 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006855 0, /*tp_repr*/
6856 0, /*tp_as_number*/
6857 0, /*tp_as_sequence*/
6858 0, /*tp_as_mapping*/
6859 0, /*tp_hash*/
6860 0, /*tp_call*/
6861 0, /*tp_str*/
6862 0, /*tp_getattro*/
6863 0, /*tp_setattro*/
6864 0, /*tp_as_buffer*/
6865 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006866 _pickle_Unpickler___init____doc__, /*tp_doc*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006867 (traverseproc)Unpickler_traverse, /*tp_traverse*/
6868 (inquiry)Unpickler_clear, /*tp_clear*/
6869 0, /*tp_richcompare*/
6870 0, /*tp_weaklistoffset*/
6871 0, /*tp_iter*/
6872 0, /*tp_iternext*/
6873 Unpickler_methods, /*tp_methods*/
6874 0, /*tp_members*/
6875 Unpickler_getsets, /*tp_getset*/
6876 0, /*tp_base*/
6877 0, /*tp_dict*/
6878 0, /*tp_descr_get*/
6879 0, /*tp_descr_set*/
6880 0, /*tp_dictoffset*/
Larry Hastingsb7ccb202014-01-18 23:50:21 -08006881 _pickle_Unpickler___init__, /*tp_init*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006882 PyType_GenericAlloc, /*tp_alloc*/
6883 PyType_GenericNew, /*tp_new*/
6884 PyObject_GC_Del, /*tp_free*/
6885 0, /*tp_is_gc*/
6886};
6887
Larry Hastings61272b72014-01-07 12:41:53 -08006888/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006889
6890_pickle.dump
6891
6892 obj: object
6893 file: object
6894 protocol: object = NULL
6895 *
6896 fix_imports: bool = True
6897
6898Write a pickled representation of obj to the open file object file.
6899
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006900This is equivalent to ``Pickler(file, protocol).dump(obj)``, but may
6901be more efficient.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006902
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006903The optional *protocol* argument tells the pickler to use the given
6904protocol supported protocols are 0, 1, 2, 3 and 4. The default
6905protocol is 3; a backward-incompatible protocol designed for Python 3.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006906
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006907Specifying a negative protocol version selects the highest protocol
6908version supported. The higher the protocol used, the more recent the
6909version of Python needed to read the pickle produced.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006910
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006911The *file* argument must have a write() method that accepts a single
6912bytes argument. It can thus be a file object opened for binary
6913writing, a io.BytesIO instance, or any other custom object that meets
6914this interface.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006915
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006916If *fix_imports* is True and protocol is less than 3, pickle will try
6917to map the new Python 3 names to the old module names used in Python
69182, so that the pickle data stream is readable with Python 2.
Larry Hastings61272b72014-01-07 12:41:53 -08006919[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006920
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006921static PyObject *
6922_pickle_dump_impl(PyModuleDef *module, PyObject *obj, PyObject *file, PyObject *protocol, int fix_imports)
Larry Hastings581ee362014-01-28 05:00:08 -08006923/*[clinic end generated code: output=a606e626d553850d input=e9e5fdd48de92eae]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006924{
6925 PicklerObject *pickler = _Pickler_New();
6926
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006927 if (pickler == NULL)
6928 return NULL;
6929
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006930 if (_Pickler_SetProtocol(pickler, protocol, fix_imports) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006931 goto error;
6932
6933 if (_Pickler_SetOutputStream(pickler, file) < 0)
6934 goto error;
6935
6936 if (dump(pickler, obj) < 0)
6937 goto error;
6938
6939 if (_Pickler_FlushToFile(pickler) < 0)
6940 goto error;
6941
6942 Py_DECREF(pickler);
6943 Py_RETURN_NONE;
6944
6945 error:
6946 Py_XDECREF(pickler);
6947 return NULL;
6948}
6949
Larry Hastings61272b72014-01-07 12:41:53 -08006950/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006951
6952_pickle.dumps
6953
6954 obj: object
6955 protocol: object = NULL
6956 *
6957 fix_imports: bool = True
6958
6959Return the pickled representation of the object as a bytes object.
6960
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006961The optional *protocol* argument tells the pickler to use the given
6962protocol; supported protocols are 0, 1, 2, 3 and 4. The default
6963protocol is 3; a backward-incompatible protocol designed for Python 3.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006964
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006965Specifying a negative protocol version selects the highest protocol
6966version supported. The higher the protocol used, the more recent the
6967version of Python needed to read the pickle produced.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006968
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08006969If *fix_imports* is True and *protocol* is less than 3, pickle will
6970try to map the new Python 3 names to the old module names used in
6971Python 2, so that the pickle data stream is readable with Python 2.
Larry Hastings61272b72014-01-07 12:41:53 -08006972[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006973
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006974static PyObject *
6975_pickle_dumps_impl(PyModuleDef *module, PyObject *obj, PyObject *protocol, int fix_imports)
Larry Hastings581ee362014-01-28 05:00:08 -08006976/*[clinic end generated code: output=777f0deefe5b88ee input=293dbeda181580b7]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006977{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006978 PyObject *result;
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006979 PicklerObject *pickler = _Pickler_New();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006980
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006981 if (pickler == NULL)
6982 return NULL;
6983
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08006984 if (_Pickler_SetProtocol(pickler, protocol, fix_imports) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006985 goto error;
6986
6987 if (dump(pickler, obj) < 0)
6988 goto error;
6989
6990 result = _Pickler_GetString(pickler);
6991 Py_DECREF(pickler);
6992 return result;
6993
6994 error:
6995 Py_XDECREF(pickler);
6996 return NULL;
6997}
6998
Larry Hastings61272b72014-01-07 12:41:53 -08006999/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007000
7001_pickle.load
7002
7003 file: object
7004 *
7005 fix_imports: bool = True
7006 encoding: str = 'ASCII'
7007 errors: str = 'strict'
7008
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08007009Read and return an object from the pickle data stored in a file.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007010
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08007011This is equivalent to ``Unpickler(file).load()``, but may be more
7012efficient.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007013
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08007014The protocol version of the pickle is detected automatically, so no
7015protocol argument is needed. Bytes past the pickled object's
7016representation are ignored.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007017
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08007018The argument *file* must have two methods, a read() method that takes
7019an integer argument, and a readline() method that requires no
7020arguments. Both methods should return bytes. Thus *file* can be a
7021binary file object opened for reading, a io.BytesIO object, or any
7022other custom object that meets this interface.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007023
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08007024Optional keyword arguments are *fix_imports*, *encoding* and *errors*,
7025which are used to control compatiblity support for pickle stream
7026generated by Python 2. If *fix_imports* is True, pickle will try to
7027map the old Python 2 names to the new names used in Python 3. The
7028*encoding* and *errors* tell pickle how to decode 8-bit string
7029instances pickled by Python 2; these default to 'ASCII' and 'strict',
7030respectively. The *encoding* can be 'bytes' to read these 8-bit
7031string instances as bytes objects.
Larry Hastings61272b72014-01-07 12:41:53 -08007032[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007033
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007034static PyObject *
7035_pickle_load_impl(PyModuleDef *module, PyObject *file, int fix_imports, const char *encoding, const char *errors)
Larry Hastings581ee362014-01-28 05:00:08 -08007036/*[clinic end generated code: output=568c61356c172654 input=da97372e38e510a6]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007037{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007038 PyObject *result;
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007039 UnpicklerObject *unpickler = _Unpickler_New();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007040
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007041 if (unpickler == NULL)
7042 return NULL;
7043
7044 if (_Unpickler_SetInputStream(unpickler, file) < 0)
7045 goto error;
7046
7047 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
7048 goto error;
7049
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007050 unpickler->fix_imports = fix_imports;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007051
7052 result = load(unpickler);
7053 Py_DECREF(unpickler);
7054 return result;
7055
7056 error:
7057 Py_XDECREF(unpickler);
7058 return NULL;
7059}
7060
Larry Hastings61272b72014-01-07 12:41:53 -08007061/*[clinic input]
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007062
7063_pickle.loads
7064
7065 data: object
7066 *
7067 fix_imports: bool = True
7068 encoding: str = 'ASCII'
7069 errors: str = 'strict'
7070
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08007071Read and return an object from the given pickle data.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007072
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08007073The protocol version of the pickle is detected automatically, so no
7074protocol argument is needed. Bytes past the pickled object's
7075representation are ignored.
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007076
Alexandre Vassalottid05c9ff2013-12-07 01:09:27 -08007077Optional keyword arguments are *fix_imports*, *encoding* and *errors*,
7078which are used to control compatiblity support for pickle stream
7079generated by Python 2. If *fix_imports* is True, pickle will try to
7080map the old Python 2 names to the new names used in Python 3. The
7081*encoding* and *errors* tell pickle how to decode 8-bit string
7082instances pickled by Python 2; these default to 'ASCII' and 'strict',
7083respectively. The *encoding* can be 'bytes' to read these 8-bit
7084string instances as bytes objects.
Larry Hastings61272b72014-01-07 12:41:53 -08007085[clinic start generated code]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007086
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007087static PyObject *
7088_pickle_loads_impl(PyModuleDef *module, PyObject *data, int fix_imports, const char *encoding, const char *errors)
Larry Hastings581ee362014-01-28 05:00:08 -08007089/*[clinic end generated code: output=0b3845ad110b2522 input=f57f0fdaa2b4cb8b]*/
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007090{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007091 PyObject *result;
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007092 UnpicklerObject *unpickler = _Unpickler_New();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007093
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007094 if (unpickler == NULL)
7095 return NULL;
7096
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007097 if (_Unpickler_SetStringInput(unpickler, data) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007098 goto error;
7099
7100 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
7101 goto error;
7102
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007103 unpickler->fix_imports = fix_imports;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007104
7105 result = load(unpickler);
7106 Py_DECREF(unpickler);
7107 return result;
7108
7109 error:
7110 Py_XDECREF(unpickler);
7111 return NULL;
7112}
7113
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007114static struct PyMethodDef pickle_methods[] = {
Alexandre Vassalottied8c9062013-11-24 12:25:48 -08007115 _PICKLE_DUMP_METHODDEF
7116 _PICKLE_DUMPS_METHODDEF
7117 _PICKLE_LOAD_METHODDEF
7118 _PICKLE_LOADS_METHODDEF
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007119 {NULL, NULL} /* sentinel */
7120};
7121
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007122static int
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007123pickle_clear(PyObject *m)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007124{
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007125 _Pickle_ClearState(_Pickle_GetState(m));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007126 return 0;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007127}
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007128
Stefan Krahf483b0f2013-12-14 13:43:10 +01007129static void
7130pickle_free(PyObject *m)
7131{
7132 _Pickle_ClearState(_Pickle_GetState(m));
7133}
7134
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007135static int
7136pickle_traverse(PyObject *m, visitproc visit, void *arg)
7137{
7138 PickleState *st = _Pickle_GetState(m);
7139 Py_VISIT(st->PickleError);
7140 Py_VISIT(st->PicklingError);
7141 Py_VISIT(st->UnpicklingError);
7142 Py_VISIT(st->dispatch_table);
7143 Py_VISIT(st->extension_registry);
7144 Py_VISIT(st->extension_cache);
7145 Py_VISIT(st->inverted_registry);
7146 Py_VISIT(st->name_mapping_2to3);
7147 Py_VISIT(st->import_mapping_2to3);
7148 Py_VISIT(st->name_mapping_3to2);
7149 Py_VISIT(st->import_mapping_3to2);
7150 Py_VISIT(st->codecs_encode);
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007151 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007152}
7153
7154static struct PyModuleDef _picklemodule = {
7155 PyModuleDef_HEAD_INIT,
Stefan Krahf483b0f2013-12-14 13:43:10 +01007156 "_pickle", /* m_name */
7157 pickle_module_doc, /* m_doc */
7158 sizeof(PickleState), /* m_size */
7159 pickle_methods, /* m_methods */
7160 NULL, /* m_reload */
7161 pickle_traverse, /* m_traverse */
7162 pickle_clear, /* m_clear */
7163 (freefunc)pickle_free /* m_free */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007164};
7165
7166PyMODINIT_FUNC
7167PyInit__pickle(void)
7168{
7169 PyObject *m;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007170 PickleState *st;
7171
7172 m = PyState_FindModule(&_picklemodule);
7173 if (m) {
7174 Py_INCREF(m);
7175 return m;
7176 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007177
7178 if (PyType_Ready(&Unpickler_Type) < 0)
7179 return NULL;
7180 if (PyType_Ready(&Pickler_Type) < 0)
7181 return NULL;
7182 if (PyType_Ready(&Pdata_Type) < 0)
7183 return NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00007184 if (PyType_Ready(&PicklerMemoProxyType) < 0)
7185 return NULL;
7186 if (PyType_Ready(&UnpicklerMemoProxyType) < 0)
7187 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007188
7189 /* Create the module and add the functions. */
7190 m = PyModule_Create(&_picklemodule);
7191 if (m == NULL)
7192 return NULL;
7193
Antoine Pitrou8391cf42011-07-15 21:01:21 +02007194 Py_INCREF(&Pickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007195 if (PyModule_AddObject(m, "Pickler", (PyObject *)&Pickler_Type) < 0)
7196 return NULL;
Antoine Pitrou8391cf42011-07-15 21:01:21 +02007197 Py_INCREF(&Unpickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007198 if (PyModule_AddObject(m, "Unpickler", (PyObject *)&Unpickler_Type) < 0)
7199 return NULL;
7200
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007201 st = _Pickle_GetState(m);
7202
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007203 /* Initialize the exceptions. */
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007204 st->PickleError = PyErr_NewException("_pickle.PickleError", NULL, NULL);
7205 if (st->PickleError == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007206 return NULL;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007207 st->PicklingError = \
7208 PyErr_NewException("_pickle.PicklingError", st->PickleError, NULL);
7209 if (st->PicklingError == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007210 return NULL;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007211 st->UnpicklingError = \
7212 PyErr_NewException("_pickle.UnpicklingError", st->PickleError, NULL);
7213 if (st->UnpicklingError == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007214 return NULL;
7215
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007216 Py_INCREF(st->PickleError);
7217 if (PyModule_AddObject(m, "PickleError", st->PickleError) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007218 return NULL;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007219 Py_INCREF(st->PicklingError);
7220 if (PyModule_AddObject(m, "PicklingError", st->PicklingError) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007221 return NULL;
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007222 Py_INCREF(st->UnpicklingError);
7223 if (PyModule_AddObject(m, "UnpicklingError", st->UnpicklingError) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007224 return NULL;
7225
Alexandre Vassalotti23bdd832013-11-27 19:36:52 -08007226 if (_Pickle_InitState(st) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00007227 return NULL;
7228
7229 return m;
7230}