blob: ae801f74effeb60e78d181855f44855bd0fcbb7f [file] [log] [blame]
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001#include "Python.h"
2#include "structmember.h"
3
4PyDoc_STRVAR(pickle_module_doc,
5"Optimized C implementation for the Python pickle module.");
6
7/* Bump this when new opcodes are added to the pickle protocol. */
8enum {
9 HIGHEST_PROTOCOL = 3,
10 DEFAULT_PROTOCOL = 3
11};
12
Alexandre Vassalottica2d6102008-06-12 18:26:05 +000013/* Pickle opcodes. These must be kept updated with pickle.py.
14 Extensive docs are in pickletools.py. */
15enum opcode {
16 MARK = '(',
17 STOP = '.',
18 POP = '0',
19 POP_MARK = '1',
20 DUP = '2',
21 FLOAT = 'F',
22 INT = 'I',
23 BININT = 'J',
24 BININT1 = 'K',
25 LONG = 'L',
26 BININT2 = 'M',
27 NONE = 'N',
28 PERSID = 'P',
29 BINPERSID = 'Q',
30 REDUCE = 'R',
31 STRING = 'S',
32 BINSTRING = 'T',
33 SHORT_BINSTRING = 'U',
34 UNICODE = 'V',
35 BINUNICODE = 'X',
36 APPEND = 'a',
37 BUILD = 'b',
38 GLOBAL = 'c',
39 DICT = 'd',
40 EMPTY_DICT = '}',
41 APPENDS = 'e',
42 GET = 'g',
43 BINGET = 'h',
44 INST = 'i',
45 LONG_BINGET = 'j',
46 LIST = 'l',
47 EMPTY_LIST = ']',
48 OBJ = 'o',
49 PUT = 'p',
50 BINPUT = 'q',
51 LONG_BINPUT = 'r',
52 SETITEM = 's',
53 TUPLE = 't',
54 EMPTY_TUPLE = ')',
55 SETITEMS = 'u',
56 BINFLOAT = 'G',
57
58 /* Protocol 2. */
59 PROTO = '\x80',
60 NEWOBJ = '\x81',
61 EXT1 = '\x82',
62 EXT2 = '\x83',
63 EXT4 = '\x84',
64 TUPLE1 = '\x85',
65 TUPLE2 = '\x86',
66 TUPLE3 = '\x87',
67 NEWTRUE = '\x88',
68 NEWFALSE = '\x89',
69 LONG1 = '\x8a',
70 LONG4 = '\x8b',
71
72 /* Protocol 3 (Python 3.x) */
73 BINBYTES = 'B',
Victor Stinner132ef6c2010-11-09 09:39:41 +000074 SHORT_BINBYTES = 'C'
Alexandre Vassalottica2d6102008-06-12 18:26:05 +000075};
76
77/* These aren't opcodes -- they're ways to pickle bools before protocol 2
78 * so that unpicklers written before bools were introduced unpickle them
79 * as ints, but unpicklers after can recognize that bools were intended.
80 * Note that protocol 2 added direct ways to pickle bools.
81 */
82#undef TRUE
83#define TRUE "I01\n"
84#undef FALSE
85#define FALSE "I00\n"
86
87enum {
88 /* Keep in synch with pickle.Pickler._BATCHSIZE. This is how many elements
89 batch_list/dict() pumps out before doing APPENDS/SETITEMS. Nothing will
90 break if this gets out of synch with pickle.py, but it's unclear that would
91 help anything either. */
92 BATCHSIZE = 1000,
93
94 /* Nesting limit until Pickler, when running in "fast mode", starts
95 checking for self-referential data-structures. */
96 FAST_NESTING_LIMIT = 50,
97
Antoine Pitrouea99c5c2010-09-09 18:33:21 +000098 /* Initial size of the write buffer of Pickler. */
99 WRITE_BUF_SIZE = 4096,
100
101 /* Maximum size of the write buffer of Pickler when pickling to a
102 stream. This is ignored for in-memory pickling. */
103 MAX_WRITE_BUF_SIZE = 64 * 1024,
Antoine Pitrou04248a82010-10-12 20:51:21 +0000104
105 /* Prefetch size when unpickling (disabled on unpeekable streams) */
Victor Stinner132ef6c2010-11-09 09:39:41 +0000106 PREFETCH = 8192 * 16
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000107};
108
109/* Exception classes for pickle. These should override the ones defined in
110 pickle.py, when the C-optimized Pickler and Unpickler are used. */
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000111static PyObject *PickleError = NULL;
112static PyObject *PicklingError = NULL;
113static PyObject *UnpicklingError = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000114
115/* copyreg.dispatch_table, {type_object: pickling_function} */
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000116static PyObject *dispatch_table = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000117/* For EXT[124] opcodes. */
118/* copyreg._extension_registry, {(module_name, function_name): code} */
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000119static PyObject *extension_registry = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000120/* copyreg._inverted_registry, {code: (module_name, function_name)} */
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000121static PyObject *inverted_registry = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000122/* copyreg._extension_cache, {code: object} */
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000123static PyObject *extension_cache = NULL;
124
125/* _compat_pickle.NAME_MAPPING, {(oldmodule, oldname): (newmodule, newname)} */
126static PyObject *name_mapping_2to3 = NULL;
127/* _compat_pickle.IMPORT_MAPPING, {oldmodule: newmodule} */
128static PyObject *import_mapping_2to3 = NULL;
129/* Same, but with REVERSE_NAME_MAPPING / REVERSE_IMPORT_MAPPING */
130static PyObject *name_mapping_3to2 = NULL;
131static PyObject *import_mapping_3to2 = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000132
133/* XXX: Are these really nescessary? */
134/* As the name says, an empty tuple. */
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000135static PyObject *empty_tuple = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000136/* For looking up name pairs in copyreg._extension_registry. */
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000137static PyObject *two_tuple = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000138
139static int
140stack_underflow(void)
141{
142 PyErr_SetString(UnpicklingError, "unpickling stack underflow");
143 return -1;
144}
145
146/* Internal data type used as the unpickling stack. */
147typedef struct {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000148 PyObject_VAR_HEAD
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000149 PyObject **data;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000150 Py_ssize_t allocated; /* number of slots in data allocated */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000151} Pdata;
152
153static void
154Pdata_dealloc(Pdata *self)
155{
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200156 Py_ssize_t i = Py_SIZE(self);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000157 while (--i >= 0) {
158 Py_DECREF(self->data[i]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000159 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000160 PyMem_FREE(self->data);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000161 PyObject_Del(self);
162}
163
164static PyTypeObject Pdata_Type = {
165 PyVarObject_HEAD_INIT(NULL, 0)
166 "_pickle.Pdata", /*tp_name*/
167 sizeof(Pdata), /*tp_basicsize*/
168 0, /*tp_itemsize*/
169 (destructor)Pdata_dealloc, /*tp_dealloc*/
170};
171
172static PyObject *
173Pdata_New(void)
174{
175 Pdata *self;
176
177 if (!(self = PyObject_New(Pdata, &Pdata_Type)))
178 return NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000179 Py_SIZE(self) = 0;
180 self->allocated = 8;
181 self->data = PyMem_MALLOC(self->allocated * sizeof(PyObject *));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000182 if (self->data)
183 return (PyObject *)self;
184 Py_DECREF(self);
185 return PyErr_NoMemory();
186}
187
188
189/* Retain only the initial clearto items. If clearto >= the current
190 * number of items, this is a (non-erroneous) NOP.
191 */
192static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200193Pdata_clear(Pdata *self, Py_ssize_t clearto)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000194{
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200195 Py_ssize_t i = Py_SIZE(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000196
197 if (clearto < 0)
198 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000199 if (clearto >= i)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000200 return 0;
201
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000202 while (--i >= clearto) {
203 Py_CLEAR(self->data[i]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000204 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000205 Py_SIZE(self) = clearto;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000206 return 0;
207}
208
209static int
210Pdata_grow(Pdata *self)
211{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000212 PyObject **data = self->data;
213 Py_ssize_t allocated = self->allocated;
214 Py_ssize_t new_allocated;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000215
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000216 new_allocated = (allocated >> 3) + 6;
217 /* check for integer overflow */
218 if (new_allocated > PY_SSIZE_T_MAX - allocated)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000219 goto nomemory;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000220 new_allocated += allocated;
221 if (new_allocated > (PY_SSIZE_T_MAX / sizeof(PyObject *)))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000222 goto nomemory;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000223 data = PyMem_REALLOC(data, new_allocated * sizeof(PyObject *));
224 if (data == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000225 goto nomemory;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000226
227 self->data = data;
228 self->allocated = new_allocated;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000229 return 0;
230
231 nomemory:
232 PyErr_NoMemory();
233 return -1;
234}
235
236/* D is a Pdata*. Pop the topmost element and store it into V, which
237 * must be an lvalue holding PyObject*. On stack underflow, UnpicklingError
238 * is raised and V is set to NULL.
239 */
240static PyObject *
241Pdata_pop(Pdata *self)
242{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000243 if (Py_SIZE(self) == 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000244 PyErr_SetString(UnpicklingError, "bad pickle data");
245 return NULL;
246 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000247 return self->data[--Py_SIZE(self)];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000248}
249#define PDATA_POP(D, V) do { (V) = Pdata_pop((D)); } while (0)
250
251static int
252Pdata_push(Pdata *self, PyObject *obj)
253{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000254 if (Py_SIZE(self) == self->allocated && Pdata_grow(self) < 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000255 return -1;
256 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000257 self->data[Py_SIZE(self)++] = obj;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000258 return 0;
259}
260
261/* Push an object on stack, transferring its ownership to the stack. */
262#define PDATA_PUSH(D, O, ER) do { \
263 if (Pdata_push((D), (O)) < 0) return (ER); } while(0)
264
265/* Push an object on stack, adding a new reference to the object. */
266#define PDATA_APPEND(D, O, ER) do { \
267 Py_INCREF((O)); \
268 if (Pdata_push((D), (O)) < 0) return (ER); } while(0)
269
270static PyObject *
271Pdata_poptuple(Pdata *self, Py_ssize_t start)
272{
273 PyObject *tuple;
274 Py_ssize_t len, i, j;
275
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000276 len = Py_SIZE(self) - start;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000277 tuple = PyTuple_New(len);
278 if (tuple == NULL)
279 return NULL;
280 for (i = start, j = 0; j < len; i++, j++)
281 PyTuple_SET_ITEM(tuple, j, self->data[i]);
282
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000283 Py_SIZE(self) = start;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000284 return tuple;
285}
286
287static PyObject *
288Pdata_poplist(Pdata *self, Py_ssize_t start)
289{
290 PyObject *list;
291 Py_ssize_t len, i, j;
292
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000293 len = Py_SIZE(self) - start;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000294 list = PyList_New(len);
295 if (list == NULL)
296 return NULL;
297 for (i = start, j = 0; j < len; i++, j++)
298 PyList_SET_ITEM(list, j, self->data[i]);
299
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000300 Py_SIZE(self) = start;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000301 return list;
302}
303
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000304typedef struct {
305 PyObject *me_key;
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200306 Py_ssize_t me_value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000307} PyMemoEntry;
308
309typedef struct {
310 Py_ssize_t mt_mask;
311 Py_ssize_t mt_used;
312 Py_ssize_t mt_allocated;
313 PyMemoEntry *mt_table;
314} PyMemoTable;
315
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000316typedef struct PicklerObject {
317 PyObject_HEAD
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000318 PyMemoTable *memo; /* Memo table, keep track of the seen
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000319 objects to support self-referential objects
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000320 pickling. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000321 PyObject *pers_func; /* persistent_id() method, can be NULL */
Antoine Pitrou8d3c2902012-03-04 18:31:48 +0100322 PyObject *dispatch_table; /* private dispatch_table, can be NULL */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000323 PyObject *arg;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000324
325 PyObject *write; /* write() method of the output stream. */
326 PyObject *output_buffer; /* Write into a local bytearray buffer before
327 flushing to the stream. */
328 Py_ssize_t output_len; /* Length of output_buffer. */
329 Py_ssize_t max_output_len; /* Allocation size of output_buffer. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000330 int proto; /* Pickle protocol number, >= 0 */
331 int bin; /* Boolean, true if proto > 0 */
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200332 Py_ssize_t buf_size; /* Size of the current buffered pickle data */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000333 int fast; /* Enable fast mode if set to a true value.
334 The fast mode disable the usage of memo,
335 therefore speeding the pickling process by
336 not generating superfluous PUT opcodes. It
337 should not be used if with self-referential
338 objects. */
339 int fast_nesting;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000340 int fix_imports; /* Indicate whether Pickler should fix
341 the name of globals for Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000342 PyObject *fast_memo;
343} PicklerObject;
344
345typedef struct UnpicklerObject {
346 PyObject_HEAD
347 Pdata *stack; /* Pickle data stack, store unpickled objects. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000348
349 /* The unpickler memo is just an array of PyObject *s. Using a dict
350 is unnecessary, since the keys are contiguous ints. */
351 PyObject **memo;
352 Py_ssize_t memo_size;
353
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000354 PyObject *arg;
355 PyObject *pers_func; /* persistent_load() method, can be NULL. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000356
357 Py_buffer buffer;
358 char *input_buffer;
359 char *input_line;
360 Py_ssize_t input_len;
361 Py_ssize_t next_read_idx;
Antoine Pitrou04248a82010-10-12 20:51:21 +0000362 Py_ssize_t prefetched_idx; /* index of first prefetched byte */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000363 PyObject *read; /* read() method of the input stream. */
364 PyObject *readline; /* readline() method of the input stream. */
Antoine Pitrou04248a82010-10-12 20:51:21 +0000365 PyObject *peek; /* peek() method of the input stream, or NULL */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000366
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000367 char *encoding; /* Name of the encoding to be used for
368 decoding strings pickled using Python
369 2.x. The default value is "ASCII" */
370 char *errors; /* Name of errors handling scheme to used when
371 decoding strings. The default value is
372 "strict". */
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -0500373 Py_ssize_t *marks; /* Mark stack, used for unpickling container
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000374 objects. */
375 Py_ssize_t num_marks; /* Number of marks in the mark stack. */
376 Py_ssize_t marks_size; /* Current allocated size of the mark stack. */
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000377 int proto; /* Protocol of the pickle loaded. */
378 int fix_imports; /* Indicate whether Unpickler should fix
379 the name of globals pickled by Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000380} UnpicklerObject;
381
382/* Forward declarations */
383static int save(PicklerObject *, PyObject *, int);
384static int save_reduce(PicklerObject *, PyObject *, PyObject *);
385static PyTypeObject Pickler_Type;
386static PyTypeObject Unpickler_Type;
387
388
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000389/*************************************************************************
Serhiy Storchaka95949422013-08-27 19:40:23 +0300390 A custom hashtable mapping void* to Python ints. This is used by the pickler
391 for memoization. Using a custom hashtable rather than PyDict allows us to skip
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000392 a bunch of unnecessary object creation. This makes a huge performance
393 difference. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000394
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000395#define MT_MINSIZE 8
396#define PERTURB_SHIFT 5
397
398
399static PyMemoTable *
400PyMemoTable_New(void)
401{
402 PyMemoTable *memo = PyMem_MALLOC(sizeof(PyMemoTable));
403 if (memo == NULL) {
404 PyErr_NoMemory();
405 return NULL;
406 }
407
408 memo->mt_used = 0;
409 memo->mt_allocated = MT_MINSIZE;
410 memo->mt_mask = MT_MINSIZE - 1;
411 memo->mt_table = PyMem_MALLOC(MT_MINSIZE * sizeof(PyMemoEntry));
412 if (memo->mt_table == NULL) {
413 PyMem_FREE(memo);
414 PyErr_NoMemory();
415 return NULL;
416 }
417 memset(memo->mt_table, 0, MT_MINSIZE * sizeof(PyMemoEntry));
418
419 return memo;
420}
421
422static PyMemoTable *
423PyMemoTable_Copy(PyMemoTable *self)
424{
425 Py_ssize_t i;
426 PyMemoTable *new = PyMemoTable_New();
427 if (new == NULL)
428 return NULL;
429
430 new->mt_used = self->mt_used;
431 new->mt_allocated = self->mt_allocated;
432 new->mt_mask = self->mt_mask;
433 /* The table we get from _New() is probably smaller than we wanted.
434 Free it and allocate one that's the right size. */
435 PyMem_FREE(new->mt_table);
436 new->mt_table = PyMem_MALLOC(self->mt_allocated * sizeof(PyMemoEntry));
437 if (new->mt_table == NULL) {
438 PyMem_FREE(new);
439 return NULL;
440 }
441 for (i = 0; i < self->mt_allocated; i++) {
442 Py_XINCREF(self->mt_table[i].me_key);
443 }
444 memcpy(new->mt_table, self->mt_table,
445 sizeof(PyMemoEntry) * self->mt_allocated);
446
447 return new;
448}
449
450static Py_ssize_t
451PyMemoTable_Size(PyMemoTable *self)
452{
453 return self->mt_used;
454}
455
456static int
457PyMemoTable_Clear(PyMemoTable *self)
458{
459 Py_ssize_t i = self->mt_allocated;
460
461 while (--i >= 0) {
462 Py_XDECREF(self->mt_table[i].me_key);
463 }
464 self->mt_used = 0;
465 memset(self->mt_table, 0, self->mt_allocated * sizeof(PyMemoEntry));
466 return 0;
467}
468
469static void
470PyMemoTable_Del(PyMemoTable *self)
471{
472 if (self == NULL)
473 return;
474 PyMemoTable_Clear(self);
475
476 PyMem_FREE(self->mt_table);
477 PyMem_FREE(self);
478}
479
480/* Since entries cannot be deleted from this hashtable, _PyMemoTable_Lookup()
481 can be considerably simpler than dictobject.c's lookdict(). */
482static PyMemoEntry *
483_PyMemoTable_Lookup(PyMemoTable *self, PyObject *key)
484{
485 size_t i;
486 size_t perturb;
487 size_t mask = (size_t)self->mt_mask;
488 PyMemoEntry *table = self->mt_table;
489 PyMemoEntry *entry;
Benjamin Peterson8f67d082010-10-17 20:54:53 +0000490 Py_hash_t hash = (Py_hash_t)key >> 3;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000491
492 i = hash & mask;
493 entry = &table[i];
494 if (entry->me_key == NULL || entry->me_key == key)
495 return entry;
496
497 for (perturb = hash; ; perturb >>= PERTURB_SHIFT) {
498 i = (i << 2) + i + perturb + 1;
499 entry = &table[i & mask];
500 if (entry->me_key == NULL || entry->me_key == key)
501 return entry;
502 }
503 assert(0); /* Never reached */
504 return NULL;
505}
506
507/* Returns -1 on failure, 0 on success. */
508static int
509_PyMemoTable_ResizeTable(PyMemoTable *self, Py_ssize_t min_size)
510{
511 PyMemoEntry *oldtable = NULL;
512 PyMemoEntry *oldentry, *newentry;
513 Py_ssize_t new_size = MT_MINSIZE;
514 Py_ssize_t to_process;
515
516 assert(min_size > 0);
517
518 /* Find the smallest valid table size >= min_size. */
519 while (new_size < min_size && new_size > 0)
520 new_size <<= 1;
521 if (new_size <= 0) {
522 PyErr_NoMemory();
523 return -1;
524 }
525 /* new_size needs to be a power of two. */
526 assert((new_size & (new_size - 1)) == 0);
527
528 /* Allocate new table. */
529 oldtable = self->mt_table;
530 self->mt_table = PyMem_MALLOC(new_size * sizeof(PyMemoEntry));
531 if (self->mt_table == NULL) {
532 PyMem_FREE(oldtable);
533 PyErr_NoMemory();
534 return -1;
535 }
536 self->mt_allocated = new_size;
537 self->mt_mask = new_size - 1;
538 memset(self->mt_table, 0, sizeof(PyMemoEntry) * new_size);
539
540 /* Copy entries from the old table. */
541 to_process = self->mt_used;
542 for (oldentry = oldtable; to_process > 0; oldentry++) {
543 if (oldentry->me_key != NULL) {
544 to_process--;
545 /* newentry is a pointer to a chunk of the new
546 mt_table, so we're setting the key:value pair
547 in-place. */
548 newentry = _PyMemoTable_Lookup(self, oldentry->me_key);
549 newentry->me_key = oldentry->me_key;
550 newentry->me_value = oldentry->me_value;
551 }
552 }
553
554 /* Deallocate the old table. */
555 PyMem_FREE(oldtable);
556 return 0;
557}
558
559/* Returns NULL on failure, a pointer to the value otherwise. */
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200560static Py_ssize_t *
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000561PyMemoTable_Get(PyMemoTable *self, PyObject *key)
562{
563 PyMemoEntry *entry = _PyMemoTable_Lookup(self, key);
564 if (entry->me_key == NULL)
565 return NULL;
566 return &entry->me_value;
567}
568
569/* Returns -1 on failure, 0 on success. */
570static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200571PyMemoTable_Set(PyMemoTable *self, PyObject *key, Py_ssize_t value)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000572{
573 PyMemoEntry *entry;
574
575 assert(key != NULL);
576
577 entry = _PyMemoTable_Lookup(self, key);
578 if (entry->me_key != NULL) {
579 entry->me_value = value;
580 return 0;
581 }
582 Py_INCREF(key);
583 entry->me_key = key;
584 entry->me_value = value;
585 self->mt_used++;
586
587 /* If we added a key, we can safely resize. Otherwise just return!
588 * If used >= 2/3 size, adjust size. Normally, this quaduples the size.
589 *
590 * Quadrupling the size improves average table sparseness
591 * (reducing collisions) at the cost of some memory. It also halves
592 * the number of expensive resize operations in a growing memo table.
593 *
594 * Very large memo tables (over 50K items) use doubling instead.
595 * This may help applications with severe memory constraints.
596 */
597 if (!(self->mt_used * 3 >= (self->mt_mask + 1) * 2))
598 return 0;
599 return _PyMemoTable_ResizeTable(self,
600 (self->mt_used > 50000 ? 2 : 4) * self->mt_used);
601}
602
603#undef MT_MINSIZE
604#undef PERTURB_SHIFT
605
606/*************************************************************************/
607
608/* Helpers for creating the argument tuple passed to functions. This has the
Victor Stinner121aab42011-09-29 23:40:53 +0200609 performance advantage of calling PyTuple_New() only once.
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000610
611 XXX(avassalotti): Inline directly in _Pickler_FastCall() and
612 _Unpickler_FastCall(). */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000613#define ARG_TUP(self, obj) do { \
614 if ((self)->arg || ((self)->arg=PyTuple_New(1))) { \
615 Py_XDECREF(PyTuple_GET_ITEM((self)->arg, 0)); \
616 PyTuple_SET_ITEM((self)->arg, 0, (obj)); \
617 } \
618 else { \
619 Py_DECREF((obj)); \
620 } \
621 } while (0)
622
623#define FREE_ARG_TUP(self) do { \
624 if ((self)->arg->ob_refcnt > 1) \
625 Py_CLEAR((self)->arg); \
626 } while (0)
627
628/* A temporary cleaner API for fast single argument function call.
629
630 XXX: Does caching the argument tuple provides any real performance benefits?
631
632 A quick benchmark, on a 2.0GHz Athlon64 3200+ running Linux 2.6.24 with
633 glibc 2.7, tells me that it takes roughly 20,000,000 PyTuple_New(1) calls
634 when the tuple is retrieved from the freelist (i.e, call PyTuple_New() then
635 immediately DECREF it) and 1,200,000 calls when allocating brand new tuples
636 (i.e, call PyTuple_New() and store the returned value in an array), to save
637 one second (wall clock time). Either ways, the loading time a pickle stream
638 large enough to generate this number of calls would be massively
639 overwhelmed by other factors, like I/O throughput, the GC traversal and
640 object allocation overhead. So, I really doubt these functions provide any
641 real benefits.
642
643 On the other hand, oprofile reports that pickle spends a lot of time in
644 these functions. But, that is probably more related to the function call
645 overhead, than the argument tuple allocation.
646
647 XXX: And, what is the reference behavior of these? Steal, borrow? At first
648 glance, it seems to steal the reference of 'arg' and borrow the reference
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000649 of 'func'. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000650static PyObject *
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000651_Pickler_FastCall(PicklerObject *self, PyObject *func, PyObject *arg)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000652{
653 PyObject *result = NULL;
654
655 ARG_TUP(self, arg);
656 if (self->arg) {
657 result = PyObject_Call(func, self->arg, NULL);
658 FREE_ARG_TUP(self);
659 }
660 return result;
661}
662
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000663static int
664_Pickler_ClearBuffer(PicklerObject *self)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000665{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000666 Py_CLEAR(self->output_buffer);
667 self->output_buffer =
668 PyBytes_FromStringAndSize(NULL, self->max_output_len);
669 if (self->output_buffer == NULL)
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +0000670 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000671 self->output_len = 0;
672 return 0;
673}
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +0000674
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000675static PyObject *
676_Pickler_GetString(PicklerObject *self)
677{
678 PyObject *output_buffer = self->output_buffer;
679
680 assert(self->output_buffer != NULL);
681 self->output_buffer = NULL;
682 /* Resize down to exact size */
683 if (_PyBytes_Resize(&output_buffer, self->output_len) < 0)
684 return NULL;
685 return output_buffer;
686}
687
688static int
689_Pickler_FlushToFile(PicklerObject *self)
690{
691 PyObject *output, *result;
692
693 assert(self->write != NULL);
694
695 output = _Pickler_GetString(self);
696 if (output == NULL)
697 return -1;
698
699 result = _Pickler_FastCall(self, self->write, output);
700 Py_XDECREF(result);
701 return (result == NULL) ? -1 : 0;
702}
703
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200704static Py_ssize_t
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000705_Pickler_Write(PicklerObject *self, const char *s, Py_ssize_t n)
706{
707 Py_ssize_t i, required;
708 char *buffer;
709
710 assert(s != NULL);
711
712 required = self->output_len + n;
713 if (required > self->max_output_len) {
714 if (self->write != NULL && required > MAX_WRITE_BUF_SIZE) {
715 /* XXX This reallocates a new buffer every time, which is a bit
716 wasteful. */
717 if (_Pickler_FlushToFile(self) < 0)
718 return -1;
719 if (_Pickler_ClearBuffer(self) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000720 return -1;
721 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000722 if (self->write != NULL && n > MAX_WRITE_BUF_SIZE) {
723 /* we already flushed above, so the buffer is empty */
724 PyObject *result;
725 /* XXX we could spare an intermediate copy and pass
726 a memoryview instead */
727 PyObject *output = PyBytes_FromStringAndSize(s, n);
728 if (s == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000729 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000730 result = _Pickler_FastCall(self, self->write, output);
731 Py_XDECREF(result);
732 return (result == NULL) ? -1 : 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000733 }
734 else {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000735 if (self->output_len >= PY_SSIZE_T_MAX / 2 - n) {
736 PyErr_NoMemory();
737 return -1;
738 }
Antoine Pitrou82be19f2011-08-29 23:09:33 +0200739 self->max_output_len = (self->output_len + n) / 2 * 3;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000740 if (_PyBytes_Resize(&self->output_buffer, self->max_output_len) < 0)
741 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000742 }
743 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000744 buffer = PyBytes_AS_STRING(self->output_buffer);
745 if (n < 8) {
746 /* This is faster than memcpy when the string is short. */
747 for (i = 0; i < n; i++) {
748 buffer[self->output_len + i] = s[i];
749 }
750 }
751 else {
752 memcpy(buffer + self->output_len, s, n);
753 }
754 self->output_len += n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000755 return n;
756}
757
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000758static PicklerObject *
759_Pickler_New(void)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000760{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000761 PicklerObject *self;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000762
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000763 self = PyObject_GC_New(PicklerObject, &Pickler_Type);
764 if (self == NULL)
765 return NULL;
766
767 self->pers_func = NULL;
Antoine Pitrou8d3c2902012-03-04 18:31:48 +0100768 self->dispatch_table = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000769 self->arg = NULL;
770 self->write = NULL;
771 self->proto = 0;
772 self->bin = 0;
773 self->fast = 0;
774 self->fast_nesting = 0;
775 self->fix_imports = 0;
776 self->fast_memo = NULL;
777
778 self->memo = PyMemoTable_New();
779 if (self->memo == NULL) {
780 Py_DECREF(self);
781 return NULL;
782 }
783 self->max_output_len = WRITE_BUF_SIZE;
784 self->output_len = 0;
785 self->output_buffer = PyBytes_FromStringAndSize(NULL,
786 self->max_output_len);
787 if (self->output_buffer == NULL) {
788 Py_DECREF(self);
789 return NULL;
790 }
791 return self;
792}
793
794static int
795_Pickler_SetProtocol(PicklerObject *self, PyObject *proto_obj,
796 PyObject *fix_imports_obj)
797{
798 long proto = 0;
799 int fix_imports;
800
801 if (proto_obj == NULL || proto_obj == Py_None)
802 proto = DEFAULT_PROTOCOL;
803 else {
804 proto = PyLong_AsLong(proto_obj);
805 if (proto == -1 && PyErr_Occurred())
806 return -1;
807 }
808 if (proto < 0)
809 proto = HIGHEST_PROTOCOL;
810 if (proto > HIGHEST_PROTOCOL) {
811 PyErr_Format(PyExc_ValueError, "pickle protocol must be <= %d",
812 HIGHEST_PROTOCOL);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000813 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000814 }
815 fix_imports = PyObject_IsTrue(fix_imports_obj);
816 if (fix_imports == -1)
817 return -1;
Victor Stinner121aab42011-09-29 23:40:53 +0200818
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000819 self->proto = proto;
820 self->bin = proto > 0;
821 self->fix_imports = fix_imports && proto < 3;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000822
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000823 return 0;
824}
825
826/* Returns -1 (with an exception set) on failure, 0 on success. This may
827 be called once on a freshly created Pickler. */
828static int
829_Pickler_SetOutputStream(PicklerObject *self, PyObject *file)
830{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +0200831 _Py_IDENTIFIER(write);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000832 assert(file != NULL);
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +0200833 self->write = _PyObject_GetAttrId(file, &PyId_write);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000834 if (self->write == NULL) {
835 if (PyErr_ExceptionMatches(PyExc_AttributeError))
836 PyErr_SetString(PyExc_TypeError,
837 "file must have a 'write' attribute");
838 return -1;
839 }
840
841 return 0;
842}
843
844/* See documentation for _Pickler_FastCall(). */
845static PyObject *
846_Unpickler_FastCall(UnpicklerObject *self, PyObject *func, PyObject *arg)
847{
848 PyObject *result = NULL;
849
850 ARG_TUP(self, arg);
851 if (self->arg) {
852 result = PyObject_Call(func, self->arg, NULL);
853 FREE_ARG_TUP(self);
854 }
855 return result;
856}
857
858/* Returns the size of the input on success, -1 on failure. This takes its
859 own reference to `input`. */
860static Py_ssize_t
861_Unpickler_SetStringInput(UnpicklerObject *self, PyObject *input)
862{
863 if (self->buffer.buf != NULL)
864 PyBuffer_Release(&self->buffer);
865 if (PyObject_GetBuffer(input, &self->buffer, PyBUF_CONTIG_RO) < 0)
866 return -1;
867 self->input_buffer = self->buffer.buf;
868 self->input_len = self->buffer.len;
869 self->next_read_idx = 0;
Antoine Pitrou04248a82010-10-12 20:51:21 +0000870 self->prefetched_idx = self->input_len;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000871 return self->input_len;
872}
873
Antoine Pitrou04248a82010-10-12 20:51:21 +0000874static int
875_Unpickler_SkipConsumed(UnpicklerObject *self)
876{
877 Py_ssize_t consumed = self->next_read_idx - self->prefetched_idx;
878
879 if (consumed > 0) {
880 PyObject *r;
881 assert(self->peek); /* otherwise we did something wrong */
882 /* This makes an useless copy... */
883 r = PyObject_CallFunction(self->read, "n", consumed);
884 if (r == NULL)
885 return -1;
886 Py_DECREF(r);
887 self->prefetched_idx = self->next_read_idx;
888 }
889 return 0;
890}
891
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000892static const Py_ssize_t READ_WHOLE_LINE = -1;
893
894/* If reading from a file, we need to only pull the bytes we need, since there
895 may be multiple pickle objects arranged contiguously in the same input
896 buffer.
897
898 If `n` is READ_WHOLE_LINE, read a whole line. Otherwise, read up to `n`
899 bytes from the input stream/buffer.
900
901 Update the unpickler's input buffer with the newly-read data. Returns -1 on
902 failure; on success, returns the number of bytes read from the file.
903
904 On success, self->input_len will be 0; this is intentional so that when
905 unpickling from a file, the "we've run out of data" code paths will trigger,
906 causing the Unpickler to go back to the file for more data. Use the returned
907 size to tell you how much data you can process. */
908static Py_ssize_t
909_Unpickler_ReadFromFile(UnpicklerObject *self, Py_ssize_t n)
910{
911 PyObject *data;
Antoine Pitrou04248a82010-10-12 20:51:21 +0000912 Py_ssize_t read_size, prefetched_size = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000913
914 assert(self->read != NULL);
Victor Stinner121aab42011-09-29 23:40:53 +0200915
Antoine Pitrou04248a82010-10-12 20:51:21 +0000916 if (_Unpickler_SkipConsumed(self) < 0)
917 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000918
919 if (n == READ_WHOLE_LINE)
920 data = PyObject_Call(self->readline, empty_tuple, NULL);
921 else {
922 PyObject *len = PyLong_FromSsize_t(n);
923 if (len == NULL)
924 return -1;
925 data = _Unpickler_FastCall(self, self->read, len);
926 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000927 if (data == NULL)
928 return -1;
929
Antoine Pitrou04248a82010-10-12 20:51:21 +0000930 /* Prefetch some data without advancing the file pointer, if possible */
931 if (self->peek) {
932 PyObject *len, *prefetched;
933 len = PyLong_FromSsize_t(PREFETCH);
934 if (len == NULL) {
935 Py_DECREF(data);
936 return -1;
937 }
938 prefetched = _Unpickler_FastCall(self, self->peek, len);
939 if (prefetched == NULL) {
940 if (PyErr_ExceptionMatches(PyExc_NotImplementedError)) {
941 /* peek() is probably not supported by the given file object */
942 PyErr_Clear();
943 Py_CLEAR(self->peek);
944 }
945 else {
946 Py_DECREF(data);
947 return -1;
948 }
949 }
950 else {
951 assert(PyBytes_Check(prefetched));
952 prefetched_size = PyBytes_GET_SIZE(prefetched);
953 PyBytes_ConcatAndDel(&data, prefetched);
954 if (data == NULL)
955 return -1;
956 }
957 }
958
959 read_size = _Unpickler_SetStringInput(self, data) - prefetched_size;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000960 Py_DECREF(data);
Antoine Pitrou04248a82010-10-12 20:51:21 +0000961 self->prefetched_idx = read_size;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000962 return read_size;
963}
964
965/* Read `n` bytes from the unpickler's data source, storing the result in `*s`.
966
967 This should be used for all data reads, rather than accessing the unpickler's
968 input buffer directly. This method deals correctly with reading from input
969 streams, which the input buffer doesn't deal with.
970
971 Note that when reading from a file-like object, self->next_read_idx won't
972 be updated (it should remain at 0 for the entire unpickling process). You
973 should use this function's return value to know how many bytes you can
974 consume.
975
976 Returns -1 (with an exception set) on failure. On success, return the
977 number of chars read. */
978static Py_ssize_t
979_Unpickler_Read(UnpicklerObject *self, char **s, Py_ssize_t n)
980{
Antoine Pitrou04248a82010-10-12 20:51:21 +0000981 Py_ssize_t num_read;
982
Antoine Pitrou04248a82010-10-12 20:51:21 +0000983 if (self->next_read_idx + n <= self->input_len) {
984 *s = self->input_buffer + self->next_read_idx;
985 self->next_read_idx += n;
986 return n;
987 }
988 if (!self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000989 PyErr_Format(PyExc_EOFError, "Ran out of input");
Amaury Forgeot d'Arc3e4e72f2008-11-11 20:05:06 +0000990 return -1;
991 }
Antoine Pitrou04248a82010-10-12 20:51:21 +0000992 num_read = _Unpickler_ReadFromFile(self, n);
993 if (num_read < 0)
994 return -1;
995 if (num_read < n) {
996 PyErr_Format(PyExc_EOFError, "Ran out of input");
997 return -1;
998 }
999 *s = self->input_buffer;
1000 self->next_read_idx = n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001001 return n;
1002}
1003
1004static Py_ssize_t
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001005_Unpickler_CopyLine(UnpicklerObject *self, char *line, Py_ssize_t len,
1006 char **result)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001007{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001008 char *input_line = PyMem_Realloc(self->input_line, len + 1);
1009 if (input_line == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001010 return -1;
1011
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001012 memcpy(input_line, line, len);
1013 input_line[len] = '\0';
1014 self->input_line = input_line;
1015 *result = self->input_line;
1016 return len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001017}
1018
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001019/* Read a line from the input stream/buffer. If we run off the end of the input
1020 before hitting \n, return the data we found.
1021
1022 Returns the number of chars read, or -1 on failure. */
1023static Py_ssize_t
1024_Unpickler_Readline(UnpicklerObject *self, char **result)
1025{
1026 Py_ssize_t i, num_read;
1027
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001028 for (i = self->next_read_idx; i < self->input_len; i++) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001029 if (self->input_buffer[i] == '\n') {
1030 char *line_start = self->input_buffer + self->next_read_idx;
1031 num_read = i - self->next_read_idx + 1;
1032 self->next_read_idx = i + 1;
1033 return _Unpickler_CopyLine(self, line_start, num_read, result);
1034 }
1035 }
1036 if (self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001037 num_read = _Unpickler_ReadFromFile(self, READ_WHOLE_LINE);
1038 if (num_read < 0)
1039 return -1;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001040 self->next_read_idx = num_read;
Antoine Pitrouf6c7a852011-08-11 21:04:02 +02001041 return _Unpickler_CopyLine(self, self->input_buffer, num_read, result);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001042 }
Victor Stinner121aab42011-09-29 23:40:53 +02001043
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001044 /* If we get here, we've run off the end of the input string. Return the
1045 remaining string and let the caller figure it out. */
1046 *result = self->input_buffer + self->next_read_idx;
1047 num_read = i - self->next_read_idx;
1048 self->next_read_idx = i;
1049 return num_read;
1050}
1051
1052/* Returns -1 (with an exception set) on failure, 0 on success. The memo array
1053 will be modified in place. */
1054static int
1055_Unpickler_ResizeMemoList(UnpicklerObject *self, Py_ssize_t new_size)
1056{
1057 Py_ssize_t i;
1058 PyObject **memo;
1059
1060 assert(new_size > self->memo_size);
1061
1062 memo = PyMem_REALLOC(self->memo, new_size * sizeof(PyObject *));
1063 if (memo == NULL) {
1064 PyErr_NoMemory();
1065 return -1;
1066 }
1067 self->memo = memo;
1068 for (i = self->memo_size; i < new_size; i++)
1069 self->memo[i] = NULL;
1070 self->memo_size = new_size;
1071 return 0;
1072}
1073
1074/* Returns NULL if idx is out of bounds. */
1075static PyObject *
1076_Unpickler_MemoGet(UnpicklerObject *self, Py_ssize_t idx)
1077{
1078 if (idx < 0 || idx >= self->memo_size)
1079 return NULL;
1080
1081 return self->memo[idx];
1082}
1083
1084/* Returns -1 (with an exception set) on failure, 0 on success.
1085 This takes its own reference to `value`. */
1086static int
1087_Unpickler_MemoPut(UnpicklerObject *self, Py_ssize_t idx, PyObject *value)
1088{
1089 PyObject *old_item;
1090
1091 if (idx >= self->memo_size) {
1092 if (_Unpickler_ResizeMemoList(self, idx * 2) < 0)
1093 return -1;
1094 assert(idx < self->memo_size);
1095 }
1096 Py_INCREF(value);
1097 old_item = self->memo[idx];
1098 self->memo[idx] = value;
1099 Py_XDECREF(old_item);
1100 return 0;
1101}
1102
1103static PyObject **
1104_Unpickler_NewMemo(Py_ssize_t new_size)
1105{
1106 PyObject **memo = PyMem_MALLOC(new_size * sizeof(PyObject *));
1107 if (memo == NULL)
1108 return NULL;
1109 memset(memo, 0, new_size * sizeof(PyObject *));
1110 return memo;
1111}
1112
1113/* Free the unpickler's memo, taking care to decref any items left in it. */
1114static void
1115_Unpickler_MemoCleanup(UnpicklerObject *self)
1116{
1117 Py_ssize_t i;
1118 PyObject **memo = self->memo;
1119
1120 if (self->memo == NULL)
1121 return;
1122 self->memo = NULL;
1123 i = self->memo_size;
1124 while (--i >= 0) {
1125 Py_XDECREF(memo[i]);
1126 }
1127 PyMem_FREE(memo);
1128}
1129
1130static UnpicklerObject *
1131_Unpickler_New(void)
1132{
1133 UnpicklerObject *self;
1134
1135 self = PyObject_GC_New(UnpicklerObject, &Unpickler_Type);
1136 if (self == NULL)
1137 return NULL;
1138
1139 self->stack = (Pdata *)Pdata_New();
1140 if (self->stack == NULL) {
1141 Py_DECREF(self);
1142 return NULL;
1143 }
1144 memset(&self->buffer, 0, sizeof(Py_buffer));
1145
1146 self->memo_size = 32;
1147 self->memo = _Unpickler_NewMemo(self->memo_size);
1148 if (self->memo == NULL) {
1149 Py_DECREF(self);
1150 return NULL;
1151 }
1152
1153 self->arg = NULL;
1154 self->pers_func = NULL;
1155 self->input_buffer = NULL;
1156 self->input_line = NULL;
1157 self->input_len = 0;
1158 self->next_read_idx = 0;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001159 self->prefetched_idx = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001160 self->read = NULL;
1161 self->readline = NULL;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001162 self->peek = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001163 self->encoding = NULL;
1164 self->errors = NULL;
1165 self->marks = NULL;
1166 self->num_marks = 0;
1167 self->marks_size = 0;
1168 self->proto = 0;
1169 self->fix_imports = 0;
1170
1171 return self;
1172}
1173
1174/* Returns -1 (with an exception set) on failure, 0 on success. This may
1175 be called once on a freshly created Pickler. */
1176static int
1177_Unpickler_SetInputStream(UnpicklerObject *self, PyObject *file)
1178{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02001179 _Py_IDENTIFIER(peek);
1180 _Py_IDENTIFIER(read);
1181 _Py_IDENTIFIER(readline);
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02001182
1183 self->peek = _PyObject_GetAttrId(file, &PyId_peek);
Antoine Pitrou04248a82010-10-12 20:51:21 +00001184 if (self->peek == NULL) {
1185 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1186 PyErr_Clear();
1187 else
1188 return -1;
1189 }
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02001190 self->read = _PyObject_GetAttrId(file, &PyId_read);
1191 self->readline = _PyObject_GetAttrId(file, &PyId_readline);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001192 if (self->readline == NULL || self->read == NULL) {
1193 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1194 PyErr_SetString(PyExc_TypeError,
1195 "file must have 'read' and 'readline' attributes");
1196 Py_CLEAR(self->read);
1197 Py_CLEAR(self->readline);
Antoine Pitrou04248a82010-10-12 20:51:21 +00001198 Py_CLEAR(self->peek);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001199 return -1;
1200 }
1201 return 0;
1202}
1203
1204/* Returns -1 (with an exception set) on failure, 0 on success. This may
1205 be called once on a freshly created Pickler. */
1206static int
1207_Unpickler_SetInputEncoding(UnpicklerObject *self,
1208 const char *encoding,
1209 const char *errors)
1210{
1211 if (encoding == NULL)
1212 encoding = "ASCII";
1213 if (errors == NULL)
1214 errors = "strict";
1215
1216 self->encoding = strdup(encoding);
1217 self->errors = strdup(errors);
1218 if (self->encoding == NULL || self->errors == NULL) {
1219 PyErr_NoMemory();
1220 return -1;
1221 }
1222 return 0;
1223}
1224
1225/* Generate a GET opcode for an object stored in the memo. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001226static int
1227memo_get(PicklerObject *self, PyObject *key)
1228{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001229 Py_ssize_t *value;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001230 char pdata[30];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001231 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001232
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001233 value = PyMemoTable_Get(self->memo, key);
1234 if (value == NULL) {
1235 PyErr_SetObject(PyExc_KeyError, key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001236 return -1;
1237 }
1238
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001239 if (!self->bin) {
1240 pdata[0] = GET;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001241 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1,
1242 "%" PY_FORMAT_SIZE_T "d\n", *value);
1243 len = strlen(pdata);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001244 }
1245 else {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001246 if (*value < 256) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001247 pdata[0] = BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001248 pdata[1] = (unsigned char)(*value & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001249 len = 2;
1250 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001251 else if (*value <= 0xffffffffL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001252 pdata[0] = LONG_BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001253 pdata[1] = (unsigned char)(*value & 0xff);
1254 pdata[2] = (unsigned char)((*value >> 8) & 0xff);
1255 pdata[3] = (unsigned char)((*value >> 16) & 0xff);
1256 pdata[4] = (unsigned char)((*value >> 24) & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001257 len = 5;
1258 }
1259 else { /* unlikely */
1260 PyErr_SetString(PicklingError,
1261 "memo id too large for LONG_BINGET");
1262 return -1;
1263 }
1264 }
1265
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001266 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001267 return -1;
1268
1269 return 0;
1270}
1271
1272/* Store an object in the memo, assign it a new unique ID based on the number
1273 of objects currently stored in the memo and generate a PUT opcode. */
1274static int
1275memo_put(PicklerObject *self, PyObject *obj)
1276{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001277 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001278 char pdata[30];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001279 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001280 int status = 0;
1281
1282 if (self->fast)
1283 return 0;
1284
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001285 x = PyMemoTable_Size(self->memo);
1286 if (PyMemoTable_Set(self->memo, obj, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001287 goto error;
1288
1289 if (!self->bin) {
1290 pdata[0] = PUT;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001291 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1,
1292 "%" PY_FORMAT_SIZE_T "d\n", x);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001293 len = strlen(pdata);
1294 }
1295 else {
1296 if (x < 256) {
1297 pdata[0] = BINPUT;
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00001298 pdata[1] = (unsigned char)x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001299 len = 2;
1300 }
1301 else if (x <= 0xffffffffL) {
1302 pdata[0] = LONG_BINPUT;
1303 pdata[1] = (unsigned char)(x & 0xff);
1304 pdata[2] = (unsigned char)((x >> 8) & 0xff);
1305 pdata[3] = (unsigned char)((x >> 16) & 0xff);
1306 pdata[4] = (unsigned char)((x >> 24) & 0xff);
1307 len = 5;
1308 }
1309 else { /* unlikely */
1310 PyErr_SetString(PicklingError,
1311 "memo id too large for LONG_BINPUT");
1312 return -1;
1313 }
1314 }
1315
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001316 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001317 goto error;
1318
1319 if (0) {
1320 error:
1321 status = -1;
1322 }
1323
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001324 return status;
1325}
1326
1327static PyObject *
1328whichmodule(PyObject *global, PyObject *global_name)
1329{
1330 Py_ssize_t i, j;
1331 static PyObject *module_str = NULL;
1332 static PyObject *main_str = NULL;
1333 PyObject *module_name;
1334 PyObject *modules_dict;
1335 PyObject *module;
1336 PyObject *obj;
1337
1338 if (module_str == NULL) {
1339 module_str = PyUnicode_InternFromString("__module__");
1340 if (module_str == NULL)
1341 return NULL;
1342 main_str = PyUnicode_InternFromString("__main__");
1343 if (main_str == NULL)
1344 return NULL;
1345 }
1346
1347 module_name = PyObject_GetAttr(global, module_str);
1348
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00001349 /* In some rare cases (e.g., bound methods of extension types),
1350 __module__ can be None. If it is so, then search sys.modules
1351 for the module of global. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001352 if (module_name == Py_None) {
1353 Py_DECREF(module_name);
1354 goto search;
1355 }
1356
1357 if (module_name) {
1358 return module_name;
1359 }
1360 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1361 PyErr_Clear();
1362 else
1363 return NULL;
1364
1365 search:
1366 modules_dict = PySys_GetObject("modules");
1367 if (modules_dict == NULL)
1368 return NULL;
1369
1370 i = 0;
1371 module_name = NULL;
1372 while ((j = PyDict_Next(modules_dict, &i, &module_name, &module))) {
Mark Dickinson211c6252009-02-01 10:28:51 +00001373 if (PyObject_RichCompareBool(module_name, main_str, Py_EQ) == 1)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001374 continue;
1375
1376 obj = PyObject_GetAttr(module, global_name);
1377 if (obj == NULL) {
1378 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1379 PyErr_Clear();
1380 else
1381 return NULL;
1382 continue;
1383 }
1384
1385 if (obj != global) {
1386 Py_DECREF(obj);
1387 continue;
1388 }
1389
1390 Py_DECREF(obj);
1391 break;
1392 }
1393
1394 /* If no module is found, use __main__. */
1395 if (!j) {
1396 module_name = main_str;
1397 }
1398
1399 Py_INCREF(module_name);
1400 return module_name;
1401}
1402
1403/* fast_save_enter() and fast_save_leave() are guards against recursive
1404 objects when Pickler is used with the "fast mode" (i.e., with object
1405 memoization disabled). If the nesting of a list or dict object exceed
1406 FAST_NESTING_LIMIT, these guards will start keeping an internal
1407 reference to the seen list or dict objects and check whether these objects
1408 are recursive. These are not strictly necessary, since save() has a
1409 hard-coded recursion limit, but they give a nicer error message than the
1410 typical RuntimeError. */
1411static int
1412fast_save_enter(PicklerObject *self, PyObject *obj)
1413{
1414 /* if fast_nesting < 0, we're doing an error exit. */
1415 if (++self->fast_nesting >= FAST_NESTING_LIMIT) {
1416 PyObject *key = NULL;
1417 if (self->fast_memo == NULL) {
1418 self->fast_memo = PyDict_New();
1419 if (self->fast_memo == NULL) {
1420 self->fast_nesting = -1;
1421 return 0;
1422 }
1423 }
1424 key = PyLong_FromVoidPtr(obj);
1425 if (key == NULL)
1426 return 0;
1427 if (PyDict_GetItem(self->fast_memo, key)) {
1428 Py_DECREF(key);
1429 PyErr_Format(PyExc_ValueError,
1430 "fast mode: can't pickle cyclic objects "
1431 "including object type %.200s at %p",
1432 obj->ob_type->tp_name, obj);
1433 self->fast_nesting = -1;
1434 return 0;
1435 }
1436 if (PyDict_SetItem(self->fast_memo, key, Py_None) < 0) {
1437 Py_DECREF(key);
1438 self->fast_nesting = -1;
1439 return 0;
1440 }
1441 Py_DECREF(key);
1442 }
1443 return 1;
1444}
1445
1446static int
1447fast_save_leave(PicklerObject *self, PyObject *obj)
1448{
1449 if (self->fast_nesting-- >= FAST_NESTING_LIMIT) {
1450 PyObject *key = PyLong_FromVoidPtr(obj);
1451 if (key == NULL)
1452 return 0;
1453 if (PyDict_DelItem(self->fast_memo, key) < 0) {
1454 Py_DECREF(key);
1455 return 0;
1456 }
1457 Py_DECREF(key);
1458 }
1459 return 1;
1460}
1461
1462static int
1463save_none(PicklerObject *self, PyObject *obj)
1464{
1465 const char none_op = NONE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001466 if (_Pickler_Write(self, &none_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001467 return -1;
1468
1469 return 0;
1470}
1471
1472static int
1473save_bool(PicklerObject *self, PyObject *obj)
1474{
1475 static const char *buf[2] = { FALSE, TRUE };
1476 const char len[2] = {sizeof(FALSE) - 1, sizeof(TRUE) - 1};
1477 int p = (obj == Py_True);
1478
1479 if (self->proto >= 2) {
1480 const char bool_op = p ? NEWTRUE : NEWFALSE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001481 if (_Pickler_Write(self, &bool_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001482 return -1;
1483 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001484 else if (_Pickler_Write(self, buf[p], len[p]) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001485 return -1;
1486
1487 return 0;
1488}
1489
1490static int
1491save_int(PicklerObject *self, long x)
1492{
1493 char pdata[32];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001494 Py_ssize_t len = 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001495
1496 if (!self->bin
1497#if SIZEOF_LONG > 4
1498 || x > 0x7fffffffL || x < -0x80000000L
1499#endif
1500 ) {
1501 /* Text-mode pickle, or long too big to fit in the 4-byte
1502 * signed BININT format: store as a string.
1503 */
Mark Dickinson8dd05142009-01-20 20:43:58 +00001504 pdata[0] = LONG; /* use LONG for consistency with pickle.py */
1505 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1, "%ldL\n", x);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001506 if (_Pickler_Write(self, pdata, strlen(pdata)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001507 return -1;
1508 }
1509 else {
1510 /* Binary pickle and x fits in a signed 4-byte int. */
1511 pdata[1] = (unsigned char)(x & 0xff);
1512 pdata[2] = (unsigned char)((x >> 8) & 0xff);
1513 pdata[3] = (unsigned char)((x >> 16) & 0xff);
1514 pdata[4] = (unsigned char)((x >> 24) & 0xff);
1515
1516 if ((pdata[4] == 0) && (pdata[3] == 0)) {
1517 if (pdata[2] == 0) {
1518 pdata[0] = BININT1;
1519 len = 2;
1520 }
1521 else {
1522 pdata[0] = BININT2;
1523 len = 3;
1524 }
1525 }
1526 else {
1527 pdata[0] = BININT;
1528 len = 5;
1529 }
1530
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001531 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001532 return -1;
1533 }
1534
1535 return 0;
1536}
1537
1538static int
1539save_long(PicklerObject *self, PyObject *obj)
1540{
1541 PyObject *repr = NULL;
1542 Py_ssize_t size;
1543 long val = PyLong_AsLong(obj);
1544 int status = 0;
1545
1546 const char long_op = LONG;
1547
1548 if (val == -1 && PyErr_Occurred()) {
1549 /* out of range for int pickling */
1550 PyErr_Clear();
1551 }
Antoine Pitroue58bffb2011-08-13 20:40:32 +02001552 else
1553#if SIZEOF_LONG > 4
1554 if (val <= 0x7fffffffL && val >= -0x80000000L)
1555#endif
1556 return save_int(self, val);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001557
1558 if (self->proto >= 2) {
1559 /* Linear-time pickling. */
1560 size_t nbits;
1561 size_t nbytes;
1562 unsigned char *pdata;
1563 char header[5];
1564 int i;
1565 int sign = _PyLong_Sign(obj);
1566
1567 if (sign == 0) {
1568 header[0] = LONG1;
1569 header[1] = 0; /* It's 0 -- an empty bytestring. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001570 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001571 goto error;
1572 return 0;
1573 }
1574 nbits = _PyLong_NumBits(obj);
1575 if (nbits == (size_t)-1 && PyErr_Occurred())
1576 goto error;
1577 /* How many bytes do we need? There are nbits >> 3 full
1578 * bytes of data, and nbits & 7 leftover bits. If there
1579 * are any leftover bits, then we clearly need another
1580 * byte. Wnat's not so obvious is that we *probably*
1581 * need another byte even if there aren't any leftovers:
1582 * the most-significant bit of the most-significant byte
1583 * acts like a sign bit, and it's usually got a sense
Serhiy Storchaka95949422013-08-27 19:40:23 +03001584 * opposite of the one we need. The exception is ints
1585 * of the form -(2**(8*j-1)) for j > 0. Such an int is
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001586 * its own 256's-complement, so has the right sign bit
1587 * even without the extra byte. That's a pain to check
1588 * for in advance, though, so we always grab an extra
1589 * byte at the start, and cut it back later if possible.
1590 */
1591 nbytes = (nbits >> 3) + 1;
Antoine Pitroubf6ecf92012-11-24 20:40:21 +01001592 if (nbytes > 0x7fffffffL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001593 PyErr_SetString(PyExc_OverflowError,
Serhiy Storchaka95949422013-08-27 19:40:23 +03001594 "int too large to pickle");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001595 goto error;
1596 }
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001597 repr = PyBytes_FromStringAndSize(NULL, (Py_ssize_t)nbytes);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001598 if (repr == NULL)
1599 goto error;
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001600 pdata = (unsigned char *)PyBytes_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001601 i = _PyLong_AsByteArray((PyLongObject *)obj,
1602 pdata, nbytes,
1603 1 /* little endian */ , 1 /* signed */ );
1604 if (i < 0)
1605 goto error;
Serhiy Storchaka95949422013-08-27 19:40:23 +03001606 /* If the int is negative, this may be a byte more than
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001607 * needed. This is so iff the MSB is all redundant sign
1608 * bits.
1609 */
1610 if (sign < 0 &&
Victor Stinner121aab42011-09-29 23:40:53 +02001611 nbytes > 1 &&
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001612 pdata[nbytes - 1] == 0xff &&
1613 (pdata[nbytes - 2] & 0x80) != 0) {
1614 nbytes--;
1615 }
1616
1617 if (nbytes < 256) {
1618 header[0] = LONG1;
1619 header[1] = (unsigned char)nbytes;
1620 size = 2;
1621 }
1622 else {
1623 header[0] = LONG4;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001624 size = (Py_ssize_t) nbytes;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001625 for (i = 1; i < 5; i++) {
1626 header[i] = (unsigned char)(size & 0xff);
1627 size >>= 8;
1628 }
1629 size = 5;
1630 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001631 if (_Pickler_Write(self, header, size) < 0 ||
1632 _Pickler_Write(self, (char *)pdata, (int)nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001633 goto error;
1634 }
1635 else {
1636 char *string;
1637
Mark Dickinson8dd05142009-01-20 20:43:58 +00001638 /* proto < 2: write the repr and newline. This is quadratic-time (in
1639 the number of digits), in both directions. We add a trailing 'L'
1640 to the repr, for compatibility with Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001641
1642 repr = PyObject_Repr(obj);
1643 if (repr == NULL)
1644 goto error;
1645
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00001646 string = _PyUnicode_AsStringAndSize(repr, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001647 if (string == NULL)
1648 goto error;
1649
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001650 if (_Pickler_Write(self, &long_op, 1) < 0 ||
1651 _Pickler_Write(self, string, size) < 0 ||
1652 _Pickler_Write(self, "L\n", 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001653 goto error;
1654 }
1655
1656 if (0) {
1657 error:
1658 status = -1;
1659 }
1660 Py_XDECREF(repr);
1661
1662 return status;
1663}
1664
1665static int
1666save_float(PicklerObject *self, PyObject *obj)
1667{
1668 double x = PyFloat_AS_DOUBLE((PyFloatObject *)obj);
1669
1670 if (self->bin) {
1671 char pdata[9];
1672 pdata[0] = BINFLOAT;
1673 if (_PyFloat_Pack8(x, (unsigned char *)&pdata[1], 0) < 0)
1674 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001675 if (_Pickler_Write(self, pdata, 9) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001676 return -1;
Victor Stinner121aab42011-09-29 23:40:53 +02001677 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001678 else {
Eric Smith0923d1d2009-04-16 20:16:10 +00001679 int result = -1;
1680 char *buf = NULL;
1681 char op = FLOAT;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001682
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001683 if (_Pickler_Write(self, &op, 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001684 goto done;
1685
Mark Dickinson3e09f432009-04-17 08:41:23 +00001686 buf = PyOS_double_to_string(x, 'g', 17, 0, NULL);
Eric Smith0923d1d2009-04-16 20:16:10 +00001687 if (!buf) {
1688 PyErr_NoMemory();
1689 goto done;
1690 }
1691
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001692 if (_Pickler_Write(self, buf, strlen(buf)) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001693 goto done;
1694
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001695 if (_Pickler_Write(self, "\n", 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001696 goto done;
1697
1698 result = 0;
1699done:
1700 PyMem_Free(buf);
1701 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001702 }
1703
1704 return 0;
1705}
1706
1707static int
1708save_bytes(PicklerObject *self, PyObject *obj)
1709{
1710 if (self->proto < 3) {
1711 /* Older pickle protocols do not have an opcode for pickling bytes
1712 objects. Therefore, we need to fake the copy protocol (i.e.,
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001713 the __reduce__ method) to permit bytes object unpickling.
1714
1715 Here we use a hack to be compatible with Python 2. Since in Python
1716 2 'bytes' is just an alias for 'str' (which has different
1717 parameters than the actual bytes object), we use codecs.encode
1718 to create the appropriate 'str' object when unpickled using
1719 Python 2 *and* the appropriate 'bytes' object when unpickled
1720 using Python 3. Again this is a hack and we don't need to do this
1721 with newer protocols. */
1722 static PyObject *codecs_encode = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001723 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001724 int status;
1725
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001726 if (codecs_encode == NULL) {
1727 PyObject *codecs_module = PyImport_ImportModule("codecs");
1728 if (codecs_module == NULL) {
1729 return -1;
1730 }
1731 codecs_encode = PyObject_GetAttrString(codecs_module, "encode");
1732 Py_DECREF(codecs_module);
1733 if (codecs_encode == NULL) {
1734 return -1;
1735 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001736 }
1737
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001738 if (PyBytes_GET_SIZE(obj) == 0) {
1739 reduce_value = Py_BuildValue("(O())", (PyObject*)&PyBytes_Type);
1740 }
1741 else {
1742 static PyObject *latin1 = NULL;
1743 PyObject *unicode_str =
1744 PyUnicode_DecodeLatin1(PyBytes_AS_STRING(obj),
1745 PyBytes_GET_SIZE(obj),
1746 "strict");
1747 if (unicode_str == NULL)
1748 return -1;
1749 if (latin1 == NULL) {
1750 latin1 = PyUnicode_InternFromString("latin1");
1751 if (latin1 == NULL)
1752 return -1;
1753 }
1754 reduce_value = Py_BuildValue("(O(OO))",
1755 codecs_encode, unicode_str, latin1);
1756 Py_DECREF(unicode_str);
1757 }
1758
1759 if (reduce_value == NULL)
1760 return -1;
1761
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001762 /* save_reduce() will memoize the object automatically. */
1763 status = save_reduce(self, reduce_value, obj);
1764 Py_DECREF(reduce_value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001765 return status;
1766 }
1767 else {
1768 Py_ssize_t size;
1769 char header[5];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001770 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001771
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001772 size = PyBytes_GET_SIZE(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001773 if (size < 0)
1774 return -1;
1775
1776 if (size < 256) {
1777 header[0] = SHORT_BINBYTES;
1778 header[1] = (unsigned char)size;
1779 len = 2;
1780 }
1781 else if (size <= 0xffffffffL) {
1782 header[0] = BINBYTES;
1783 header[1] = (unsigned char)(size & 0xff);
1784 header[2] = (unsigned char)((size >> 8) & 0xff);
1785 header[3] = (unsigned char)((size >> 16) & 0xff);
1786 header[4] = (unsigned char)((size >> 24) & 0xff);
1787 len = 5;
1788 }
1789 else {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001790 PyErr_SetString(PyExc_OverflowError,
Serhiy Storchakaf8def282013-02-16 17:29:56 +02001791 "cannot serialize a bytes object larger than 4 GiB");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001792 return -1; /* string too large */
1793 }
1794
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001795 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001796 return -1;
1797
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001798 if (_Pickler_Write(self, PyBytes_AS_STRING(obj), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001799 return -1;
1800
1801 if (memo_put(self, obj) < 0)
1802 return -1;
1803
1804 return 0;
1805 }
1806}
1807
1808/* A copy of PyUnicode_EncodeRawUnicodeEscape() that also translates
1809 backslash and newline characters to \uXXXX escapes. */
1810static PyObject *
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001811raw_unicode_escape(PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001812{
1813 PyObject *repr, *result;
1814 char *p;
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001815 Py_ssize_t i, size, expandsize;
1816 void *data;
1817 unsigned int kind;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001818
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001819 if (PyUnicode_READY(obj))
1820 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001821
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001822 size = PyUnicode_GET_LENGTH(obj);
1823 data = PyUnicode_DATA(obj);
1824 kind = PyUnicode_KIND(obj);
1825 if (kind == PyUnicode_4BYTE_KIND)
1826 expandsize = 10;
1827 else
1828 expandsize = 6;
Victor Stinner121aab42011-09-29 23:40:53 +02001829
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001830 if (size > PY_SSIZE_T_MAX / expandsize)
1831 return PyErr_NoMemory();
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001832 repr = PyByteArray_FromStringAndSize(NULL, expandsize * size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001833 if (repr == NULL)
1834 return NULL;
1835 if (size == 0)
1836 goto done;
1837
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001838 p = PyByteArray_AS_STRING(repr);
1839 for (i=0; i < size; i++) {
1840 Py_UCS4 ch = PyUnicode_READ(kind, data, i);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001841 /* Map 32-bit characters to '\Uxxxxxxxx' */
1842 if (ch >= 0x10000) {
1843 *p++ = '\\';
1844 *p++ = 'U';
Victor Stinnerf5cff562011-10-14 02:13:11 +02001845 *p++ = Py_hexdigits[(ch >> 28) & 0xf];
1846 *p++ = Py_hexdigits[(ch >> 24) & 0xf];
1847 *p++ = Py_hexdigits[(ch >> 20) & 0xf];
1848 *p++ = Py_hexdigits[(ch >> 16) & 0xf];
1849 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
1850 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
1851 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
1852 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001853 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001854 /* Map 16-bit characters to '\uxxxx' */
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001855 else if (ch >= 256 || ch == '\\' || ch == '\n') {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001856 *p++ = '\\';
1857 *p++ = 'u';
Victor Stinnerf5cff562011-10-14 02:13:11 +02001858 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
1859 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
1860 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
1861 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001862 }
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001863 /* Copy everything else as-is */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001864 else
1865 *p++ = (char) ch;
1866 }
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001867 size = p - PyByteArray_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001868
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001869done:
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001870 result = PyBytes_FromStringAndSize(PyByteArray_AS_STRING(repr), size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001871 Py_DECREF(repr);
1872 return result;
1873}
1874
1875static int
1876save_unicode(PicklerObject *self, PyObject *obj)
1877{
1878 Py_ssize_t size;
1879 PyObject *encoded = NULL;
1880
1881 if (self->bin) {
1882 char pdata[5];
1883
Martin v. Löwisd63a3b82011-09-28 07:41:54 +02001884 encoded = PyUnicode_AsEncodedString(obj, "utf-8", "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001885 if (encoded == NULL)
1886 goto error;
1887
1888 size = PyBytes_GET_SIZE(encoded);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001889 if (size > 0xffffffffL) {
1890 PyErr_SetString(PyExc_OverflowError,
Serhiy Storchakaf8def282013-02-16 17:29:56 +02001891 "cannot serialize a string larger than 4 GiB");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001892 goto error; /* string too large */
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001893 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001894
1895 pdata[0] = BINUNICODE;
1896 pdata[1] = (unsigned char)(size & 0xff);
1897 pdata[2] = (unsigned char)((size >> 8) & 0xff);
1898 pdata[3] = (unsigned char)((size >> 16) & 0xff);
1899 pdata[4] = (unsigned char)((size >> 24) & 0xff);
1900
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001901 if (_Pickler_Write(self, pdata, 5) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001902 goto error;
1903
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001904 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001905 goto error;
1906 }
1907 else {
1908 const char unicode_op = UNICODE;
1909
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001910 encoded = raw_unicode_escape(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001911 if (encoded == NULL)
1912 goto error;
1913
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001914 if (_Pickler_Write(self, &unicode_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001915 goto error;
1916
1917 size = PyBytes_GET_SIZE(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001918 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001919 goto error;
1920
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001921 if (_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001922 goto error;
1923 }
1924 if (memo_put(self, obj) < 0)
1925 goto error;
1926
1927 Py_DECREF(encoded);
1928 return 0;
1929
1930 error:
1931 Py_XDECREF(encoded);
1932 return -1;
1933}
1934
1935/* A helper for save_tuple. Push the len elements in tuple t on the stack. */
1936static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001937store_tuple_elements(PicklerObject *self, PyObject *t, Py_ssize_t len)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001938{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001939 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001940
1941 assert(PyTuple_Size(t) == len);
1942
1943 for (i = 0; i < len; i++) {
1944 PyObject *element = PyTuple_GET_ITEM(t, i);
1945
1946 if (element == NULL)
1947 return -1;
1948 if (save(self, element, 0) < 0)
1949 return -1;
1950 }
1951
1952 return 0;
1953}
1954
1955/* Tuples are ubiquitous in the pickle protocols, so many techniques are
1956 * used across protocols to minimize the space needed to pickle them.
1957 * Tuples are also the only builtin immutable type that can be recursive
1958 * (a tuple can be reached from itself), and that requires some subtle
1959 * magic so that it works in all cases. IOW, this is a long routine.
1960 */
1961static int
1962save_tuple(PicklerObject *self, PyObject *obj)
1963{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001964 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001965
1966 const char mark_op = MARK;
1967 const char tuple_op = TUPLE;
1968 const char pop_op = POP;
1969 const char pop_mark_op = POP_MARK;
1970 const char len2opcode[] = {EMPTY_TUPLE, TUPLE1, TUPLE2, TUPLE3};
1971
1972 if ((len = PyTuple_Size(obj)) < 0)
1973 return -1;
1974
1975 if (len == 0) {
1976 char pdata[2];
1977
1978 if (self->proto) {
1979 pdata[0] = EMPTY_TUPLE;
1980 len = 1;
1981 }
1982 else {
1983 pdata[0] = MARK;
1984 pdata[1] = TUPLE;
1985 len = 2;
1986 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001987 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001988 return -1;
1989 return 0;
1990 }
1991
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001992 /* The tuple isn't in the memo now. If it shows up there after
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001993 * saving the tuple elements, the tuple must be recursive, in
1994 * which case we'll pop everything we put on the stack, and fetch
1995 * its value from the memo.
1996 */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001997 if (len <= 3 && self->proto >= 2) {
1998 /* Use TUPLE{1,2,3} opcodes. */
1999 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002000 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002001
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002002 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002003 /* pop the len elements */
2004 for (i = 0; i < len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002005 if (_Pickler_Write(self, &pop_op, 1) < 0)
2006 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002007 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002008 if (memo_get(self, obj) < 0)
2009 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002010
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002011 return 0;
2012 }
2013 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002014 if (_Pickler_Write(self, len2opcode + len, 1) < 0)
2015 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002016 }
2017 goto memoize;
2018 }
2019
2020 /* proto < 2 and len > 0, or proto >= 2 and len > 3.
2021 * Generate MARK e1 e2 ... TUPLE
2022 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002023 if (_Pickler_Write(self, &mark_op, 1) < 0)
2024 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002025
2026 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002027 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002028
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002029 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002030 /* pop the stack stuff we pushed */
2031 if (self->bin) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002032 if (_Pickler_Write(self, &pop_mark_op, 1) < 0)
2033 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002034 }
2035 else {
2036 /* Note that we pop one more than len, to remove
2037 * the MARK too.
2038 */
2039 for (i = 0; i <= len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002040 if (_Pickler_Write(self, &pop_op, 1) < 0)
2041 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002042 }
2043 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002044 if (memo_get(self, obj) < 0)
2045 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002046
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002047 return 0;
2048 }
2049 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002050 if (_Pickler_Write(self, &tuple_op, 1) < 0)
2051 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002052 }
2053
2054 memoize:
2055 if (memo_put(self, obj) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002056 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002057
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002058 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002059}
2060
2061/* iter is an iterator giving items, and we batch up chunks of
2062 * MARK item item ... item APPENDS
2063 * opcode sequences. Calling code should have arranged to first create an
2064 * empty list, or list-like object, for the APPENDS to operate on.
2065 * Returns 0 on success, <0 on error.
2066 */
2067static int
2068batch_list(PicklerObject *self, PyObject *iter)
2069{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002070 PyObject *obj = NULL;
2071 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002072 int i, n;
2073
2074 const char mark_op = MARK;
2075 const char append_op = APPEND;
2076 const char appends_op = APPENDS;
2077
2078 assert(iter != NULL);
2079
2080 /* XXX: I think this function could be made faster by avoiding the
2081 iterator interface and fetching objects directly from list using
2082 PyList_GET_ITEM.
2083 */
2084
2085 if (self->proto == 0) {
2086 /* APPENDS isn't available; do one at a time. */
2087 for (;;) {
2088 obj = PyIter_Next(iter);
2089 if (obj == NULL) {
2090 if (PyErr_Occurred())
2091 return -1;
2092 break;
2093 }
2094 i = save(self, obj, 0);
2095 Py_DECREF(obj);
2096 if (i < 0)
2097 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002098 if (_Pickler_Write(self, &append_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002099 return -1;
2100 }
2101 return 0;
2102 }
2103
2104 /* proto > 0: write in batches of BATCHSIZE. */
2105 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002106 /* Get first item */
2107 firstitem = PyIter_Next(iter);
2108 if (firstitem == NULL) {
2109 if (PyErr_Occurred())
2110 goto error;
2111
2112 /* nothing more to add */
2113 break;
2114 }
2115
2116 /* Try to get a second item */
2117 obj = PyIter_Next(iter);
2118 if (obj == NULL) {
2119 if (PyErr_Occurred())
2120 goto error;
2121
2122 /* Only one item to write */
2123 if (save(self, firstitem, 0) < 0)
2124 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002125 if (_Pickler_Write(self, &append_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002126 goto error;
2127 Py_CLEAR(firstitem);
2128 break;
2129 }
2130
2131 /* More than one item to write */
2132
2133 /* Pump out MARK, items, APPENDS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002134 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002135 goto error;
2136
2137 if (save(self, firstitem, 0) < 0)
2138 goto error;
2139 Py_CLEAR(firstitem);
2140 n = 1;
2141
2142 /* Fetch and save up to BATCHSIZE items */
2143 while (obj) {
2144 if (save(self, obj, 0) < 0)
2145 goto error;
2146 Py_CLEAR(obj);
2147 n += 1;
2148
2149 if (n == BATCHSIZE)
2150 break;
2151
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002152 obj = PyIter_Next(iter);
2153 if (obj == NULL) {
2154 if (PyErr_Occurred())
2155 goto error;
2156 break;
2157 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002158 }
2159
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002160 if (_Pickler_Write(self, &appends_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002161 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002162
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002163 } while (n == BATCHSIZE);
2164 return 0;
2165
2166 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002167 Py_XDECREF(firstitem);
2168 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002169 return -1;
2170}
2171
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002172/* This is a variant of batch_list() above, specialized for lists (with no
2173 * support for list subclasses). Like batch_list(), we batch up chunks of
2174 * MARK item item ... item APPENDS
2175 * opcode sequences. Calling code should have arranged to first create an
2176 * empty list, or list-like object, for the APPENDS to operate on.
2177 * Returns 0 on success, -1 on error.
2178 *
2179 * This version is considerably faster than batch_list(), if less general.
2180 *
2181 * Note that this only works for protocols > 0.
2182 */
2183static int
2184batch_list_exact(PicklerObject *self, PyObject *obj)
2185{
2186 PyObject *item = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002187 Py_ssize_t this_batch, total;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002188
2189 const char append_op = APPEND;
2190 const char appends_op = APPENDS;
2191 const char mark_op = MARK;
2192
2193 assert(obj != NULL);
2194 assert(self->proto > 0);
2195 assert(PyList_CheckExact(obj));
2196
2197 if (PyList_GET_SIZE(obj) == 1) {
2198 item = PyList_GET_ITEM(obj, 0);
2199 if (save(self, item, 0) < 0)
2200 return -1;
2201 if (_Pickler_Write(self, &append_op, 1) < 0)
2202 return -1;
2203 return 0;
2204 }
2205
2206 /* Write in batches of BATCHSIZE. */
2207 total = 0;
2208 do {
2209 this_batch = 0;
2210 if (_Pickler_Write(self, &mark_op, 1) < 0)
2211 return -1;
2212 while (total < PyList_GET_SIZE(obj)) {
2213 item = PyList_GET_ITEM(obj, total);
2214 if (save(self, item, 0) < 0)
2215 return -1;
2216 total++;
2217 if (++this_batch == BATCHSIZE)
2218 break;
2219 }
2220 if (_Pickler_Write(self, &appends_op, 1) < 0)
2221 return -1;
2222
2223 } while (total < PyList_GET_SIZE(obj));
2224
2225 return 0;
2226}
2227
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002228static int
2229save_list(PicklerObject *self, PyObject *obj)
2230{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002231 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002232 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002233 int status = 0;
2234
2235 if (self->fast && !fast_save_enter(self, obj))
2236 goto error;
2237
2238 /* Create an empty list. */
2239 if (self->bin) {
2240 header[0] = EMPTY_LIST;
2241 len = 1;
2242 }
2243 else {
2244 header[0] = MARK;
2245 header[1] = LIST;
2246 len = 2;
2247 }
2248
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002249 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002250 goto error;
2251
2252 /* Get list length, and bow out early if empty. */
2253 if ((len = PyList_Size(obj)) < 0)
2254 goto error;
2255
2256 if (memo_put(self, obj) < 0)
2257 goto error;
2258
2259 if (len != 0) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002260 /* Materialize the list elements. */
2261 if (PyList_CheckExact(obj) && self->proto > 0) {
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002262 if (Py_EnterRecursiveCall(" while pickling an object"))
2263 goto error;
2264 status = batch_list_exact(self, obj);
2265 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002266 } else {
2267 PyObject *iter = PyObject_GetIter(obj);
2268 if (iter == NULL)
2269 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002270
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002271 if (Py_EnterRecursiveCall(" while pickling an object")) {
2272 Py_DECREF(iter);
2273 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002274 }
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002275 status = batch_list(self, iter);
2276 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002277 Py_DECREF(iter);
2278 }
2279 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002280 if (0) {
2281 error:
2282 status = -1;
2283 }
2284
2285 if (self->fast && !fast_save_leave(self, obj))
2286 status = -1;
2287
2288 return status;
2289}
2290
2291/* iter is an iterator giving (key, value) pairs, and we batch up chunks of
2292 * MARK key value ... key value SETITEMS
2293 * opcode sequences. Calling code should have arranged to first create an
2294 * empty dict, or dict-like object, for the SETITEMS to operate on.
2295 * Returns 0 on success, <0 on error.
2296 *
2297 * This is very much like batch_list(). The difference between saving
2298 * elements directly, and picking apart two-tuples, is so long-winded at
2299 * the C level, though, that attempts to combine these routines were too
2300 * ugly to bear.
2301 */
2302static int
2303batch_dict(PicklerObject *self, PyObject *iter)
2304{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002305 PyObject *obj = NULL;
2306 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002307 int i, n;
2308
2309 const char mark_op = MARK;
2310 const char setitem_op = SETITEM;
2311 const char setitems_op = SETITEMS;
2312
2313 assert(iter != NULL);
2314
2315 if (self->proto == 0) {
2316 /* SETITEMS isn't available; do one at a time. */
2317 for (;;) {
2318 obj = PyIter_Next(iter);
2319 if (obj == NULL) {
2320 if (PyErr_Occurred())
2321 return -1;
2322 break;
2323 }
2324 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2325 PyErr_SetString(PyExc_TypeError, "dict items "
2326 "iterator must return 2-tuples");
2327 return -1;
2328 }
2329 i = save(self, PyTuple_GET_ITEM(obj, 0), 0);
2330 if (i >= 0)
2331 i = save(self, PyTuple_GET_ITEM(obj, 1), 0);
2332 Py_DECREF(obj);
2333 if (i < 0)
2334 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002335 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002336 return -1;
2337 }
2338 return 0;
2339 }
2340
2341 /* proto > 0: write in batches of BATCHSIZE. */
2342 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002343 /* Get first item */
2344 firstitem = PyIter_Next(iter);
2345 if (firstitem == NULL) {
2346 if (PyErr_Occurred())
2347 goto error;
2348
2349 /* nothing more to add */
2350 break;
2351 }
2352 if (!PyTuple_Check(firstitem) || PyTuple_Size(firstitem) != 2) {
2353 PyErr_SetString(PyExc_TypeError, "dict items "
2354 "iterator must return 2-tuples");
2355 goto error;
2356 }
2357
2358 /* Try to get a second item */
2359 obj = PyIter_Next(iter);
2360 if (obj == NULL) {
2361 if (PyErr_Occurred())
2362 goto error;
2363
2364 /* Only one item to write */
2365 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2366 goto error;
2367 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2368 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002369 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002370 goto error;
2371 Py_CLEAR(firstitem);
2372 break;
2373 }
2374
2375 /* More than one item to write */
2376
2377 /* Pump out MARK, items, SETITEMS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002378 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002379 goto error;
2380
2381 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2382 goto error;
2383 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2384 goto error;
2385 Py_CLEAR(firstitem);
2386 n = 1;
2387
2388 /* Fetch and save up to BATCHSIZE items */
2389 while (obj) {
2390 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2391 PyErr_SetString(PyExc_TypeError, "dict items "
2392 "iterator must return 2-tuples");
2393 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002394 }
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002395 if (save(self, PyTuple_GET_ITEM(obj, 0), 0) < 0 ||
2396 save(self, PyTuple_GET_ITEM(obj, 1), 0) < 0)
2397 goto error;
2398 Py_CLEAR(obj);
2399 n += 1;
2400
2401 if (n == BATCHSIZE)
2402 break;
2403
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002404 obj = PyIter_Next(iter);
2405 if (obj == NULL) {
2406 if (PyErr_Occurred())
2407 goto error;
2408 break;
2409 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002410 }
2411
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002412 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002413 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002414
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002415 } while (n == BATCHSIZE);
2416 return 0;
2417
2418 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002419 Py_XDECREF(firstitem);
2420 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002421 return -1;
2422}
2423
Collin Winter5c9b02d2009-05-25 05:43:30 +00002424/* This is a variant of batch_dict() above that specializes for dicts, with no
2425 * support for dict subclasses. Like batch_dict(), we batch up chunks of
2426 * MARK key value ... key value SETITEMS
2427 * opcode sequences. Calling code should have arranged to first create an
2428 * empty dict, or dict-like object, for the SETITEMS to operate on.
2429 * Returns 0 on success, -1 on error.
2430 *
2431 * Note that this currently doesn't work for protocol 0.
2432 */
2433static int
2434batch_dict_exact(PicklerObject *self, PyObject *obj)
2435{
2436 PyObject *key = NULL, *value = NULL;
2437 int i;
2438 Py_ssize_t dict_size, ppos = 0;
2439
Alexandre Vassalottif70b1292009-05-25 18:00:52 +00002440 const char mark_op = MARK;
2441 const char setitem_op = SETITEM;
2442 const char setitems_op = SETITEMS;
Collin Winter5c9b02d2009-05-25 05:43:30 +00002443
2444 assert(obj != NULL);
2445 assert(self->proto > 0);
2446
2447 dict_size = PyDict_Size(obj);
2448
2449 /* Special-case len(d) == 1 to save space. */
2450 if (dict_size == 1) {
2451 PyDict_Next(obj, &ppos, &key, &value);
2452 if (save(self, key, 0) < 0)
2453 return -1;
2454 if (save(self, value, 0) < 0)
2455 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002456 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002457 return -1;
2458 return 0;
2459 }
2460
2461 /* Write in batches of BATCHSIZE. */
2462 do {
2463 i = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002464 if (_Pickler_Write(self, &mark_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002465 return -1;
2466 while (PyDict_Next(obj, &ppos, &key, &value)) {
2467 if (save(self, key, 0) < 0)
2468 return -1;
2469 if (save(self, value, 0) < 0)
2470 return -1;
2471 if (++i == BATCHSIZE)
2472 break;
2473 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002474 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002475 return -1;
2476 if (PyDict_Size(obj) != dict_size) {
2477 PyErr_Format(
2478 PyExc_RuntimeError,
2479 "dictionary changed size during iteration");
2480 return -1;
2481 }
2482
2483 } while (i == BATCHSIZE);
2484 return 0;
2485}
2486
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002487static int
2488save_dict(PicklerObject *self, PyObject *obj)
2489{
2490 PyObject *items, *iter;
2491 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002492 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002493 int status = 0;
2494
2495 if (self->fast && !fast_save_enter(self, obj))
2496 goto error;
2497
2498 /* Create an empty dict. */
2499 if (self->bin) {
2500 header[0] = EMPTY_DICT;
2501 len = 1;
2502 }
2503 else {
2504 header[0] = MARK;
2505 header[1] = DICT;
2506 len = 2;
2507 }
2508
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002509 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002510 goto error;
2511
2512 /* Get dict size, and bow out early if empty. */
2513 if ((len = PyDict_Size(obj)) < 0)
2514 goto error;
2515
2516 if (memo_put(self, obj) < 0)
2517 goto error;
2518
2519 if (len != 0) {
2520 /* Save the dict items. */
Collin Winter5c9b02d2009-05-25 05:43:30 +00002521 if (PyDict_CheckExact(obj) && self->proto > 0) {
2522 /* We can take certain shortcuts if we know this is a dict and
2523 not a dict subclass. */
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002524 if (Py_EnterRecursiveCall(" while pickling an object"))
2525 goto error;
2526 status = batch_dict_exact(self, obj);
2527 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002528 } else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02002529 _Py_IDENTIFIER(items);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02002530
2531 items = _PyObject_CallMethodId(obj, &PyId_items, "()");
Collin Winter5c9b02d2009-05-25 05:43:30 +00002532 if (items == NULL)
2533 goto error;
2534 iter = PyObject_GetIter(items);
2535 Py_DECREF(items);
2536 if (iter == NULL)
2537 goto error;
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002538 if (Py_EnterRecursiveCall(" while pickling an object")) {
2539 Py_DECREF(iter);
2540 goto error;
2541 }
Collin Winter5c9b02d2009-05-25 05:43:30 +00002542 status = batch_dict(self, iter);
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002543 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002544 Py_DECREF(iter);
2545 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002546 }
2547
2548 if (0) {
2549 error:
2550 status = -1;
2551 }
2552
2553 if (self->fast && !fast_save_leave(self, obj))
2554 status = -1;
2555
2556 return status;
2557}
2558
2559static int
2560save_global(PicklerObject *self, PyObject *obj, PyObject *name)
2561{
2562 static PyObject *name_str = NULL;
2563 PyObject *global_name = NULL;
2564 PyObject *module_name = NULL;
2565 PyObject *module = NULL;
2566 PyObject *cls;
2567 int status = 0;
2568
2569 const char global_op = GLOBAL;
2570
2571 if (name_str == NULL) {
2572 name_str = PyUnicode_InternFromString("__name__");
2573 if (name_str == NULL)
2574 goto error;
2575 }
2576
2577 if (name) {
2578 global_name = name;
2579 Py_INCREF(global_name);
2580 }
2581 else {
2582 global_name = PyObject_GetAttr(obj, name_str);
2583 if (global_name == NULL)
2584 goto error;
2585 }
2586
2587 module_name = whichmodule(obj, global_name);
2588 if (module_name == NULL)
2589 goto error;
2590
2591 /* XXX: Change to use the import C API directly with level=0 to disallow
2592 relative imports.
2593
2594 XXX: PyImport_ImportModuleLevel could be used. However, this bypasses
2595 builtins.__import__. Therefore, _pickle, unlike pickle.py, will ignore
2596 custom import functions (IMHO, this would be a nice security
2597 feature). The import C API would need to be extended to support the
2598 extra parameters of __import__ to fix that. */
2599 module = PyImport_Import(module_name);
2600 if (module == NULL) {
2601 PyErr_Format(PicklingError,
2602 "Can't pickle %R: import of module %R failed",
2603 obj, module_name);
2604 goto error;
2605 }
2606 cls = PyObject_GetAttr(module, global_name);
2607 if (cls == NULL) {
2608 PyErr_Format(PicklingError,
2609 "Can't pickle %R: attribute lookup %S.%S failed",
2610 obj, module_name, global_name);
2611 goto error;
2612 }
2613 if (cls != obj) {
2614 Py_DECREF(cls);
2615 PyErr_Format(PicklingError,
2616 "Can't pickle %R: it's not the same object as %S.%S",
2617 obj, module_name, global_name);
2618 goto error;
2619 }
2620 Py_DECREF(cls);
2621
2622 if (self->proto >= 2) {
2623 /* See whether this is in the extension registry, and if
2624 * so generate an EXT opcode.
2625 */
2626 PyObject *code_obj; /* extension code as Python object */
2627 long code; /* extension code as C value */
2628 char pdata[5];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002629 Py_ssize_t n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002630
2631 PyTuple_SET_ITEM(two_tuple, 0, module_name);
2632 PyTuple_SET_ITEM(two_tuple, 1, global_name);
2633 code_obj = PyDict_GetItem(extension_registry, two_tuple);
2634 /* The object is not registered in the extension registry.
2635 This is the most likely code path. */
2636 if (code_obj == NULL)
2637 goto gen_global;
2638
2639 /* XXX: pickle.py doesn't check neither the type, nor the range
2640 of the value returned by the extension_registry. It should for
2641 consistency. */
2642
2643 /* Verify code_obj has the right type and value. */
2644 if (!PyLong_Check(code_obj)) {
2645 PyErr_Format(PicklingError,
2646 "Can't pickle %R: extension code %R isn't an integer",
2647 obj, code_obj);
2648 goto error;
2649 }
2650 code = PyLong_AS_LONG(code_obj);
2651 if (code <= 0 || code > 0x7fffffffL) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002652 if (!PyErr_Occurred())
2653 PyErr_Format(PicklingError,
2654 "Can't pickle %R: extension code %ld is out of range",
2655 obj, code);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002656 goto error;
2657 }
2658
2659 /* Generate an EXT opcode. */
2660 if (code <= 0xff) {
2661 pdata[0] = EXT1;
2662 pdata[1] = (unsigned char)code;
2663 n = 2;
2664 }
2665 else if (code <= 0xffff) {
2666 pdata[0] = EXT2;
2667 pdata[1] = (unsigned char)(code & 0xff);
2668 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2669 n = 3;
2670 }
2671 else {
2672 pdata[0] = EXT4;
2673 pdata[1] = (unsigned char)(code & 0xff);
2674 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2675 pdata[3] = (unsigned char)((code >> 16) & 0xff);
2676 pdata[4] = (unsigned char)((code >> 24) & 0xff);
2677 n = 5;
2678 }
2679
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002680 if (_Pickler_Write(self, pdata, n) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002681 goto error;
2682 }
2683 else {
2684 /* Generate a normal global opcode if we are using a pickle
2685 protocol <= 2, or if the object is not registered in the
2686 extension registry. */
2687 PyObject *encoded;
2688 PyObject *(*unicode_encoder)(PyObject *);
2689
2690 gen_global:
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002691 if (_Pickler_Write(self, &global_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002692 goto error;
2693
2694 /* Since Python 3.0 now supports non-ASCII identifiers, we encode both
2695 the module name and the global name using UTF-8. We do so only when
2696 we are using the pickle protocol newer than version 3. This is to
2697 ensure compatibility with older Unpickler running on Python 2.x. */
2698 if (self->proto >= 3) {
2699 unicode_encoder = PyUnicode_AsUTF8String;
2700 }
2701 else {
2702 unicode_encoder = PyUnicode_AsASCIIString;
2703 }
2704
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00002705 /* For protocol < 3 and if the user didn't request against doing so,
2706 we convert module names to the old 2.x module names. */
2707 if (self->fix_imports) {
2708 PyObject *key;
2709 PyObject *item;
2710
2711 key = PyTuple_Pack(2, module_name, global_name);
2712 if (key == NULL)
2713 goto error;
2714 item = PyDict_GetItemWithError(name_mapping_3to2, key);
2715 Py_DECREF(key);
2716 if (item) {
2717 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
2718 PyErr_Format(PyExc_RuntimeError,
2719 "_compat_pickle.REVERSE_NAME_MAPPING values "
2720 "should be 2-tuples, not %.200s",
2721 Py_TYPE(item)->tp_name);
2722 goto error;
2723 }
2724 Py_CLEAR(module_name);
2725 Py_CLEAR(global_name);
2726 module_name = PyTuple_GET_ITEM(item, 0);
2727 global_name = PyTuple_GET_ITEM(item, 1);
2728 if (!PyUnicode_Check(module_name) ||
2729 !PyUnicode_Check(global_name)) {
2730 PyErr_Format(PyExc_RuntimeError,
2731 "_compat_pickle.REVERSE_NAME_MAPPING values "
2732 "should be pairs of str, not (%.200s, %.200s)",
2733 Py_TYPE(module_name)->tp_name,
2734 Py_TYPE(global_name)->tp_name);
2735 goto error;
2736 }
2737 Py_INCREF(module_name);
2738 Py_INCREF(global_name);
2739 }
2740 else if (PyErr_Occurred()) {
2741 goto error;
2742 }
2743
2744 item = PyDict_GetItemWithError(import_mapping_3to2, module_name);
2745 if (item) {
2746 if (!PyUnicode_Check(item)) {
2747 PyErr_Format(PyExc_RuntimeError,
2748 "_compat_pickle.REVERSE_IMPORT_MAPPING values "
2749 "should be strings, not %.200s",
2750 Py_TYPE(item)->tp_name);
2751 goto error;
2752 }
2753 Py_CLEAR(module_name);
2754 module_name = item;
2755 Py_INCREF(module_name);
2756 }
2757 else if (PyErr_Occurred()) {
2758 goto error;
2759 }
2760 }
2761
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002762 /* Save the name of the module. */
2763 encoded = unicode_encoder(module_name);
2764 if (encoded == NULL) {
2765 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2766 PyErr_Format(PicklingError,
2767 "can't pickle module identifier '%S' using "
2768 "pickle protocol %i", module_name, self->proto);
2769 goto error;
2770 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002771 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002772 PyBytes_GET_SIZE(encoded)) < 0) {
2773 Py_DECREF(encoded);
2774 goto error;
2775 }
2776 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002777 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002778 goto error;
2779
2780 /* Save the name of the module. */
2781 encoded = unicode_encoder(global_name);
2782 if (encoded == NULL) {
2783 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2784 PyErr_Format(PicklingError,
2785 "can't pickle global identifier '%S' using "
2786 "pickle protocol %i", global_name, self->proto);
2787 goto error;
2788 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002789 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002790 PyBytes_GET_SIZE(encoded)) < 0) {
2791 Py_DECREF(encoded);
2792 goto error;
2793 }
2794 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002795 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002796 goto error;
2797
2798 /* Memoize the object. */
2799 if (memo_put(self, obj) < 0)
2800 goto error;
2801 }
2802
2803 if (0) {
2804 error:
2805 status = -1;
2806 }
2807 Py_XDECREF(module_name);
2808 Py_XDECREF(global_name);
2809 Py_XDECREF(module);
2810
2811 return status;
2812}
2813
2814static int
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002815save_ellipsis(PicklerObject *self, PyObject *obj)
2816{
Łukasz Langadbd78252012-03-12 22:59:11 +01002817 PyObject *str = PyUnicode_FromString("Ellipsis");
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002818 int res;
Łukasz Langadbd78252012-03-12 22:59:11 +01002819 if (str == NULL)
Łukasz Langacad1a072012-03-12 23:41:07 +01002820 return -1;
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002821 res = save_global(self, Py_Ellipsis, str);
2822 Py_DECREF(str);
2823 return res;
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002824}
2825
2826static int
2827save_notimplemented(PicklerObject *self, PyObject *obj)
2828{
Łukasz Langadbd78252012-03-12 22:59:11 +01002829 PyObject *str = PyUnicode_FromString("NotImplemented");
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002830 int res;
Łukasz Langadbd78252012-03-12 22:59:11 +01002831 if (str == NULL)
Łukasz Langacad1a072012-03-12 23:41:07 +01002832 return -1;
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002833 res = save_global(self, Py_NotImplemented, str);
2834 Py_DECREF(str);
2835 return res;
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002836}
2837
2838static int
Alexandre Vassalotti19b6fa62013-11-30 16:06:39 -08002839save_singleton_type(PicklerObject *self, PyObject *obj, PyObject *singleton)
2840{
2841 PyObject *reduce_value;
2842 int status;
2843
2844 reduce_value = Py_BuildValue("O(O)", &PyType_Type, singleton);
2845 if (reduce_value == NULL) {
2846 return -1;
2847 }
2848 status = save_reduce(self, reduce_value, obj);
2849 Py_DECREF(reduce_value);
2850 return status;
2851}
2852
2853static int
2854save_type(PicklerObject *self, PyObject *obj)
2855{
Alexandre Vassalotti65846c62013-11-30 17:55:48 -08002856 if (obj == (PyObject *)&_PyNone_Type) {
Alexandre Vassalotti19b6fa62013-11-30 16:06:39 -08002857 return save_singleton_type(self, obj, Py_None);
2858 }
2859 else if (obj == (PyObject *)&PyEllipsis_Type) {
2860 return save_singleton_type(self, obj, Py_Ellipsis);
2861 }
Alexandre Vassalotti65846c62013-11-30 17:55:48 -08002862 else if (obj == (PyObject *)&_PyNotImplemented_Type) {
Alexandre Vassalotti19b6fa62013-11-30 16:06:39 -08002863 return save_singleton_type(self, obj, Py_NotImplemented);
2864 }
2865 return save_global(self, obj, NULL);
2866}
2867
2868static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002869save_pers(PicklerObject *self, PyObject *obj, PyObject *func)
2870{
2871 PyObject *pid = NULL;
2872 int status = 0;
2873
2874 const char persid_op = PERSID;
2875 const char binpersid_op = BINPERSID;
2876
2877 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002878 pid = _Pickler_FastCall(self, func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002879 if (pid == NULL)
2880 return -1;
2881
2882 if (pid != Py_None) {
2883 if (self->bin) {
2884 if (save(self, pid, 1) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002885 _Pickler_Write(self, &binpersid_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002886 goto error;
2887 }
2888 else {
2889 PyObject *pid_str = NULL;
2890 char *pid_ascii_bytes;
2891 Py_ssize_t size;
2892
2893 pid_str = PyObject_Str(pid);
2894 if (pid_str == NULL)
2895 goto error;
2896
2897 /* XXX: Should it check whether the persistent id only contains
2898 ASCII characters? And what if the pid contains embedded
2899 newlines? */
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00002900 pid_ascii_bytes = _PyUnicode_AsStringAndSize(pid_str, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002901 Py_DECREF(pid_str);
2902 if (pid_ascii_bytes == NULL)
2903 goto error;
2904
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002905 if (_Pickler_Write(self, &persid_op, 1) < 0 ||
2906 _Pickler_Write(self, pid_ascii_bytes, size) < 0 ||
2907 _Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002908 goto error;
2909 }
2910 status = 1;
2911 }
2912
2913 if (0) {
2914 error:
2915 status = -1;
2916 }
2917 Py_XDECREF(pid);
2918
2919 return status;
2920}
2921
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002922static PyObject *
2923get_class(PyObject *obj)
2924{
2925 PyObject *cls;
2926 static PyObject *str_class;
2927
2928 if (str_class == NULL) {
2929 str_class = PyUnicode_InternFromString("__class__");
2930 if (str_class == NULL)
2931 return NULL;
2932 }
2933 cls = PyObject_GetAttr(obj, str_class);
2934 if (cls == NULL) {
2935 if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
2936 PyErr_Clear();
2937 cls = (PyObject *) Py_TYPE(obj);
2938 Py_INCREF(cls);
2939 }
2940 }
2941 return cls;
2942}
2943
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002944/* We're saving obj, and args is the 2-thru-5 tuple returned by the
2945 * appropriate __reduce__ method for obj.
2946 */
2947static int
2948save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
2949{
2950 PyObject *callable;
2951 PyObject *argtup;
2952 PyObject *state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002953 PyObject *listitems = Py_None;
2954 PyObject *dictitems = Py_None;
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002955 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002956
2957 int use_newobj = self->proto >= 2;
2958
2959 const char reduce_op = REDUCE;
2960 const char build_op = BUILD;
2961 const char newobj_op = NEWOBJ;
2962
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002963 size = PyTuple_Size(args);
2964 if (size < 2 || size > 5) {
2965 PyErr_SetString(PicklingError, "tuple returned by "
2966 "__reduce__ must contain 2 through 5 elements");
2967 return -1;
2968 }
2969
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002970 if (!PyArg_UnpackTuple(args, "save_reduce", 2, 5,
2971 &callable, &argtup, &state, &listitems, &dictitems))
2972 return -1;
2973
2974 if (!PyCallable_Check(callable)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002975 PyErr_SetString(PicklingError, "first item of the tuple "
2976 "returned by __reduce__ must be callable");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002977 return -1;
2978 }
2979 if (!PyTuple_Check(argtup)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002980 PyErr_SetString(PicklingError, "second item of the tuple "
2981 "returned by __reduce__ must be a tuple");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002982 return -1;
2983 }
2984
2985 if (state == Py_None)
2986 state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002987
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002988 if (listitems == Py_None)
2989 listitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002990 else if (!PyIter_Check(listitems)) {
Alexandre Vassalotti00d83f22013-04-14 01:28:01 -07002991 PyErr_Format(PicklingError, "fourth element of the tuple "
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002992 "returned by __reduce__ must be an iterator, not %s",
2993 Py_TYPE(listitems)->tp_name);
2994 return -1;
2995 }
2996
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002997 if (dictitems == Py_None)
2998 dictitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002999 else if (!PyIter_Check(dictitems)) {
Alexandre Vassalotti00d83f22013-04-14 01:28:01 -07003000 PyErr_Format(PicklingError, "fifth element of the tuple "
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003001 "returned by __reduce__ must be an iterator, not %s",
3002 Py_TYPE(dictitems)->tp_name);
3003 return -1;
3004 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003005
3006 /* Protocol 2 special case: if callable's name is __newobj__, use
3007 NEWOBJ. */
3008 if (use_newobj) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003009 static PyObject *newobj_str = NULL, *name_str = NULL;
3010 PyObject *name;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003011
3012 if (newobj_str == NULL) {
3013 newobj_str = PyUnicode_InternFromString("__newobj__");
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003014 name_str = PyUnicode_InternFromString("__name__");
3015 if (newobj_str == NULL || name_str == NULL)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003016 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003017 }
3018
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003019 name = PyObject_GetAttr(callable, name_str);
3020 if (name == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003021 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3022 PyErr_Clear();
3023 else
3024 return -1;
3025 use_newobj = 0;
3026 }
3027 else {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003028 use_newobj = PyUnicode_Check(name) &&
3029 PyUnicode_Compare(name, newobj_str) == 0;
3030 Py_DECREF(name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003031 }
3032 }
3033 if (use_newobj) {
3034 PyObject *cls;
3035 PyObject *newargtup;
3036 PyObject *obj_class;
3037 int p;
3038
3039 /* Sanity checks. */
3040 if (Py_SIZE(argtup) < 1) {
3041 PyErr_SetString(PicklingError, "__newobj__ arglist is empty");
3042 return -1;
3043 }
3044
3045 cls = PyTuple_GET_ITEM(argtup, 0);
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003046 if (!PyType_Check(cls)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003047 PyErr_SetString(PicklingError, "args[0] from "
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003048 "__newobj__ args is not a type");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003049 return -1;
3050 }
3051
3052 if (obj != NULL) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003053 obj_class = get_class(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003054 p = obj_class != cls; /* true iff a problem */
3055 Py_DECREF(obj_class);
3056 if (p) {
3057 PyErr_SetString(PicklingError, "args[0] from "
3058 "__newobj__ args has the wrong class");
3059 return -1;
3060 }
3061 }
3062 /* XXX: These calls save() are prone to infinite recursion. Imagine
3063 what happen if the value returned by the __reduce__() method of
3064 some extension type contains another object of the same type. Ouch!
3065
3066 Here is a quick example, that I ran into, to illustrate what I
3067 mean:
3068
3069 >>> import pickle, copyreg
3070 >>> copyreg.dispatch_table.pop(complex)
3071 >>> pickle.dumps(1+2j)
3072 Traceback (most recent call last):
3073 ...
3074 RuntimeError: maximum recursion depth exceeded
3075
3076 Removing the complex class from copyreg.dispatch_table made the
3077 __reduce_ex__() method emit another complex object:
3078
3079 >>> (1+1j).__reduce_ex__(2)
3080 (<function __newobj__ at 0xb7b71c3c>,
3081 (<class 'complex'>, (1+1j)), None, None, None)
3082
3083 Thus when save() was called on newargstup (the 2nd item) recursion
3084 ensued. Of course, the bug was in the complex class which had a
3085 broken __getnewargs__() that emitted another complex object. But,
3086 the point, here, is it is quite easy to end up with a broken reduce
3087 function. */
3088
3089 /* Save the class and its __new__ arguments. */
3090 if (save(self, cls, 0) < 0)
3091 return -1;
3092
3093 newargtup = PyTuple_GetSlice(argtup, 1, Py_SIZE(argtup));
3094 if (newargtup == NULL)
3095 return -1;
3096
3097 p = save(self, newargtup, 0);
3098 Py_DECREF(newargtup);
3099 if (p < 0)
3100 return -1;
3101
3102 /* Add NEWOBJ opcode. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003103 if (_Pickler_Write(self, &newobj_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003104 return -1;
3105 }
3106 else { /* Not using NEWOBJ. */
3107 if (save(self, callable, 0) < 0 ||
3108 save(self, argtup, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003109 _Pickler_Write(self, &reduce_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003110 return -1;
3111 }
3112
3113 /* obj can be NULL when save_reduce() is used directly. A NULL obj means
3114 the caller do not want to memoize the object. Not particularly useful,
3115 but that is to mimic the behavior save_reduce() in pickle.py when
3116 obj is None. */
3117 if (obj && memo_put(self, obj) < 0)
3118 return -1;
3119
3120 if (listitems && batch_list(self, listitems) < 0)
3121 return -1;
3122
3123 if (dictitems && batch_dict(self, dictitems) < 0)
3124 return -1;
3125
3126 if (state) {
Victor Stinner121aab42011-09-29 23:40:53 +02003127 if (save(self, state, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003128 _Pickler_Write(self, &build_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003129 return -1;
3130 }
3131
3132 return 0;
3133}
3134
3135static int
3136save(PicklerObject *self, PyObject *obj, int pers_save)
3137{
3138 PyTypeObject *type;
3139 PyObject *reduce_func = NULL;
3140 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003141 int status = 0;
3142
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00003143 if (Py_EnterRecursiveCall(" while pickling an object"))
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003144 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003145
3146 /* The extra pers_save argument is necessary to avoid calling save_pers()
3147 on its returned object. */
3148 if (!pers_save && self->pers_func) {
3149 /* save_pers() returns:
3150 -1 to signal an error;
3151 0 if it did nothing successfully;
3152 1 if a persistent id was saved.
3153 */
3154 if ((status = save_pers(self, obj, self->pers_func)) != 0)
3155 goto done;
3156 }
3157
3158 type = Py_TYPE(obj);
3159
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003160 /* The old cPickle had an optimization that used switch-case statement
3161 dispatching on the first letter of the type name. This has was removed
3162 since benchmarks shown that this optimization was actually slowing
3163 things down. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003164
3165 /* Atom types; these aren't memoized, so don't check the memo. */
3166
3167 if (obj == Py_None) {
3168 status = save_none(self, obj);
3169 goto done;
3170 }
Łukasz Langaf3078fb2012-03-12 19:46:12 +01003171 else if (obj == Py_Ellipsis) {
3172 status = save_ellipsis(self, obj);
3173 goto done;
3174 }
3175 else if (obj == Py_NotImplemented) {
3176 status = save_notimplemented(self, obj);
3177 goto done;
3178 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003179 else if (obj == Py_False || obj == Py_True) {
3180 status = save_bool(self, obj);
3181 goto done;
3182 }
3183 else if (type == &PyLong_Type) {
3184 status = save_long(self, obj);
3185 goto done;
3186 }
3187 else if (type == &PyFloat_Type) {
3188 status = save_float(self, obj);
3189 goto done;
3190 }
3191
3192 /* Check the memo to see if it has the object. If so, generate
3193 a GET (or BINGET) opcode, instead of pickling the object
3194 once again. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003195 if (PyMemoTable_Get(self->memo, obj)) {
3196 if (memo_get(self, obj) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003197 goto error;
3198 goto done;
3199 }
3200
3201 if (type == &PyBytes_Type) {
3202 status = save_bytes(self, obj);
3203 goto done;
3204 }
3205 else if (type == &PyUnicode_Type) {
3206 status = save_unicode(self, obj);
3207 goto done;
3208 }
3209 else if (type == &PyDict_Type) {
3210 status = save_dict(self, obj);
3211 goto done;
3212 }
3213 else if (type == &PyList_Type) {
3214 status = save_list(self, obj);
3215 goto done;
3216 }
3217 else if (type == &PyTuple_Type) {
3218 status = save_tuple(self, obj);
3219 goto done;
3220 }
3221 else if (type == &PyType_Type) {
Alexandre Vassalotti19b6fa62013-11-30 16:06:39 -08003222 status = save_type(self, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003223 goto done;
3224 }
3225 else if (type == &PyFunction_Type) {
3226 status = save_global(self, obj, NULL);
3227 if (status < 0 && PyErr_ExceptionMatches(PickleError)) {
3228 /* fall back to reduce */
3229 PyErr_Clear();
3230 }
3231 else {
3232 goto done;
3233 }
3234 }
3235 else if (type == &PyCFunction_Type) {
3236 status = save_global(self, obj, NULL);
3237 goto done;
3238 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003239
3240 /* XXX: This part needs some unit tests. */
3241
3242 /* Get a reduction callable, and call it. This may come from
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003243 * self.dispatch_table, copyreg.dispatch_table, the object's
3244 * __reduce_ex__ method, or the object's __reduce__ method.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003245 */
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003246 if (self->dispatch_table == NULL) {
3247 reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
3248 /* PyDict_GetItem() unlike PyObject_GetItem() and
3249 PyObject_GetAttr() returns a borrowed ref */
3250 Py_XINCREF(reduce_func);
3251 } else {
3252 reduce_func = PyObject_GetItem(self->dispatch_table, (PyObject *)type);
3253 if (reduce_func == NULL) {
3254 if (PyErr_ExceptionMatches(PyExc_KeyError))
3255 PyErr_Clear();
3256 else
3257 goto error;
3258 }
3259 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003260 if (reduce_func != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003261 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003262 reduce_value = _Pickler_FastCall(self, reduce_func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003263 }
Antoine Pitrouffd41d92011-10-04 09:23:04 +02003264 else if (PyType_IsSubtype(type, &PyType_Type)) {
3265 status = save_global(self, obj, NULL);
3266 goto done;
3267 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003268 else {
3269 static PyObject *reduce_str = NULL;
3270 static PyObject *reduce_ex_str = NULL;
3271
3272 /* Cache the name of the reduce methods. */
3273 if (reduce_str == NULL) {
3274 reduce_str = PyUnicode_InternFromString("__reduce__");
3275 if (reduce_str == NULL)
3276 goto error;
3277 reduce_ex_str = PyUnicode_InternFromString("__reduce_ex__");
3278 if (reduce_ex_str == NULL)
3279 goto error;
3280 }
3281
3282 /* XXX: If the __reduce__ method is defined, __reduce_ex__ is
3283 automatically defined as __reduce__. While this is convenient, this
3284 make it impossible to know which method was actually called. Of
3285 course, this is not a big deal. But still, it would be nice to let
3286 the user know which method was called when something go
3287 wrong. Incidentally, this means if __reduce_ex__ is not defined, we
3288 don't actually have to check for a __reduce__ method. */
3289
3290 /* Check for a __reduce_ex__ method. */
3291 reduce_func = PyObject_GetAttr(obj, reduce_ex_str);
3292 if (reduce_func != NULL) {
3293 PyObject *proto;
3294 proto = PyLong_FromLong(self->proto);
3295 if (proto != NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003296 reduce_value = _Pickler_FastCall(self, reduce_func, proto);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003297 }
3298 }
3299 else {
3300 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3301 PyErr_Clear();
3302 else
3303 goto error;
3304 /* Check for a __reduce__ method. */
3305 reduce_func = PyObject_GetAttr(obj, reduce_str);
3306 if (reduce_func != NULL) {
3307 reduce_value = PyObject_Call(reduce_func, empty_tuple, NULL);
3308 }
3309 else {
3310 PyErr_Format(PicklingError, "can't pickle '%.200s' object: %R",
3311 type->tp_name, obj);
3312 goto error;
3313 }
3314 }
3315 }
3316
3317 if (reduce_value == NULL)
3318 goto error;
3319
3320 if (PyUnicode_Check(reduce_value)) {
3321 status = save_global(self, obj, reduce_value);
3322 goto done;
3323 }
3324
3325 if (!PyTuple_Check(reduce_value)) {
3326 PyErr_SetString(PicklingError,
3327 "__reduce__ must return a string or tuple");
3328 goto error;
3329 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003330
3331 status = save_reduce(self, reduce_value, obj);
3332
3333 if (0) {
3334 error:
3335 status = -1;
3336 }
3337 done:
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003338 Py_LeaveRecursiveCall();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003339 Py_XDECREF(reduce_func);
3340 Py_XDECREF(reduce_value);
3341
3342 return status;
3343}
3344
3345static int
3346dump(PicklerObject *self, PyObject *obj)
3347{
3348 const char stop_op = STOP;
3349
3350 if (self->proto >= 2) {
3351 char header[2];
3352
3353 header[0] = PROTO;
3354 assert(self->proto >= 0 && self->proto < 256);
3355 header[1] = (unsigned char)self->proto;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003356 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003357 return -1;
3358 }
3359
3360 if (save(self, obj, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003361 _Pickler_Write(self, &stop_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003362 return -1;
3363
3364 return 0;
3365}
3366
3367PyDoc_STRVAR(Pickler_clear_memo_doc,
3368"clear_memo() -> None. Clears the pickler's \"memo\"."
3369"\n"
3370"The memo is the data structure that remembers which objects the\n"
3371"pickler has already seen, so that shared or recursive objects are\n"
3372"pickled by reference and not by value. This method is useful when\n"
3373"re-using picklers.");
3374
3375static PyObject *
3376Pickler_clear_memo(PicklerObject *self)
3377{
3378 if (self->memo)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003379 PyMemoTable_Clear(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003380
3381 Py_RETURN_NONE;
3382}
3383
3384PyDoc_STRVAR(Pickler_dump_doc,
3385"dump(obj) -> None. Write a pickled representation of obj to the open file.");
3386
3387static PyObject *
3388Pickler_dump(PicklerObject *self, PyObject *args)
3389{
3390 PyObject *obj;
3391
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003392 /* Check whether the Pickler was initialized correctly (issue3664).
3393 Developers often forget to call __init__() in their subclasses, which
3394 would trigger a segfault without this check. */
3395 if (self->write == NULL) {
Victor Stinner121aab42011-09-29 23:40:53 +02003396 PyErr_Format(PicklingError,
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003397 "Pickler.__init__() was not called by %s.__init__()",
3398 Py_TYPE(self)->tp_name);
3399 return NULL;
3400 }
3401
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003402 if (!PyArg_ParseTuple(args, "O:dump", &obj))
3403 return NULL;
3404
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003405 if (_Pickler_ClearBuffer(self) < 0)
3406 return NULL;
3407
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003408 if (dump(self, obj) < 0)
3409 return NULL;
3410
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003411 if (_Pickler_FlushToFile(self) < 0)
3412 return NULL;
3413
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003414 Py_RETURN_NONE;
3415}
3416
3417static struct PyMethodDef Pickler_methods[] = {
3418 {"dump", (PyCFunction)Pickler_dump, METH_VARARGS,
3419 Pickler_dump_doc},
3420 {"clear_memo", (PyCFunction)Pickler_clear_memo, METH_NOARGS,
3421 Pickler_clear_memo_doc},
3422 {NULL, NULL} /* sentinel */
3423};
3424
3425static void
3426Pickler_dealloc(PicklerObject *self)
3427{
3428 PyObject_GC_UnTrack(self);
3429
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003430 Py_XDECREF(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003431 Py_XDECREF(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003432 Py_XDECREF(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003433 Py_XDECREF(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003434 Py_XDECREF(self->arg);
3435 Py_XDECREF(self->fast_memo);
3436
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003437 PyMemoTable_Del(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003438
3439 Py_TYPE(self)->tp_free((PyObject *)self);
3440}
3441
3442static int
3443Pickler_traverse(PicklerObject *self, visitproc visit, void *arg)
3444{
3445 Py_VISIT(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003446 Py_VISIT(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003447 Py_VISIT(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003448 Py_VISIT(self->arg);
3449 Py_VISIT(self->fast_memo);
3450 return 0;
3451}
3452
3453static int
3454Pickler_clear(PicklerObject *self)
3455{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003456 Py_CLEAR(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003457 Py_CLEAR(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003458 Py_CLEAR(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003459 Py_CLEAR(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003460 Py_CLEAR(self->arg);
3461 Py_CLEAR(self->fast_memo);
3462
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003463 if (self->memo != NULL) {
3464 PyMemoTable *memo = self->memo;
3465 self->memo = NULL;
3466 PyMemoTable_Del(memo);
3467 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003468 return 0;
3469}
3470
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003471
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003472PyDoc_STRVAR(Pickler_doc,
3473"Pickler(file, protocol=None)"
3474"\n"
3475"This takes a binary file for writing a pickle data stream.\n"
3476"\n"
3477"The optional protocol argument tells the pickler to use the\n"
3478"given protocol; supported protocols are 0, 1, 2, 3. The default\n"
3479"protocol is 3; a backward-incompatible protocol designed for\n"
3480"Python 3.0.\n"
3481"\n"
3482"Specifying a negative protocol version selects the highest\n"
3483"protocol version supported. The higher the protocol used, the\n"
3484"more recent the version of Python needed to read the pickle\n"
3485"produced.\n"
3486"\n"
3487"The file argument must have a write() method that accepts a single\n"
3488"bytes argument. It can thus be a file object opened for binary\n"
3489"writing, a io.BytesIO instance, or any other custom object that\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003490"meets this interface.\n"
3491"\n"
3492"If fix_imports is True and protocol is less than 3, pickle will try to\n"
3493"map the new Python 3.x names to the old module names used in Python\n"
3494"2.x, so that the pickle data stream is readable with Python 2.x.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003495
3496static int
3497Pickler_init(PicklerObject *self, PyObject *args, PyObject *kwds)
3498{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003499 static char *kwlist[] = {"file", "protocol", "fix_imports", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003500 PyObject *file;
3501 PyObject *proto_obj = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003502 PyObject *fix_imports = Py_True;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02003503 _Py_IDENTIFIER(persistent_id);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003504 _Py_IDENTIFIER(dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003505
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003506 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:Pickler",
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003507 kwlist, &file, &proto_obj, &fix_imports))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003508 return -1;
3509
3510 /* In case of multiple __init__() calls, clear previous content. */
3511 if (self->write != NULL)
3512 (void)Pickler_clear(self);
3513
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003514 if (_Pickler_SetProtocol(self, proto_obj, fix_imports) < 0)
3515 return -1;
3516
3517 if (_Pickler_SetOutputStream(self, file) < 0)
3518 return -1;
3519
3520 /* memo and output_buffer may have already been created in _Pickler_New */
3521 if (self->memo == NULL) {
3522 self->memo = PyMemoTable_New();
3523 if (self->memo == NULL)
3524 return -1;
3525 }
3526 self->output_len = 0;
3527 if (self->output_buffer == NULL) {
3528 self->max_output_len = WRITE_BUF_SIZE;
3529 self->output_buffer = PyBytes_FromStringAndSize(NULL,
3530 self->max_output_len);
3531 if (self->output_buffer == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003532 return -1;
3533 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003534
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003535 self->arg = NULL;
3536 self->fast = 0;
3537 self->fast_nesting = 0;
3538 self->fast_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003539 self->pers_func = NULL;
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02003540 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_id)) {
3541 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
3542 &PyId_persistent_id);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003543 if (self->pers_func == NULL)
3544 return -1;
3545 }
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003546 self->dispatch_table = NULL;
3547 if (_PyObject_HasAttrId((PyObject *)self, &PyId_dispatch_table)) {
3548 self->dispatch_table = _PyObject_GetAttrId((PyObject *)self,
3549 &PyId_dispatch_table);
3550 if (self->dispatch_table == NULL)
3551 return -1;
3552 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003553 return 0;
3554}
3555
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003556/* Define a proxy object for the Pickler's internal memo object. This is to
3557 * avoid breaking code like:
3558 * pickler.memo.clear()
3559 * and
3560 * pickler.memo = saved_memo
3561 * Is this a good idea? Not really, but we don't want to break code that uses
3562 * it. Note that we don't implement the entire mapping API here. This is
3563 * intentional, as these should be treated as black-box implementation details.
3564 */
3565
3566typedef struct {
3567 PyObject_HEAD
3568 PicklerObject *pickler; /* Pickler whose memo table we're proxying. */
3569} PicklerMemoProxyObject;
3570
3571PyDoc_STRVAR(pmp_clear_doc,
3572"memo.clear() -> None. Remove all items from memo.");
3573
3574static PyObject *
3575pmp_clear(PicklerMemoProxyObject *self)
3576{
3577 if (self->pickler->memo)
3578 PyMemoTable_Clear(self->pickler->memo);
3579 Py_RETURN_NONE;
3580}
3581
3582PyDoc_STRVAR(pmp_copy_doc,
3583"memo.copy() -> new_memo. Copy the memo to a new object.");
3584
3585static PyObject *
3586pmp_copy(PicklerMemoProxyObject *self)
3587{
3588 Py_ssize_t i;
3589 PyMemoTable *memo;
3590 PyObject *new_memo = PyDict_New();
3591 if (new_memo == NULL)
3592 return NULL;
3593
3594 memo = self->pickler->memo;
3595 for (i = 0; i < memo->mt_allocated; ++i) {
3596 PyMemoEntry entry = memo->mt_table[i];
3597 if (entry.me_key != NULL) {
3598 int status;
3599 PyObject *key, *value;
3600
3601 key = PyLong_FromVoidPtr(entry.me_key);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003602 value = Py_BuildValue("nO", entry.me_value, entry.me_key);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003603
3604 if (key == NULL || value == NULL) {
3605 Py_XDECREF(key);
3606 Py_XDECREF(value);
3607 goto error;
3608 }
3609 status = PyDict_SetItem(new_memo, key, value);
3610 Py_DECREF(key);
3611 Py_DECREF(value);
3612 if (status < 0)
3613 goto error;
3614 }
3615 }
3616 return new_memo;
3617
3618 error:
3619 Py_XDECREF(new_memo);
3620 return NULL;
3621}
3622
3623PyDoc_STRVAR(pmp_reduce_doc,
3624"memo.__reduce__(). Pickling support.");
3625
3626static PyObject *
3627pmp_reduce(PicklerMemoProxyObject *self, PyObject *args)
3628{
3629 PyObject *reduce_value, *dict_args;
3630 PyObject *contents = pmp_copy(self);
3631 if (contents == NULL)
3632 return NULL;
3633
3634 reduce_value = PyTuple_New(2);
3635 if (reduce_value == NULL) {
3636 Py_DECREF(contents);
3637 return NULL;
3638 }
3639 dict_args = PyTuple_New(1);
3640 if (dict_args == NULL) {
3641 Py_DECREF(contents);
3642 Py_DECREF(reduce_value);
3643 return NULL;
3644 }
3645 PyTuple_SET_ITEM(dict_args, 0, contents);
3646 Py_INCREF((PyObject *)&PyDict_Type);
3647 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
3648 PyTuple_SET_ITEM(reduce_value, 1, dict_args);
3649 return reduce_value;
3650}
3651
3652static PyMethodDef picklerproxy_methods[] = {
3653 {"clear", (PyCFunction)pmp_clear, METH_NOARGS, pmp_clear_doc},
3654 {"copy", (PyCFunction)pmp_copy, METH_NOARGS, pmp_copy_doc},
3655 {"__reduce__", (PyCFunction)pmp_reduce, METH_VARARGS, pmp_reduce_doc},
3656 {NULL, NULL} /* sentinel */
3657};
3658
3659static void
3660PicklerMemoProxy_dealloc(PicklerMemoProxyObject *self)
3661{
3662 PyObject_GC_UnTrack(self);
3663 Py_XDECREF(self->pickler);
3664 PyObject_GC_Del((PyObject *)self);
3665}
3666
3667static int
3668PicklerMemoProxy_traverse(PicklerMemoProxyObject *self,
3669 visitproc visit, void *arg)
3670{
3671 Py_VISIT(self->pickler);
3672 return 0;
3673}
3674
3675static int
3676PicklerMemoProxy_clear(PicklerMemoProxyObject *self)
3677{
3678 Py_CLEAR(self->pickler);
3679 return 0;
3680}
3681
3682static PyTypeObject PicklerMemoProxyType = {
3683 PyVarObject_HEAD_INIT(NULL, 0)
3684 "_pickle.PicklerMemoProxy", /*tp_name*/
3685 sizeof(PicklerMemoProxyObject), /*tp_basicsize*/
3686 0,
3687 (destructor)PicklerMemoProxy_dealloc, /* tp_dealloc */
3688 0, /* tp_print */
3689 0, /* tp_getattr */
3690 0, /* tp_setattr */
3691 0, /* tp_compare */
3692 0, /* tp_repr */
3693 0, /* tp_as_number */
3694 0, /* tp_as_sequence */
3695 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00003696 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003697 0, /* tp_call */
3698 0, /* tp_str */
3699 PyObject_GenericGetAttr, /* tp_getattro */
3700 PyObject_GenericSetAttr, /* tp_setattro */
3701 0, /* tp_as_buffer */
3702 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3703 0, /* tp_doc */
3704 (traverseproc)PicklerMemoProxy_traverse, /* tp_traverse */
3705 (inquiry)PicklerMemoProxy_clear, /* tp_clear */
3706 0, /* tp_richcompare */
3707 0, /* tp_weaklistoffset */
3708 0, /* tp_iter */
3709 0, /* tp_iternext */
3710 picklerproxy_methods, /* tp_methods */
3711};
3712
3713static PyObject *
3714PicklerMemoProxy_New(PicklerObject *pickler)
3715{
3716 PicklerMemoProxyObject *self;
3717
3718 self = PyObject_GC_New(PicklerMemoProxyObject, &PicklerMemoProxyType);
3719 if (self == NULL)
3720 return NULL;
3721 Py_INCREF(pickler);
3722 self->pickler = pickler;
3723 PyObject_GC_Track(self);
3724 return (PyObject *)self;
3725}
3726
3727/*****************************************************************************/
3728
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003729static PyObject *
3730Pickler_get_memo(PicklerObject *self)
3731{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003732 return PicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003733}
3734
3735static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003736Pickler_set_memo(PicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003737{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003738 PyMemoTable *new_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003739
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003740 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003741 PyErr_SetString(PyExc_TypeError,
3742 "attribute deletion is not supported");
3743 return -1;
3744 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003745
3746 if (Py_TYPE(obj) == &PicklerMemoProxyType) {
3747 PicklerObject *pickler =
3748 ((PicklerMemoProxyObject *)obj)->pickler;
3749
3750 new_memo = PyMemoTable_Copy(pickler->memo);
3751 if (new_memo == NULL)
3752 return -1;
3753 }
3754 else if (PyDict_Check(obj)) {
3755 Py_ssize_t i = 0;
3756 PyObject *key, *value;
3757
3758 new_memo = PyMemoTable_New();
3759 if (new_memo == NULL)
3760 return -1;
3761
3762 while (PyDict_Next(obj, &i, &key, &value)) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003763 Py_ssize_t memo_id;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003764 PyObject *memo_obj;
3765
3766 if (!PyTuple_Check(value) || Py_SIZE(value) != 2) {
3767 PyErr_SetString(PyExc_TypeError,
3768 "'memo' values must be 2-item tuples");
3769 goto error;
3770 }
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003771 memo_id = PyLong_AsSsize_t(PyTuple_GET_ITEM(value, 0));
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003772 if (memo_id == -1 && PyErr_Occurred())
3773 goto error;
3774 memo_obj = PyTuple_GET_ITEM(value, 1);
3775 if (PyMemoTable_Set(new_memo, memo_obj, memo_id) < 0)
3776 goto error;
3777 }
3778 }
3779 else {
3780 PyErr_Format(PyExc_TypeError,
3781 "'memo' attribute must be an PicklerMemoProxy object"
3782 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003783 return -1;
3784 }
3785
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003786 PyMemoTable_Del(self->memo);
3787 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003788
3789 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003790
3791 error:
3792 if (new_memo)
3793 PyMemoTable_Del(new_memo);
3794 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003795}
3796
3797static PyObject *
3798Pickler_get_persid(PicklerObject *self)
3799{
3800 if (self->pers_func == NULL)
3801 PyErr_SetString(PyExc_AttributeError, "persistent_id");
3802 else
3803 Py_INCREF(self->pers_func);
3804 return self->pers_func;
3805}
3806
3807static int
3808Pickler_set_persid(PicklerObject *self, PyObject *value)
3809{
3810 PyObject *tmp;
3811
3812 if (value == NULL) {
3813 PyErr_SetString(PyExc_TypeError,
3814 "attribute deletion is not supported");
3815 return -1;
3816 }
3817 if (!PyCallable_Check(value)) {
3818 PyErr_SetString(PyExc_TypeError,
3819 "persistent_id must be a callable taking one argument");
3820 return -1;
3821 }
3822
3823 tmp = self->pers_func;
3824 Py_INCREF(value);
3825 self->pers_func = value;
3826 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
3827
3828 return 0;
3829}
3830
3831static PyMemberDef Pickler_members[] = {
3832 {"bin", T_INT, offsetof(PicklerObject, bin)},
3833 {"fast", T_INT, offsetof(PicklerObject, fast)},
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003834 {"dispatch_table", T_OBJECT_EX, offsetof(PicklerObject, dispatch_table)},
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003835 {NULL}
3836};
3837
3838static PyGetSetDef Pickler_getsets[] = {
3839 {"memo", (getter)Pickler_get_memo,
3840 (setter)Pickler_set_memo},
3841 {"persistent_id", (getter)Pickler_get_persid,
3842 (setter)Pickler_set_persid},
3843 {NULL}
3844};
3845
3846static PyTypeObject Pickler_Type = {
3847 PyVarObject_HEAD_INIT(NULL, 0)
3848 "_pickle.Pickler" , /*tp_name*/
3849 sizeof(PicklerObject), /*tp_basicsize*/
3850 0, /*tp_itemsize*/
3851 (destructor)Pickler_dealloc, /*tp_dealloc*/
3852 0, /*tp_print*/
3853 0, /*tp_getattr*/
3854 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00003855 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003856 0, /*tp_repr*/
3857 0, /*tp_as_number*/
3858 0, /*tp_as_sequence*/
3859 0, /*tp_as_mapping*/
3860 0, /*tp_hash*/
3861 0, /*tp_call*/
3862 0, /*tp_str*/
3863 0, /*tp_getattro*/
3864 0, /*tp_setattro*/
3865 0, /*tp_as_buffer*/
3866 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3867 Pickler_doc, /*tp_doc*/
3868 (traverseproc)Pickler_traverse, /*tp_traverse*/
3869 (inquiry)Pickler_clear, /*tp_clear*/
3870 0, /*tp_richcompare*/
3871 0, /*tp_weaklistoffset*/
3872 0, /*tp_iter*/
3873 0, /*tp_iternext*/
3874 Pickler_methods, /*tp_methods*/
3875 Pickler_members, /*tp_members*/
3876 Pickler_getsets, /*tp_getset*/
3877 0, /*tp_base*/
3878 0, /*tp_dict*/
3879 0, /*tp_descr_get*/
3880 0, /*tp_descr_set*/
3881 0, /*tp_dictoffset*/
3882 (initproc)Pickler_init, /*tp_init*/
3883 PyType_GenericAlloc, /*tp_alloc*/
3884 PyType_GenericNew, /*tp_new*/
3885 PyObject_GC_Del, /*tp_free*/
3886 0, /*tp_is_gc*/
3887};
3888
Victor Stinner121aab42011-09-29 23:40:53 +02003889/* Temporary helper for calling self.find_class().
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003890
3891 XXX: It would be nice to able to avoid Python function call overhead, by
3892 using directly the C version of find_class(), when find_class() is not
3893 overridden by a subclass. Although, this could become rather hackish. A
3894 simpler optimization would be to call the C function when self is not a
3895 subclass instance. */
3896static PyObject *
3897find_class(UnpicklerObject *self, PyObject *module_name, PyObject *global_name)
3898{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02003899 _Py_IDENTIFIER(find_class);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02003900
3901 return _PyObject_CallMethodId((PyObject *)self, &PyId_find_class, "OO",
3902 module_name, global_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003903}
3904
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003905static Py_ssize_t
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003906marker(UnpicklerObject *self)
3907{
3908 if (self->num_marks < 1) {
3909 PyErr_SetString(UnpicklingError, "could not find MARK");
3910 return -1;
3911 }
3912
3913 return self->marks[--self->num_marks];
3914}
3915
3916static int
3917load_none(UnpicklerObject *self)
3918{
3919 PDATA_APPEND(self->stack, Py_None, -1);
3920 return 0;
3921}
3922
3923static int
3924bad_readline(void)
3925{
3926 PyErr_SetString(UnpicklingError, "pickle data was truncated");
3927 return -1;
3928}
3929
3930static int
3931load_int(UnpicklerObject *self)
3932{
3933 PyObject *value;
3934 char *endptr, *s;
3935 Py_ssize_t len;
3936 long x;
3937
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003938 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003939 return -1;
3940 if (len < 2)
3941 return bad_readline();
3942
3943 errno = 0;
Victor Stinner121aab42011-09-29 23:40:53 +02003944 /* XXX: Should the base argument of strtol() be explicitly set to 10?
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003945 XXX(avassalotti): Should this uses PyOS_strtol()? */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003946 x = strtol(s, &endptr, 0);
3947
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003948 if (errno || (*endptr != '\n' && *endptr != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003949 /* Hm, maybe we've got something long. Let's try reading
Serhiy Storchaka95949422013-08-27 19:40:23 +03003950 * it as a Python int object. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003951 errno = 0;
3952 /* XXX: Same thing about the base here. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003953 value = PyLong_FromString(s, NULL, 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003954 if (value == NULL) {
3955 PyErr_SetString(PyExc_ValueError,
3956 "could not convert string to int");
3957 return -1;
3958 }
3959 }
3960 else {
3961 if (len == 3 && (x == 0 || x == 1)) {
3962 if ((value = PyBool_FromLong(x)) == NULL)
3963 return -1;
3964 }
3965 else {
3966 if ((value = PyLong_FromLong(x)) == NULL)
3967 return -1;
3968 }
3969 }
3970
3971 PDATA_PUSH(self->stack, value, -1);
3972 return 0;
3973}
3974
3975static int
3976load_bool(UnpicklerObject *self, PyObject *boolean)
3977{
3978 assert(boolean == Py_True || boolean == Py_False);
3979 PDATA_APPEND(self->stack, boolean, -1);
3980 return 0;
3981}
3982
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003983/* s contains x bytes of an unsigned little-endian integer. Return its value
3984 * as a C Py_ssize_t, or -1 if it's higher than PY_SSIZE_T_MAX.
3985 */
3986static Py_ssize_t
3987calc_binsize(char *bytes, int size)
3988{
3989 unsigned char *s = (unsigned char *)bytes;
3990 size_t x = 0;
3991
3992 assert(size == 4);
3993
3994 x = (size_t) s[0];
3995 x |= (size_t) s[1] << 8;
3996 x |= (size_t) s[2] << 16;
3997 x |= (size_t) s[3] << 24;
3998
3999 if (x > PY_SSIZE_T_MAX)
4000 return -1;
4001 else
4002 return (Py_ssize_t) x;
4003}
4004
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004005/* s contains x bytes of a little-endian integer. Return its value as a
4006 * C int. Obscure: when x is 1 or 2, this is an unsigned little-endian
4007 * int, but when x is 4 it's a signed one. This is an historical source
4008 * of x-platform bugs.
4009 */
4010static long
4011calc_binint(char *bytes, int size)
4012{
4013 unsigned char *s = (unsigned char *)bytes;
4014 int i = size;
4015 long x = 0;
4016
4017 for (i = 0; i < size; i++) {
4018 x |= (long)s[i] << (i * 8);
4019 }
4020
4021 /* Unlike BININT1 and BININT2, BININT (more accurately BININT4)
4022 * is signed, so on a box with longs bigger than 4 bytes we need
4023 * to extend a BININT's sign bit to the full width.
4024 */
4025 if (SIZEOF_LONG > 4 && size == 4) {
4026 x |= -(x & (1L << 31));
4027 }
4028
4029 return x;
4030}
4031
4032static int
4033load_binintx(UnpicklerObject *self, char *s, int size)
4034{
4035 PyObject *value;
4036 long x;
4037
4038 x = calc_binint(s, size);
4039
4040 if ((value = PyLong_FromLong(x)) == NULL)
4041 return -1;
4042
4043 PDATA_PUSH(self->stack, value, -1);
4044 return 0;
4045}
4046
4047static int
4048load_binint(UnpicklerObject *self)
4049{
4050 char *s;
4051
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004052 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004053 return -1;
4054
4055 return load_binintx(self, s, 4);
4056}
4057
4058static int
4059load_binint1(UnpicklerObject *self)
4060{
4061 char *s;
4062
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004063 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004064 return -1;
4065
4066 return load_binintx(self, s, 1);
4067}
4068
4069static int
4070load_binint2(UnpicklerObject *self)
4071{
4072 char *s;
4073
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004074 if (_Unpickler_Read(self, &s, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004075 return -1;
4076
4077 return load_binintx(self, s, 2);
4078}
4079
4080static int
4081load_long(UnpicklerObject *self)
4082{
4083 PyObject *value;
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004084 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004085 Py_ssize_t len;
4086
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004087 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004088 return -1;
4089 if (len < 2)
4090 return bad_readline();
4091
Mark Dickinson8dd05142009-01-20 20:43:58 +00004092 /* s[len-2] will usually be 'L' (and s[len-1] is '\n'); we need to remove
4093 the 'L' before calling PyLong_FromString. In order to maintain
4094 compatibility with Python 3.0.0, we don't actually *require*
4095 the 'L' to be present. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004096 if (s[len-2] == 'L')
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004097 s[len-2] = '\0';
Alexandre Vassalottie4bccb72009-01-24 01:47:57 +00004098 /* XXX: Should the base argument explicitly set to 10? */
4099 value = PyLong_FromString(s, NULL, 0);
Mark Dickinson8dd05142009-01-20 20:43:58 +00004100 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004101 return -1;
4102
4103 PDATA_PUSH(self->stack, value, -1);
4104 return 0;
4105}
4106
4107/* 'size' bytes contain the # of bytes of little-endian 256's-complement
4108 * data following.
4109 */
4110static int
4111load_counted_long(UnpicklerObject *self, int size)
4112{
4113 PyObject *value;
4114 char *nbytes;
4115 char *pdata;
4116
4117 assert(size == 1 || size == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004118 if (_Unpickler_Read(self, &nbytes, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004119 return -1;
4120
4121 size = calc_binint(nbytes, size);
4122 if (size < 0) {
4123 /* Corrupt or hostile pickle -- we never write one like this */
4124 PyErr_SetString(UnpicklingError,
4125 "LONG pickle has negative byte count");
4126 return -1;
4127 }
4128
4129 if (size == 0)
4130 value = PyLong_FromLong(0L);
4131 else {
4132 /* Read the raw little-endian bytes and convert. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004133 if (_Unpickler_Read(self, &pdata, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004134 return -1;
4135 value = _PyLong_FromByteArray((unsigned char *)pdata, (size_t)size,
4136 1 /* little endian */ , 1 /* signed */ );
4137 }
4138 if (value == NULL)
4139 return -1;
4140 PDATA_PUSH(self->stack, value, -1);
4141 return 0;
4142}
4143
4144static int
4145load_float(UnpicklerObject *self)
4146{
4147 PyObject *value;
4148 char *endptr, *s;
4149 Py_ssize_t len;
4150 double d;
4151
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004152 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004153 return -1;
4154 if (len < 2)
4155 return bad_readline();
4156
4157 errno = 0;
Mark Dickinson725bfd82009-05-03 20:33:40 +00004158 d = PyOS_string_to_double(s, &endptr, PyExc_OverflowError);
4159 if (d == -1.0 && PyErr_Occurred())
4160 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004161 if ((endptr[0] != '\n') && (endptr[0] != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004162 PyErr_SetString(PyExc_ValueError, "could not convert string to float");
4163 return -1;
4164 }
Mark Dickinson725bfd82009-05-03 20:33:40 +00004165 value = PyFloat_FromDouble(d);
4166 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004167 return -1;
4168
4169 PDATA_PUSH(self->stack, value, -1);
4170 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004171}
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004172
4173static int
4174load_binfloat(UnpicklerObject *self)
4175{
4176 PyObject *value;
4177 double x;
4178 char *s;
4179
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004180 if (_Unpickler_Read(self, &s, 8) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004181 return -1;
4182
4183 x = _PyFloat_Unpack8((unsigned char *)s, 0);
4184 if (x == -1.0 && PyErr_Occurred())
4185 return -1;
4186
4187 if ((value = PyFloat_FromDouble(x)) == NULL)
4188 return -1;
4189
4190 PDATA_PUSH(self->stack, value, -1);
4191 return 0;
4192}
4193
4194static int
4195load_string(UnpicklerObject *self)
4196{
4197 PyObject *bytes;
4198 PyObject *str = NULL;
4199 Py_ssize_t len;
4200 char *s, *p;
4201
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004202 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004203 return -1;
Antoine Pitrou3034efd2013-04-15 21:51:09 +02004204 if (len < 2)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004205 return bad_readline();
4206 if ((s = strdup(s)) == NULL) {
4207 PyErr_NoMemory();
4208 return -1;
4209 }
4210
4211 /* Strip outermost quotes */
Antoine Pitrou3034efd2013-04-15 21:51:09 +02004212 while (len > 0 && s[len - 1] <= ' ')
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004213 len--;
Antoine Pitrou3034efd2013-04-15 21:51:09 +02004214 if (len > 1 && s[0] == '"' && s[len - 1] == '"') {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004215 s[len - 1] = '\0';
4216 p = s + 1;
4217 len -= 2;
4218 }
Antoine Pitrou3034efd2013-04-15 21:51:09 +02004219 else if (len > 1 && s[0] == '\'' && s[len - 1] == '\'') {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004220 s[len - 1] = '\0';
4221 p = s + 1;
4222 len -= 2;
4223 }
4224 else {
4225 free(s);
4226 PyErr_SetString(PyExc_ValueError, "insecure string pickle");
4227 return -1;
4228 }
4229
4230 /* Use the PyBytes API to decode the string, since that is what is used
4231 to encode, and then coerce the result to Unicode. */
4232 bytes = PyBytes_DecodeEscape(p, len, NULL, 0, NULL);
4233 free(s);
4234 if (bytes == NULL)
4235 return -1;
4236 str = PyUnicode_FromEncodedObject(bytes, self->encoding, self->errors);
4237 Py_DECREF(bytes);
4238 if (str == NULL)
4239 return -1;
4240
4241 PDATA_PUSH(self->stack, str, -1);
4242 return 0;
4243}
4244
4245static int
4246load_binbytes(UnpicklerObject *self)
4247{
4248 PyObject *bytes;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004249 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004250 char *s;
4251
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004252 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004253 return -1;
4254
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004255 x = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004256 if (x < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004257 PyErr_Format(PyExc_OverflowError,
4258 "BINBYTES exceeds system's maximum size of %zd bytes",
4259 PY_SSIZE_T_MAX
4260 );
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004261 return -1;
4262 }
4263
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004264 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004265 return -1;
4266 bytes = PyBytes_FromStringAndSize(s, x);
4267 if (bytes == NULL)
4268 return -1;
4269
4270 PDATA_PUSH(self->stack, bytes, -1);
4271 return 0;
4272}
4273
4274static int
4275load_short_binbytes(UnpicklerObject *self)
4276{
4277 PyObject *bytes;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004278 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004279 char *s;
4280
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004281 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004282 return -1;
4283
4284 x = (unsigned char)s[0];
4285
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004286 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004287 return -1;
4288
4289 bytes = PyBytes_FromStringAndSize(s, x);
4290 if (bytes == NULL)
4291 return -1;
4292
4293 PDATA_PUSH(self->stack, bytes, -1);
4294 return 0;
4295}
4296
4297static int
4298load_binstring(UnpicklerObject *self)
4299{
4300 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004301 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004302 char *s;
4303
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004304 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004305 return -1;
4306
4307 x = calc_binint(s, 4);
4308 if (x < 0) {
Victor Stinner121aab42011-09-29 23:40:53 +02004309 PyErr_SetString(UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004310 "BINSTRING pickle has negative byte count");
4311 return -1;
4312 }
4313
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004314 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004315 return -1;
4316
4317 /* Convert Python 2.x strings to unicode. */
4318 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4319 if (str == NULL)
4320 return -1;
4321
4322 PDATA_PUSH(self->stack, str, -1);
4323 return 0;
4324}
4325
4326static int
4327load_short_binstring(UnpicklerObject *self)
4328{
4329 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004330 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004331 char *s;
4332
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004333 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004334 return -1;
4335
4336 x = (unsigned char)s[0];
4337
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004338 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004339 return -1;
4340
4341 /* Convert Python 2.x strings to unicode. */
4342 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4343 if (str == NULL)
4344 return -1;
4345
4346 PDATA_PUSH(self->stack, str, -1);
4347 return 0;
4348}
4349
4350static int
4351load_unicode(UnpicklerObject *self)
4352{
4353 PyObject *str;
4354 Py_ssize_t len;
4355 char *s;
4356
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004357 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004358 return -1;
4359 if (len < 1)
4360 return bad_readline();
4361
4362 str = PyUnicode_DecodeRawUnicodeEscape(s, len - 1, NULL);
4363 if (str == NULL)
4364 return -1;
4365
4366 PDATA_PUSH(self->stack, str, -1);
4367 return 0;
4368}
4369
4370static int
4371load_binunicode(UnpicklerObject *self)
4372{
4373 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004374 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004375 char *s;
4376
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004377 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004378 return -1;
4379
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004380 size = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004381 if (size < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004382 PyErr_Format(PyExc_OverflowError,
4383 "BINUNICODE exceeds system's maximum size of %zd bytes",
4384 PY_SSIZE_T_MAX
4385 );
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004386 return -1;
4387 }
4388
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004389
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004390 if (_Unpickler_Read(self, &s, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004391 return -1;
4392
Victor Stinner485fb562010-04-13 11:07:24 +00004393 str = PyUnicode_DecodeUTF8(s, size, "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004394 if (str == NULL)
4395 return -1;
4396
4397 PDATA_PUSH(self->stack, str, -1);
4398 return 0;
4399}
4400
4401static int
4402load_tuple(UnpicklerObject *self)
4403{
4404 PyObject *tuple;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004405 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004406
4407 if ((i = marker(self)) < 0)
4408 return -1;
4409
4410 tuple = Pdata_poptuple(self->stack, i);
4411 if (tuple == NULL)
4412 return -1;
4413 PDATA_PUSH(self->stack, tuple, -1);
4414 return 0;
4415}
4416
4417static int
4418load_counted_tuple(UnpicklerObject *self, int len)
4419{
4420 PyObject *tuple;
4421
4422 tuple = PyTuple_New(len);
4423 if (tuple == NULL)
4424 return -1;
4425
4426 while (--len >= 0) {
4427 PyObject *item;
4428
4429 PDATA_POP(self->stack, item);
4430 if (item == NULL)
4431 return -1;
4432 PyTuple_SET_ITEM(tuple, len, item);
4433 }
4434 PDATA_PUSH(self->stack, tuple, -1);
4435 return 0;
4436}
4437
4438static int
4439load_empty_list(UnpicklerObject *self)
4440{
4441 PyObject *list;
4442
4443 if ((list = PyList_New(0)) == NULL)
4444 return -1;
4445 PDATA_PUSH(self->stack, list, -1);
4446 return 0;
4447}
4448
4449static int
4450load_empty_dict(UnpicklerObject *self)
4451{
4452 PyObject *dict;
4453
4454 if ((dict = PyDict_New()) == NULL)
4455 return -1;
4456 PDATA_PUSH(self->stack, dict, -1);
4457 return 0;
4458}
4459
4460static int
4461load_list(UnpicklerObject *self)
4462{
4463 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004464 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004465
4466 if ((i = marker(self)) < 0)
4467 return -1;
4468
4469 list = Pdata_poplist(self->stack, i);
4470 if (list == NULL)
4471 return -1;
4472 PDATA_PUSH(self->stack, list, -1);
4473 return 0;
4474}
4475
4476static int
4477load_dict(UnpicklerObject *self)
4478{
4479 PyObject *dict, *key, *value;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004480 Py_ssize_t i, j, k;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004481
4482 if ((i = marker(self)) < 0)
4483 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004484 j = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004485
4486 if ((dict = PyDict_New()) == NULL)
4487 return -1;
4488
4489 for (k = i + 1; k < j; k += 2) {
4490 key = self->stack->data[k - 1];
4491 value = self->stack->data[k];
4492 if (PyDict_SetItem(dict, key, value) < 0) {
4493 Py_DECREF(dict);
4494 return -1;
4495 }
4496 }
4497 Pdata_clear(self->stack, i);
4498 PDATA_PUSH(self->stack, dict, -1);
4499 return 0;
4500}
4501
4502static PyObject *
4503instantiate(PyObject *cls, PyObject *args)
4504{
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004505 PyObject *result = NULL;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02004506 _Py_IDENTIFIER(__getinitargs__);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004507 /* Caller must assure args are a tuple. Normally, args come from
4508 Pdata_poptuple which packs objects from the top of the stack
4509 into a newly created tuple. */
4510 assert(PyTuple_Check(args));
4511 if (Py_SIZE(args) > 0 || !PyType_Check(cls) ||
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02004512 _PyObject_HasAttrId(cls, &PyId___getinitargs__)) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004513 result = PyObject_CallObject(cls, args);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004514 }
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004515 else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02004516 _Py_IDENTIFIER(__new__);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02004517
4518 result = _PyObject_CallMethodId(cls, &PyId___new__, "O", cls);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004519 }
4520 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004521}
4522
4523static int
4524load_obj(UnpicklerObject *self)
4525{
4526 PyObject *cls, *args, *obj = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004527 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004528
4529 if ((i = marker(self)) < 0)
4530 return -1;
4531
4532 args = Pdata_poptuple(self->stack, i + 1);
4533 if (args == NULL)
4534 return -1;
4535
4536 PDATA_POP(self->stack, cls);
4537 if (cls) {
4538 obj = instantiate(cls, args);
4539 Py_DECREF(cls);
4540 }
4541 Py_DECREF(args);
4542 if (obj == NULL)
4543 return -1;
4544
4545 PDATA_PUSH(self->stack, obj, -1);
4546 return 0;
4547}
4548
4549static int
4550load_inst(UnpicklerObject *self)
4551{
4552 PyObject *cls = NULL;
4553 PyObject *args = NULL;
4554 PyObject *obj = NULL;
4555 PyObject *module_name;
4556 PyObject *class_name;
4557 Py_ssize_t len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004558 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004559 char *s;
4560
4561 if ((i = marker(self)) < 0)
4562 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004563 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004564 return -1;
4565 if (len < 2)
4566 return bad_readline();
4567
4568 /* Here it is safe to use PyUnicode_DecodeASCII(), even though non-ASCII
4569 identifiers are permitted in Python 3.0, since the INST opcode is only
4570 supported by older protocols on Python 2.x. */
4571 module_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
4572 if (module_name == NULL)
4573 return -1;
4574
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004575 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004576 if (len < 2)
4577 return bad_readline();
4578 class_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004579 if (class_name != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004580 cls = find_class(self, module_name, class_name);
4581 Py_DECREF(class_name);
4582 }
4583 }
4584 Py_DECREF(module_name);
4585
4586 if (cls == NULL)
4587 return -1;
4588
4589 if ((args = Pdata_poptuple(self->stack, i)) != NULL) {
4590 obj = instantiate(cls, args);
4591 Py_DECREF(args);
4592 }
4593 Py_DECREF(cls);
4594
4595 if (obj == NULL)
4596 return -1;
4597
4598 PDATA_PUSH(self->stack, obj, -1);
4599 return 0;
4600}
4601
4602static int
4603load_newobj(UnpicklerObject *self)
4604{
4605 PyObject *args = NULL;
4606 PyObject *clsraw = NULL;
4607 PyTypeObject *cls; /* clsraw cast to its true type */
4608 PyObject *obj;
4609
4610 /* Stack is ... cls argtuple, and we want to call
4611 * cls.__new__(cls, *argtuple).
4612 */
4613 PDATA_POP(self->stack, args);
4614 if (args == NULL)
4615 goto error;
4616 if (!PyTuple_Check(args)) {
4617 PyErr_SetString(UnpicklingError, "NEWOBJ expected an arg " "tuple.");
4618 goto error;
4619 }
4620
4621 PDATA_POP(self->stack, clsraw);
4622 cls = (PyTypeObject *)clsraw;
4623 if (cls == NULL)
4624 goto error;
4625 if (!PyType_Check(cls)) {
4626 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4627 "isn't a type object");
4628 goto error;
4629 }
4630 if (cls->tp_new == NULL) {
4631 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4632 "has NULL tp_new");
4633 goto error;
4634 }
4635
4636 /* Call __new__. */
4637 obj = cls->tp_new(cls, args, NULL);
4638 if (obj == NULL)
4639 goto error;
4640
4641 Py_DECREF(args);
4642 Py_DECREF(clsraw);
4643 PDATA_PUSH(self->stack, obj, -1);
4644 return 0;
4645
4646 error:
4647 Py_XDECREF(args);
4648 Py_XDECREF(clsraw);
4649 return -1;
4650}
4651
4652static int
4653load_global(UnpicklerObject *self)
4654{
4655 PyObject *global = NULL;
4656 PyObject *module_name;
4657 PyObject *global_name;
4658 Py_ssize_t len;
4659 char *s;
4660
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004661 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004662 return -1;
4663 if (len < 2)
4664 return bad_readline();
4665 module_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4666 if (!module_name)
4667 return -1;
4668
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004669 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004670 if (len < 2) {
4671 Py_DECREF(module_name);
4672 return bad_readline();
4673 }
4674 global_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4675 if (global_name) {
4676 global = find_class(self, module_name, global_name);
4677 Py_DECREF(global_name);
4678 }
4679 }
4680 Py_DECREF(module_name);
4681
4682 if (global == NULL)
4683 return -1;
4684 PDATA_PUSH(self->stack, global, -1);
4685 return 0;
4686}
4687
4688static int
4689load_persid(UnpicklerObject *self)
4690{
4691 PyObject *pid;
4692 Py_ssize_t len;
4693 char *s;
4694
4695 if (self->pers_func) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004696 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004697 return -1;
Alexandre Vassalotti896414f2013-11-30 13:52:35 -08004698 if (len < 1)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004699 return bad_readline();
4700
4701 pid = PyBytes_FromStringAndSize(s, len - 1);
4702 if (pid == NULL)
4703 return -1;
4704
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004705 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004706 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004707 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004708 if (pid == NULL)
4709 return -1;
4710
4711 PDATA_PUSH(self->stack, pid, -1);
4712 return 0;
4713 }
4714 else {
4715 PyErr_SetString(UnpicklingError,
4716 "A load persistent id instruction was encountered,\n"
4717 "but no persistent_load function was specified.");
4718 return -1;
4719 }
4720}
4721
4722static int
4723load_binpersid(UnpicklerObject *self)
4724{
4725 PyObject *pid;
4726
4727 if (self->pers_func) {
4728 PDATA_POP(self->stack, pid);
4729 if (pid == NULL)
4730 return -1;
4731
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004732 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004733 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004734 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004735 if (pid == NULL)
4736 return -1;
4737
4738 PDATA_PUSH(self->stack, pid, -1);
4739 return 0;
4740 }
4741 else {
4742 PyErr_SetString(UnpicklingError,
4743 "A load persistent id instruction was encountered,\n"
4744 "but no persistent_load function was specified.");
4745 return -1;
4746 }
4747}
4748
4749static int
4750load_pop(UnpicklerObject *self)
4751{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004752 Py_ssize_t len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004753
4754 /* Note that we split the (pickle.py) stack into two stacks,
4755 * an object stack and a mark stack. We have to be clever and
4756 * pop the right one. We do this by looking at the top of the
Collin Winter8ca69de2009-05-26 16:53:41 +00004757 * mark stack first, and only signalling a stack underflow if
4758 * the object stack is empty and the mark stack doesn't match
4759 * our expectations.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004760 */
Collin Winter8ca69de2009-05-26 16:53:41 +00004761 if (self->num_marks > 0 && self->marks[self->num_marks - 1] == len) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004762 self->num_marks--;
Antoine Pitrou01a15ea2010-01-07 17:57:31 +00004763 } else if (len > 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004764 len--;
4765 Py_DECREF(self->stack->data[len]);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004766 Py_SIZE(self->stack) = len;
Collin Winter8ca69de2009-05-26 16:53:41 +00004767 } else {
4768 return stack_underflow();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004769 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004770 return 0;
4771}
4772
4773static int
4774load_pop_mark(UnpicklerObject *self)
4775{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004776 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004777
4778 if ((i = marker(self)) < 0)
4779 return -1;
4780
4781 Pdata_clear(self->stack, i);
4782
4783 return 0;
4784}
4785
4786static int
4787load_dup(UnpicklerObject *self)
4788{
4789 PyObject *last;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004790 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004791
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004792 if ((len = Py_SIZE(self->stack)) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004793 return stack_underflow();
4794 last = self->stack->data[len - 1];
4795 PDATA_APPEND(self->stack, last, -1);
4796 return 0;
4797}
4798
4799static int
4800load_get(UnpicklerObject *self)
4801{
4802 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004803 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004804 Py_ssize_t len;
4805 char *s;
4806
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004807 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004808 return -1;
4809 if (len < 2)
4810 return bad_readline();
4811
4812 key = PyLong_FromString(s, NULL, 10);
4813 if (key == NULL)
4814 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004815 idx = PyLong_AsSsize_t(key);
4816 if (idx == -1 && PyErr_Occurred()) {
4817 Py_DECREF(key);
4818 return -1;
4819 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004820
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004821 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004822 if (value == NULL) {
4823 if (!PyErr_Occurred())
4824 PyErr_SetObject(PyExc_KeyError, key);
4825 Py_DECREF(key);
4826 return -1;
4827 }
4828 Py_DECREF(key);
4829
4830 PDATA_APPEND(self->stack, value, -1);
4831 return 0;
4832}
4833
4834static int
4835load_binget(UnpicklerObject *self)
4836{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004837 PyObject *value;
4838 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004839 char *s;
4840
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004841 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004842 return -1;
4843
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004844 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004845
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004846 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004847 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004848 PyObject *key = PyLong_FromSsize_t(idx);
Christian Heimes9ee5c372013-07-26 22:45:00 +02004849 if (key != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004850 PyErr_SetObject(PyExc_KeyError, key);
Christian Heimes9ee5c372013-07-26 22:45:00 +02004851 Py_DECREF(key);
4852 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004853 return -1;
4854 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004855
4856 PDATA_APPEND(self->stack, value, -1);
4857 return 0;
4858}
4859
4860static int
4861load_long_binget(UnpicklerObject *self)
4862{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004863 PyObject *value;
4864 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004865 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004866
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004867 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004868 return -1;
4869
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004870 idx = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004871
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004872 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004873 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004874 PyObject *key = PyLong_FromSsize_t(idx);
Christian Heimes9ee5c372013-07-26 22:45:00 +02004875 if (key != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004876 PyErr_SetObject(PyExc_KeyError, key);
Christian Heimes9ee5c372013-07-26 22:45:00 +02004877 Py_DECREF(key);
4878 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004879 return -1;
4880 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004881
4882 PDATA_APPEND(self->stack, value, -1);
4883 return 0;
4884}
4885
4886/* Push an object from the extension registry (EXT[124]). nbytes is
4887 * the number of bytes following the opcode, holding the index (code) value.
4888 */
4889static int
4890load_extension(UnpicklerObject *self, int nbytes)
4891{
4892 char *codebytes; /* the nbytes bytes after the opcode */
4893 long code; /* calc_binint returns long */
4894 PyObject *py_code; /* code as a Python int */
4895 PyObject *obj; /* the object to push */
4896 PyObject *pair; /* (module_name, class_name) */
4897 PyObject *module_name, *class_name;
4898
4899 assert(nbytes == 1 || nbytes == 2 || nbytes == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004900 if (_Unpickler_Read(self, &codebytes, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004901 return -1;
4902 code = calc_binint(codebytes, nbytes);
4903 if (code <= 0) { /* note that 0 is forbidden */
4904 /* Corrupt or hostile pickle. */
4905 PyErr_SetString(UnpicklingError, "EXT specifies code <= 0");
4906 return -1;
4907 }
4908
4909 /* Look for the code in the cache. */
4910 py_code = PyLong_FromLong(code);
4911 if (py_code == NULL)
4912 return -1;
4913 obj = PyDict_GetItem(extension_cache, py_code);
4914 if (obj != NULL) {
4915 /* Bingo. */
4916 Py_DECREF(py_code);
4917 PDATA_APPEND(self->stack, obj, -1);
4918 return 0;
4919 }
4920
4921 /* Look up the (module_name, class_name) pair. */
4922 pair = PyDict_GetItem(inverted_registry, py_code);
4923 if (pair == NULL) {
4924 Py_DECREF(py_code);
4925 PyErr_Format(PyExc_ValueError, "unregistered extension "
4926 "code %ld", code);
4927 return -1;
4928 }
4929 /* Since the extension registry is manipulable via Python code,
4930 * confirm that pair is really a 2-tuple of strings.
4931 */
4932 if (!PyTuple_Check(pair) || PyTuple_Size(pair) != 2 ||
4933 !PyUnicode_Check(module_name = PyTuple_GET_ITEM(pair, 0)) ||
4934 !PyUnicode_Check(class_name = PyTuple_GET_ITEM(pair, 1))) {
4935 Py_DECREF(py_code);
4936 PyErr_Format(PyExc_ValueError, "_inverted_registry[%ld] "
4937 "isn't a 2-tuple of strings", code);
4938 return -1;
4939 }
4940 /* Load the object. */
4941 obj = find_class(self, module_name, class_name);
4942 if (obj == NULL) {
4943 Py_DECREF(py_code);
4944 return -1;
4945 }
4946 /* Cache code -> obj. */
4947 code = PyDict_SetItem(extension_cache, py_code, obj);
4948 Py_DECREF(py_code);
4949 if (code < 0) {
4950 Py_DECREF(obj);
4951 return -1;
4952 }
4953 PDATA_PUSH(self->stack, obj, -1);
4954 return 0;
4955}
4956
4957static int
4958load_put(UnpicklerObject *self)
4959{
4960 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004961 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004962 Py_ssize_t len;
4963 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004964
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004965 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004966 return -1;
4967 if (len < 2)
4968 return bad_readline();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004969 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004970 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004971 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004972
4973 key = PyLong_FromString(s, NULL, 10);
4974 if (key == NULL)
4975 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004976 idx = PyLong_AsSsize_t(key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004977 Py_DECREF(key);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02004978 if (idx < 0) {
4979 if (!PyErr_Occurred())
4980 PyErr_SetString(PyExc_ValueError,
4981 "negative PUT argument");
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004982 return -1;
Antoine Pitrou55549ec2011-08-30 00:27:10 +02004983 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004984
4985 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004986}
4987
4988static int
4989load_binput(UnpicklerObject *self)
4990{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004991 PyObject *value;
4992 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004993 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004994
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004995 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004996 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004997
4998 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004999 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005000 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005001
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005002 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005003
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005004 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005005}
5006
5007static int
5008load_long_binput(UnpicklerObject *self)
5009{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005010 PyObject *value;
5011 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005012 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005013
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005014 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005015 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005016
5017 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005018 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005019 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005020
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005021 idx = calc_binsize(s, 4);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02005022 if (idx < 0) {
5023 PyErr_SetString(PyExc_ValueError,
5024 "negative LONG_BINPUT argument");
5025 return -1;
5026 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005027
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005028 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005029}
5030
5031static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005032do_append(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005033{
5034 PyObject *value;
5035 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005036 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005037
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005038 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005039 if (x > len || x <= 0)
5040 return stack_underflow();
5041 if (len == x) /* nothing to do */
5042 return 0;
5043
5044 list = self->stack->data[x - 1];
5045
5046 if (PyList_Check(list)) {
5047 PyObject *slice;
5048 Py_ssize_t list_len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005049 int ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005050
5051 slice = Pdata_poplist(self->stack, x);
5052 if (!slice)
5053 return -1;
5054 list_len = PyList_GET_SIZE(list);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005055 ret = PyList_SetSlice(list, list_len, list_len, slice);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005056 Py_DECREF(slice);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005057 return ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005058 }
5059 else {
5060 PyObject *append_func;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005061 _Py_IDENTIFIER(append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005062
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005063 append_func = _PyObject_GetAttrId(list, &PyId_append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005064 if (append_func == NULL)
5065 return -1;
5066 for (i = x; i < len; i++) {
5067 PyObject *result;
5068
5069 value = self->stack->data[i];
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005070 result = _Unpickler_FastCall(self, append_func, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005071 if (result == NULL) {
5072 Pdata_clear(self->stack, i + 1);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005073 Py_SIZE(self->stack) = x;
Alexandre Vassalotti637c7c42013-04-20 21:28:21 -07005074 Py_DECREF(append_func);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005075 return -1;
5076 }
5077 Py_DECREF(result);
5078 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005079 Py_SIZE(self->stack) = x;
Alexandre Vassalotti637c7c42013-04-20 21:28:21 -07005080 Py_DECREF(append_func);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005081 }
5082
5083 return 0;
5084}
5085
5086static int
5087load_append(UnpicklerObject *self)
5088{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005089 return do_append(self, Py_SIZE(self->stack) - 1);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005090}
5091
5092static int
5093load_appends(UnpicklerObject *self)
5094{
5095 return do_append(self, marker(self));
5096}
5097
5098static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005099do_setitems(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005100{
5101 PyObject *value, *key;
5102 PyObject *dict;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005103 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005104 int status = 0;
5105
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005106 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005107 if (x > len || x <= 0)
5108 return stack_underflow();
5109 if (len == x) /* nothing to do */
5110 return 0;
Victor Stinner121aab42011-09-29 23:40:53 +02005111 if ((len - x) % 2 != 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005112 /* Currupt or hostile pickle -- we never write one like this. */
5113 PyErr_SetString(UnpicklingError, "odd number of items for SETITEMS");
5114 return -1;
5115 }
5116
5117 /* Here, dict does not actually need to be a PyDict; it could be anything
5118 that supports the __setitem__ attribute. */
5119 dict = self->stack->data[x - 1];
5120
5121 for (i = x + 1; i < len; i += 2) {
5122 key = self->stack->data[i - 1];
5123 value = self->stack->data[i];
5124 if (PyObject_SetItem(dict, key, value) < 0) {
5125 status = -1;
5126 break;
5127 }
5128 }
5129
5130 Pdata_clear(self->stack, x);
5131 return status;
5132}
5133
5134static int
5135load_setitem(UnpicklerObject *self)
5136{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005137 return do_setitems(self, Py_SIZE(self->stack) - 2);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005138}
5139
5140static int
5141load_setitems(UnpicklerObject *self)
5142{
5143 return do_setitems(self, marker(self));
5144}
5145
5146static int
5147load_build(UnpicklerObject *self)
5148{
5149 PyObject *state, *inst, *slotstate;
5150 PyObject *setstate;
5151 int status = 0;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005152 _Py_IDENTIFIER(__setstate__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005153
5154 /* Stack is ... instance, state. We want to leave instance at
5155 * the stack top, possibly mutated via instance.__setstate__(state).
5156 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005157 if (Py_SIZE(self->stack) < 2)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005158 return stack_underflow();
5159
5160 PDATA_POP(self->stack, state);
5161 if (state == NULL)
5162 return -1;
5163
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005164 inst = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005165
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005166 setstate = _PyObject_GetAttrId(inst, &PyId___setstate__);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005167 if (setstate == NULL) {
5168 if (PyErr_ExceptionMatches(PyExc_AttributeError))
5169 PyErr_Clear();
Antoine Pitroud79dc622008-09-05 00:03:33 +00005170 else {
5171 Py_DECREF(state);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005172 return -1;
Antoine Pitroud79dc622008-09-05 00:03:33 +00005173 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005174 }
5175 else {
5176 PyObject *result;
5177
5178 /* The explicit __setstate__ is responsible for everything. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005179 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Antoine Pitroud79dc622008-09-05 00:03:33 +00005180 reference to state first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005181 result = _Unpickler_FastCall(self, setstate, state);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005182 Py_DECREF(setstate);
5183 if (result == NULL)
5184 return -1;
5185 Py_DECREF(result);
5186 return 0;
5187 }
5188
5189 /* A default __setstate__. First see whether state embeds a
5190 * slot state dict too (a proto 2 addition).
5191 */
5192 if (PyTuple_Check(state) && Py_SIZE(state) == 2) {
5193 PyObject *tmp = state;
5194
5195 state = PyTuple_GET_ITEM(tmp, 0);
5196 slotstate = PyTuple_GET_ITEM(tmp, 1);
5197 Py_INCREF(state);
5198 Py_INCREF(slotstate);
5199 Py_DECREF(tmp);
5200 }
5201 else
5202 slotstate = NULL;
5203
5204 /* Set inst.__dict__ from the state dict (if any). */
5205 if (state != Py_None) {
5206 PyObject *dict;
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005207 PyObject *d_key, *d_value;
5208 Py_ssize_t i;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005209 _Py_IDENTIFIER(__dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005210
5211 if (!PyDict_Check(state)) {
5212 PyErr_SetString(UnpicklingError, "state is not a dictionary");
5213 goto error;
5214 }
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005215 dict = _PyObject_GetAttrId(inst, &PyId___dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005216 if (dict == NULL)
5217 goto error;
5218
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005219 i = 0;
5220 while (PyDict_Next(state, &i, &d_key, &d_value)) {
5221 /* normally the keys for instance attributes are
5222 interned. we should try to do that here. */
5223 Py_INCREF(d_key);
5224 if (PyUnicode_CheckExact(d_key))
5225 PyUnicode_InternInPlace(&d_key);
5226 if (PyObject_SetItem(dict, d_key, d_value) < 0) {
5227 Py_DECREF(d_key);
5228 goto error;
5229 }
5230 Py_DECREF(d_key);
5231 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005232 Py_DECREF(dict);
5233 }
5234
5235 /* Also set instance attributes from the slotstate dict (if any). */
5236 if (slotstate != NULL) {
5237 PyObject *d_key, *d_value;
5238 Py_ssize_t i;
5239
5240 if (!PyDict_Check(slotstate)) {
5241 PyErr_SetString(UnpicklingError,
5242 "slot state is not a dictionary");
5243 goto error;
5244 }
5245 i = 0;
5246 while (PyDict_Next(slotstate, &i, &d_key, &d_value)) {
5247 if (PyObject_SetAttr(inst, d_key, d_value) < 0)
5248 goto error;
5249 }
5250 }
5251
5252 if (0) {
5253 error:
5254 status = -1;
5255 }
5256
5257 Py_DECREF(state);
5258 Py_XDECREF(slotstate);
5259 return status;
5260}
5261
5262static int
5263load_mark(UnpicklerObject *self)
5264{
5265
5266 /* Note that we split the (pickle.py) stack into two stacks, an
5267 * object stack and a mark stack. Here we push a mark onto the
5268 * mark stack.
5269 */
5270
5271 if ((self->num_marks + 1) >= self->marks_size) {
5272 size_t alloc;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005273 Py_ssize_t *marks;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005274
5275 /* Use the size_t type to check for overflow. */
5276 alloc = ((size_t)self->num_marks << 1) + 20;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005277 if (alloc > (PY_SSIZE_T_MAX / sizeof(Py_ssize_t)) ||
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00005278 alloc <= ((size_t)self->num_marks + 1)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005279 PyErr_NoMemory();
5280 return -1;
5281 }
5282
5283 if (self->marks == NULL)
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005284 marks = (Py_ssize_t *) PyMem_Malloc(alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005285 else
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005286 marks = (Py_ssize_t *) PyMem_Realloc(self->marks,
5287 alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005288 if (marks == NULL) {
5289 PyErr_NoMemory();
5290 return -1;
5291 }
5292 self->marks = marks;
5293 self->marks_size = (Py_ssize_t)alloc;
5294 }
5295
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005296 self->marks[self->num_marks++] = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005297
5298 return 0;
5299}
5300
5301static int
5302load_reduce(UnpicklerObject *self)
5303{
5304 PyObject *callable = NULL;
5305 PyObject *argtup = NULL;
5306 PyObject *obj = NULL;
5307
5308 PDATA_POP(self->stack, argtup);
5309 if (argtup == NULL)
5310 return -1;
5311 PDATA_POP(self->stack, callable);
5312 if (callable) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005313 obj = PyObject_CallObject(callable, argtup);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005314 Py_DECREF(callable);
5315 }
5316 Py_DECREF(argtup);
5317
5318 if (obj == NULL)
5319 return -1;
5320
5321 PDATA_PUSH(self->stack, obj, -1);
5322 return 0;
5323}
5324
5325/* Just raises an error if we don't know the protocol specified. PROTO
5326 * is the first opcode for protocols >= 2.
5327 */
5328static int
5329load_proto(UnpicklerObject *self)
5330{
5331 char *s;
5332 int i;
5333
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005334 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005335 return -1;
5336
5337 i = (unsigned char)s[0];
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005338 if (i <= HIGHEST_PROTOCOL) {
5339 self->proto = i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005340 return 0;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005341 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005342
5343 PyErr_Format(PyExc_ValueError, "unsupported pickle protocol: %d", i);
5344 return -1;
5345}
5346
5347static PyObject *
5348load(UnpicklerObject *self)
5349{
5350 PyObject *err;
5351 PyObject *value = NULL;
5352 char *s;
5353
5354 self->num_marks = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005355 if (Py_SIZE(self->stack))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005356 Pdata_clear(self->stack, 0);
5357
5358 /* Convenient macros for the dispatch while-switch loop just below. */
5359#define OP(opcode, load_func) \
5360 case opcode: if (load_func(self) < 0) break; continue;
5361
5362#define OP_ARG(opcode, load_func, arg) \
5363 case opcode: if (load_func(self, (arg)) < 0) break; continue;
5364
5365 while (1) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005366 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005367 break;
5368
5369 switch ((enum opcode)s[0]) {
5370 OP(NONE, load_none)
5371 OP(BININT, load_binint)
5372 OP(BININT1, load_binint1)
5373 OP(BININT2, load_binint2)
5374 OP(INT, load_int)
5375 OP(LONG, load_long)
5376 OP_ARG(LONG1, load_counted_long, 1)
5377 OP_ARG(LONG4, load_counted_long, 4)
5378 OP(FLOAT, load_float)
5379 OP(BINFLOAT, load_binfloat)
5380 OP(BINBYTES, load_binbytes)
5381 OP(SHORT_BINBYTES, load_short_binbytes)
5382 OP(BINSTRING, load_binstring)
5383 OP(SHORT_BINSTRING, load_short_binstring)
5384 OP(STRING, load_string)
5385 OP(UNICODE, load_unicode)
5386 OP(BINUNICODE, load_binunicode)
5387 OP_ARG(EMPTY_TUPLE, load_counted_tuple, 0)
5388 OP_ARG(TUPLE1, load_counted_tuple, 1)
5389 OP_ARG(TUPLE2, load_counted_tuple, 2)
5390 OP_ARG(TUPLE3, load_counted_tuple, 3)
5391 OP(TUPLE, load_tuple)
5392 OP(EMPTY_LIST, load_empty_list)
5393 OP(LIST, load_list)
5394 OP(EMPTY_DICT, load_empty_dict)
5395 OP(DICT, load_dict)
5396 OP(OBJ, load_obj)
5397 OP(INST, load_inst)
5398 OP(NEWOBJ, load_newobj)
5399 OP(GLOBAL, load_global)
5400 OP(APPEND, load_append)
5401 OP(APPENDS, load_appends)
5402 OP(BUILD, load_build)
5403 OP(DUP, load_dup)
5404 OP(BINGET, load_binget)
5405 OP(LONG_BINGET, load_long_binget)
5406 OP(GET, load_get)
5407 OP(MARK, load_mark)
5408 OP(BINPUT, load_binput)
5409 OP(LONG_BINPUT, load_long_binput)
5410 OP(PUT, load_put)
5411 OP(POP, load_pop)
5412 OP(POP_MARK, load_pop_mark)
5413 OP(SETITEM, load_setitem)
5414 OP(SETITEMS, load_setitems)
5415 OP(PERSID, load_persid)
5416 OP(BINPERSID, load_binpersid)
5417 OP(REDUCE, load_reduce)
5418 OP(PROTO, load_proto)
5419 OP_ARG(EXT1, load_extension, 1)
5420 OP_ARG(EXT2, load_extension, 2)
5421 OP_ARG(EXT4, load_extension, 4)
5422 OP_ARG(NEWTRUE, load_bool, Py_True)
5423 OP_ARG(NEWFALSE, load_bool, Py_False)
5424
5425 case STOP:
5426 break;
5427
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005428 default:
Benjamin Petersonadde86d2011-09-23 13:41:41 -04005429 if (s[0] == '\0')
5430 PyErr_SetNone(PyExc_EOFError);
5431 else
5432 PyErr_Format(UnpicklingError,
5433 "invalid load key, '%c'.", s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005434 return NULL;
5435 }
5436
5437 break; /* and we are done! */
5438 }
5439
Antoine Pitrou04248a82010-10-12 20:51:21 +00005440 if (_Unpickler_SkipConsumed(self) < 0)
5441 return NULL;
5442
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005443 /* XXX: It is not clear what this is actually for. */
5444 if ((err = PyErr_Occurred())) {
5445 if (err == PyExc_EOFError) {
5446 PyErr_SetNone(PyExc_EOFError);
5447 }
5448 return NULL;
5449 }
5450
5451 PDATA_POP(self->stack, value);
5452 return value;
5453}
5454
5455PyDoc_STRVAR(Unpickler_load_doc,
5456"load() -> object. Load a pickle."
5457"\n"
5458"Read a pickled object representation from the open file object given in\n"
5459"the constructor, and return the reconstituted object hierarchy specified\n"
5460"therein.\n");
5461
5462static PyObject *
5463Unpickler_load(UnpicklerObject *self)
5464{
5465 /* Check whether the Unpickler was initialized correctly. This prevents
5466 segfaulting if a subclass overridden __init__ with a function that does
5467 not call Unpickler.__init__(). Here, we simply ensure that self->read
5468 is not NULL. */
5469 if (self->read == NULL) {
Victor Stinner121aab42011-09-29 23:40:53 +02005470 PyErr_Format(UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005471 "Unpickler.__init__() was not called by %s.__init__()",
5472 Py_TYPE(self)->tp_name);
5473 return NULL;
5474 }
5475
5476 return load(self);
5477}
5478
5479/* The name of find_class() is misleading. In newer pickle protocols, this
5480 function is used for loading any global (i.e., functions), not just
5481 classes. The name is kept only for backward compatibility. */
5482
5483PyDoc_STRVAR(Unpickler_find_class_doc,
5484"find_class(module_name, global_name) -> object.\n"
5485"\n"
5486"Return an object from a specified module, importing the module if\n"
5487"necessary. Subclasses may override this method (e.g. to restrict\n"
5488"unpickling of arbitrary classes and functions).\n"
5489"\n"
5490"This method is called whenever a class or a function object is\n"
5491"needed. Both arguments passed are str objects.\n");
5492
5493static PyObject *
5494Unpickler_find_class(UnpicklerObject *self, PyObject *args)
5495{
5496 PyObject *global;
5497 PyObject *modules_dict;
5498 PyObject *module;
5499 PyObject *module_name, *global_name;
5500
5501 if (!PyArg_UnpackTuple(args, "find_class", 2, 2,
5502 &module_name, &global_name))
5503 return NULL;
5504
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005505 /* Try to map the old names used in Python 2.x to the new ones used in
5506 Python 3.x. We do this only with old pickle protocols and when the
5507 user has not disabled the feature. */
5508 if (self->proto < 3 && self->fix_imports) {
5509 PyObject *key;
5510 PyObject *item;
5511
5512 /* Check if the global (i.e., a function or a class) was renamed
5513 or moved to another module. */
5514 key = PyTuple_Pack(2, module_name, global_name);
5515 if (key == NULL)
5516 return NULL;
5517 item = PyDict_GetItemWithError(name_mapping_2to3, key);
5518 Py_DECREF(key);
5519 if (item) {
5520 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
5521 PyErr_Format(PyExc_RuntimeError,
5522 "_compat_pickle.NAME_MAPPING values should be "
5523 "2-tuples, not %.200s", Py_TYPE(item)->tp_name);
5524 return NULL;
5525 }
5526 module_name = PyTuple_GET_ITEM(item, 0);
5527 global_name = PyTuple_GET_ITEM(item, 1);
5528 if (!PyUnicode_Check(module_name) ||
5529 !PyUnicode_Check(global_name)) {
5530 PyErr_Format(PyExc_RuntimeError,
5531 "_compat_pickle.NAME_MAPPING values should be "
5532 "pairs of str, not (%.200s, %.200s)",
5533 Py_TYPE(module_name)->tp_name,
5534 Py_TYPE(global_name)->tp_name);
5535 return NULL;
5536 }
5537 }
5538 else if (PyErr_Occurred()) {
5539 return NULL;
5540 }
5541
5542 /* Check if the module was renamed. */
5543 item = PyDict_GetItemWithError(import_mapping_2to3, module_name);
5544 if (item) {
5545 if (!PyUnicode_Check(item)) {
5546 PyErr_Format(PyExc_RuntimeError,
5547 "_compat_pickle.IMPORT_MAPPING values should be "
5548 "strings, not %.200s", Py_TYPE(item)->tp_name);
5549 return NULL;
5550 }
5551 module_name = item;
5552 }
5553 else if (PyErr_Occurred()) {
5554 return NULL;
5555 }
5556 }
5557
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005558 modules_dict = PySys_GetObject("modules");
5559 if (modules_dict == NULL)
5560 return NULL;
5561
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005562 module = PyDict_GetItemWithError(modules_dict, module_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005563 if (module == NULL) {
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005564 if (PyErr_Occurred())
5565 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005566 module = PyImport_Import(module_name);
5567 if (module == NULL)
5568 return NULL;
5569 global = PyObject_GetAttr(module, global_name);
5570 Py_DECREF(module);
5571 }
Victor Stinner121aab42011-09-29 23:40:53 +02005572 else {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005573 global = PyObject_GetAttr(module, global_name);
5574 }
5575 return global;
5576}
5577
5578static struct PyMethodDef Unpickler_methods[] = {
5579 {"load", (PyCFunction)Unpickler_load, METH_NOARGS,
5580 Unpickler_load_doc},
5581 {"find_class", (PyCFunction)Unpickler_find_class, METH_VARARGS,
5582 Unpickler_find_class_doc},
5583 {NULL, NULL} /* sentinel */
5584};
5585
5586static void
5587Unpickler_dealloc(UnpicklerObject *self)
5588{
5589 PyObject_GC_UnTrack((PyObject *)self);
5590 Py_XDECREF(self->readline);
5591 Py_XDECREF(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005592 Py_XDECREF(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005593 Py_XDECREF(self->stack);
5594 Py_XDECREF(self->pers_func);
5595 Py_XDECREF(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005596 if (self->buffer.buf != NULL) {
5597 PyBuffer_Release(&self->buffer);
5598 self->buffer.buf = NULL;
5599 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005600
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005601 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005602 PyMem_Free(self->marks);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005603 PyMem_Free(self->input_line);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005604 free(self->encoding);
5605 free(self->errors);
5606
5607 Py_TYPE(self)->tp_free((PyObject *)self);
5608}
5609
5610static int
5611Unpickler_traverse(UnpicklerObject *self, visitproc visit, void *arg)
5612{
5613 Py_VISIT(self->readline);
5614 Py_VISIT(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005615 Py_VISIT(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005616 Py_VISIT(self->stack);
5617 Py_VISIT(self->pers_func);
5618 Py_VISIT(self->arg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005619 return 0;
5620}
5621
5622static int
5623Unpickler_clear(UnpicklerObject *self)
5624{
5625 Py_CLEAR(self->readline);
5626 Py_CLEAR(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005627 Py_CLEAR(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005628 Py_CLEAR(self->stack);
5629 Py_CLEAR(self->pers_func);
5630 Py_CLEAR(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005631 if (self->buffer.buf != NULL) {
5632 PyBuffer_Release(&self->buffer);
5633 self->buffer.buf = NULL;
5634 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005635
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005636 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005637 PyMem_Free(self->marks);
5638 self->marks = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005639 PyMem_Free(self->input_line);
5640 self->input_line = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005641 free(self->encoding);
5642 self->encoding = NULL;
5643 free(self->errors);
5644 self->errors = NULL;
5645
5646 return 0;
5647}
5648
5649PyDoc_STRVAR(Unpickler_doc,
5650"Unpickler(file, *, encoding='ASCII', errors='strict')"
5651"\n"
5652"This takes a binary file for reading a pickle data stream.\n"
5653"\n"
5654"The protocol version of the pickle is detected automatically, so no\n"
5655"proto argument is needed.\n"
5656"\n"
5657"The file-like object must have two methods, a read() method\n"
5658"that takes an integer argument, and a readline() method that\n"
5659"requires no arguments. Both methods should return bytes.\n"
5660"Thus file-like object can be a binary file object opened for\n"
5661"reading, a BytesIO object, or any other custom object that\n"
5662"meets this interface.\n"
5663"\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005664"Optional keyword arguments are *fix_imports*, *encoding* and *errors*,\n"
5665"which are used to control compatiblity support for pickle stream\n"
5666"generated by Python 2.x. If *fix_imports* is True, pickle will try to\n"
5667"map the old Python 2.x names to the new names used in Python 3.x. The\n"
5668"*encoding* and *errors* tell pickle how to decode 8-bit string\n"
5669"instances pickled by Python 2.x; these default to 'ASCII' and\n"
5670"'strict', respectively.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005671
5672static int
5673Unpickler_init(UnpicklerObject *self, PyObject *args, PyObject *kwds)
5674{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005675 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005676 PyObject *file;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005677 PyObject *fix_imports = Py_True;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005678 char *encoding = NULL;
5679 char *errors = NULL;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005680 _Py_IDENTIFIER(persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005681
5682 /* XXX: That is an horrible error message. But, I don't know how to do
5683 better... */
5684 if (Py_SIZE(args) != 1) {
5685 PyErr_Format(PyExc_TypeError,
5686 "%s takes exactly one positional argument (%zd given)",
5687 Py_TYPE(self)->tp_name, Py_SIZE(args));
5688 return -1;
5689 }
5690
5691 /* Arguments parsing needs to be done in the __init__() method to allow
5692 subclasses to define their own __init__() method, which may (or may
5693 not) support Unpickler arguments. However, this means we need to be
5694 extra careful in the other Unpickler methods, since a subclass could
5695 forget to call Unpickler.__init__() thus breaking our internal
5696 invariants. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005697 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:Unpickler", kwlist,
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005698 &file, &fix_imports, &encoding, &errors))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005699 return -1;
5700
5701 /* In case of multiple __init__() calls, clear previous content. */
5702 if (self->read != NULL)
5703 (void)Unpickler_clear(self);
5704
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005705 if (_Unpickler_SetInputStream(self, file) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005706 return -1;
5707
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005708 if (_Unpickler_SetInputEncoding(self, encoding, errors) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005709 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005710
5711 self->fix_imports = PyObject_IsTrue(fix_imports);
5712 if (self->fix_imports == -1)
5713 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005714
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005715 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_load)) {
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005716 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
5717 &PyId_persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005718 if (self->pers_func == NULL)
5719 return -1;
5720 }
5721 else {
5722 self->pers_func = NULL;
5723 }
5724
5725 self->stack = (Pdata *)Pdata_New();
5726 if (self->stack == NULL)
5727 return -1;
5728
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005729 self->memo_size = 32;
5730 self->memo = _Unpickler_NewMemo(self->memo_size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005731 if (self->memo == NULL)
5732 return -1;
5733
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005734 self->arg = NULL;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005735 self->proto = 0;
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005736
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005737 return 0;
5738}
5739
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005740/* Define a proxy object for the Unpickler's internal memo object. This is to
5741 * avoid breaking code like:
5742 * unpickler.memo.clear()
5743 * and
5744 * unpickler.memo = saved_memo
5745 * Is this a good idea? Not really, but we don't want to break code that uses
5746 * it. Note that we don't implement the entire mapping API here. This is
5747 * intentional, as these should be treated as black-box implementation details.
5748 *
5749 * We do, however, have to implement pickling/unpickling support because of
Victor Stinner121aab42011-09-29 23:40:53 +02005750 * real-world code like cvs2svn.
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005751 */
5752
5753typedef struct {
5754 PyObject_HEAD
5755 UnpicklerObject *unpickler;
5756} UnpicklerMemoProxyObject;
5757
5758PyDoc_STRVAR(ump_clear_doc,
5759"memo.clear() -> None. Remove all items from memo.");
5760
5761static PyObject *
5762ump_clear(UnpicklerMemoProxyObject *self)
5763{
5764 _Unpickler_MemoCleanup(self->unpickler);
5765 self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size);
5766 if (self->unpickler->memo == NULL)
5767 return NULL;
5768 Py_RETURN_NONE;
5769}
5770
5771PyDoc_STRVAR(ump_copy_doc,
5772"memo.copy() -> new_memo. Copy the memo to a new object.");
5773
5774static PyObject *
5775ump_copy(UnpicklerMemoProxyObject *self)
5776{
5777 Py_ssize_t i;
5778 PyObject *new_memo = PyDict_New();
5779 if (new_memo == NULL)
5780 return NULL;
5781
5782 for (i = 0; i < self->unpickler->memo_size; i++) {
5783 int status;
5784 PyObject *key, *value;
5785
5786 value = self->unpickler->memo[i];
5787 if (value == NULL)
5788 continue;
5789
5790 key = PyLong_FromSsize_t(i);
5791 if (key == NULL)
5792 goto error;
5793 status = PyDict_SetItem(new_memo, key, value);
5794 Py_DECREF(key);
5795 if (status < 0)
5796 goto error;
5797 }
5798 return new_memo;
5799
5800error:
5801 Py_DECREF(new_memo);
5802 return NULL;
5803}
5804
5805PyDoc_STRVAR(ump_reduce_doc,
5806"memo.__reduce__(). Pickling support.");
5807
5808static PyObject *
5809ump_reduce(UnpicklerMemoProxyObject *self, PyObject *args)
5810{
5811 PyObject *reduce_value;
5812 PyObject *constructor_args;
5813 PyObject *contents = ump_copy(self);
5814 if (contents == NULL)
5815 return NULL;
5816
5817 reduce_value = PyTuple_New(2);
5818 if (reduce_value == NULL) {
5819 Py_DECREF(contents);
5820 return NULL;
5821 }
5822 constructor_args = PyTuple_New(1);
5823 if (constructor_args == NULL) {
5824 Py_DECREF(contents);
5825 Py_DECREF(reduce_value);
5826 return NULL;
5827 }
5828 PyTuple_SET_ITEM(constructor_args, 0, contents);
5829 Py_INCREF((PyObject *)&PyDict_Type);
5830 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
5831 PyTuple_SET_ITEM(reduce_value, 1, constructor_args);
5832 return reduce_value;
5833}
5834
5835static PyMethodDef unpicklerproxy_methods[] = {
5836 {"clear", (PyCFunction)ump_clear, METH_NOARGS, ump_clear_doc},
5837 {"copy", (PyCFunction)ump_copy, METH_NOARGS, ump_copy_doc},
5838 {"__reduce__", (PyCFunction)ump_reduce, METH_VARARGS, ump_reduce_doc},
5839 {NULL, NULL} /* sentinel */
5840};
5841
5842static void
5843UnpicklerMemoProxy_dealloc(UnpicklerMemoProxyObject *self)
5844{
5845 PyObject_GC_UnTrack(self);
5846 Py_XDECREF(self->unpickler);
5847 PyObject_GC_Del((PyObject *)self);
5848}
5849
5850static int
5851UnpicklerMemoProxy_traverse(UnpicklerMemoProxyObject *self,
5852 visitproc visit, void *arg)
5853{
5854 Py_VISIT(self->unpickler);
5855 return 0;
5856}
5857
5858static int
5859UnpicklerMemoProxy_clear(UnpicklerMemoProxyObject *self)
5860{
5861 Py_CLEAR(self->unpickler);
5862 return 0;
5863}
5864
5865static PyTypeObject UnpicklerMemoProxyType = {
5866 PyVarObject_HEAD_INIT(NULL, 0)
5867 "_pickle.UnpicklerMemoProxy", /*tp_name*/
5868 sizeof(UnpicklerMemoProxyObject), /*tp_basicsize*/
5869 0,
5870 (destructor)UnpicklerMemoProxy_dealloc, /* tp_dealloc */
5871 0, /* tp_print */
5872 0, /* tp_getattr */
5873 0, /* tp_setattr */
5874 0, /* tp_compare */
5875 0, /* tp_repr */
5876 0, /* tp_as_number */
5877 0, /* tp_as_sequence */
5878 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00005879 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005880 0, /* tp_call */
5881 0, /* tp_str */
5882 PyObject_GenericGetAttr, /* tp_getattro */
5883 PyObject_GenericSetAttr, /* tp_setattro */
5884 0, /* tp_as_buffer */
5885 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5886 0, /* tp_doc */
5887 (traverseproc)UnpicklerMemoProxy_traverse, /* tp_traverse */
5888 (inquiry)UnpicklerMemoProxy_clear, /* tp_clear */
5889 0, /* tp_richcompare */
5890 0, /* tp_weaklistoffset */
5891 0, /* tp_iter */
5892 0, /* tp_iternext */
5893 unpicklerproxy_methods, /* tp_methods */
5894};
5895
5896static PyObject *
5897UnpicklerMemoProxy_New(UnpicklerObject *unpickler)
5898{
5899 UnpicklerMemoProxyObject *self;
5900
5901 self = PyObject_GC_New(UnpicklerMemoProxyObject,
5902 &UnpicklerMemoProxyType);
5903 if (self == NULL)
5904 return NULL;
5905 Py_INCREF(unpickler);
5906 self->unpickler = unpickler;
5907 PyObject_GC_Track(self);
5908 return (PyObject *)self;
5909}
5910
5911/*****************************************************************************/
5912
5913
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005914static PyObject *
5915Unpickler_get_memo(UnpicklerObject *self)
5916{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005917 return UnpicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005918}
5919
5920static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005921Unpickler_set_memo(UnpicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005922{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005923 PyObject **new_memo;
5924 Py_ssize_t new_memo_size = 0;
5925 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005926
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005927 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005928 PyErr_SetString(PyExc_TypeError,
5929 "attribute deletion is not supported");
5930 return -1;
5931 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005932
5933 if (Py_TYPE(obj) == &UnpicklerMemoProxyType) {
5934 UnpicklerObject *unpickler =
5935 ((UnpicklerMemoProxyObject *)obj)->unpickler;
5936
5937 new_memo_size = unpickler->memo_size;
5938 new_memo = _Unpickler_NewMemo(new_memo_size);
5939 if (new_memo == NULL)
5940 return -1;
5941
5942 for (i = 0; i < new_memo_size; i++) {
5943 Py_XINCREF(unpickler->memo[i]);
5944 new_memo[i] = unpickler->memo[i];
5945 }
5946 }
5947 else if (PyDict_Check(obj)) {
5948 Py_ssize_t i = 0;
5949 PyObject *key, *value;
5950
5951 new_memo_size = PyDict_Size(obj);
5952 new_memo = _Unpickler_NewMemo(new_memo_size);
5953 if (new_memo == NULL)
5954 return -1;
5955
5956 while (PyDict_Next(obj, &i, &key, &value)) {
5957 Py_ssize_t idx;
5958 if (!PyLong_Check(key)) {
5959 PyErr_SetString(PyExc_TypeError,
5960 "memo key must be integers");
5961 goto error;
5962 }
5963 idx = PyLong_AsSsize_t(key);
5964 if (idx == -1 && PyErr_Occurred())
5965 goto error;
Christian Heimesa24b4d22013-07-01 15:17:45 +02005966 if (idx < 0) {
5967 PyErr_SetString(PyExc_ValueError,
Christian Heimes80878792013-07-01 15:23:39 +02005968 "memo key must be positive integers.");
Christian Heimesa24b4d22013-07-01 15:17:45 +02005969 goto error;
5970 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005971 if (_Unpickler_MemoPut(self, idx, value) < 0)
5972 goto error;
5973 }
5974 }
5975 else {
5976 PyErr_Format(PyExc_TypeError,
5977 "'memo' attribute must be an UnpicklerMemoProxy object"
5978 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005979 return -1;
5980 }
5981
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005982 _Unpickler_MemoCleanup(self);
5983 self->memo_size = new_memo_size;
5984 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005985
5986 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005987
5988 error:
5989 if (new_memo_size) {
5990 i = new_memo_size;
5991 while (--i >= 0) {
5992 Py_XDECREF(new_memo[i]);
5993 }
5994 PyMem_FREE(new_memo);
5995 }
5996 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005997}
5998
5999static PyObject *
6000Unpickler_get_persload(UnpicklerObject *self)
6001{
6002 if (self->pers_func == NULL)
6003 PyErr_SetString(PyExc_AttributeError, "persistent_load");
6004 else
6005 Py_INCREF(self->pers_func);
6006 return self->pers_func;
6007}
6008
6009static int
6010Unpickler_set_persload(UnpicklerObject *self, PyObject *value)
6011{
6012 PyObject *tmp;
6013
6014 if (value == NULL) {
6015 PyErr_SetString(PyExc_TypeError,
6016 "attribute deletion is not supported");
6017 return -1;
6018 }
6019 if (!PyCallable_Check(value)) {
6020 PyErr_SetString(PyExc_TypeError,
6021 "persistent_load must be a callable taking "
6022 "one argument");
6023 return -1;
6024 }
6025
6026 tmp = self->pers_func;
6027 Py_INCREF(value);
6028 self->pers_func = value;
6029 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
6030
6031 return 0;
6032}
6033
6034static PyGetSetDef Unpickler_getsets[] = {
6035 {"memo", (getter)Unpickler_get_memo, (setter)Unpickler_set_memo},
6036 {"persistent_load", (getter)Unpickler_get_persload,
6037 (setter)Unpickler_set_persload},
6038 {NULL}
6039};
6040
6041static PyTypeObject Unpickler_Type = {
6042 PyVarObject_HEAD_INIT(NULL, 0)
6043 "_pickle.Unpickler", /*tp_name*/
6044 sizeof(UnpicklerObject), /*tp_basicsize*/
6045 0, /*tp_itemsize*/
6046 (destructor)Unpickler_dealloc, /*tp_dealloc*/
6047 0, /*tp_print*/
6048 0, /*tp_getattr*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006049 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00006050 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006051 0, /*tp_repr*/
6052 0, /*tp_as_number*/
6053 0, /*tp_as_sequence*/
6054 0, /*tp_as_mapping*/
6055 0, /*tp_hash*/
6056 0, /*tp_call*/
6057 0, /*tp_str*/
6058 0, /*tp_getattro*/
6059 0, /*tp_setattro*/
6060 0, /*tp_as_buffer*/
6061 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
6062 Unpickler_doc, /*tp_doc*/
6063 (traverseproc)Unpickler_traverse, /*tp_traverse*/
6064 (inquiry)Unpickler_clear, /*tp_clear*/
6065 0, /*tp_richcompare*/
6066 0, /*tp_weaklistoffset*/
6067 0, /*tp_iter*/
6068 0, /*tp_iternext*/
6069 Unpickler_methods, /*tp_methods*/
6070 0, /*tp_members*/
6071 Unpickler_getsets, /*tp_getset*/
6072 0, /*tp_base*/
6073 0, /*tp_dict*/
6074 0, /*tp_descr_get*/
6075 0, /*tp_descr_set*/
6076 0, /*tp_dictoffset*/
6077 (initproc)Unpickler_init, /*tp_init*/
6078 PyType_GenericAlloc, /*tp_alloc*/
6079 PyType_GenericNew, /*tp_new*/
6080 PyObject_GC_Del, /*tp_free*/
6081 0, /*tp_is_gc*/
6082};
6083
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006084PyDoc_STRVAR(pickle_dump_doc,
6085"dump(obj, file, protocol=None, *, fix_imports=True) -> None\n"
6086"\n"
6087"Write a pickled representation of obj to the open file object file. This\n"
6088"is equivalent to ``Pickler(file, protocol).dump(obj)``, but may be more\n"
6089"efficient.\n"
6090"\n"
6091"The optional protocol argument tells the pickler to use the given protocol;\n"
6092"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
6093"backward-incompatible protocol designed for Python 3.0.\n"
6094"\n"
6095"Specifying a negative protocol version selects the highest protocol version\n"
6096"supported. The higher the protocol used, the more recent the version of\n"
6097"Python needed to read the pickle produced.\n"
6098"\n"
6099"The file argument must have a write() method that accepts a single bytes\n"
6100"argument. It can thus be a file object opened for binary writing, a\n"
6101"io.BytesIO instance, or any other custom object that meets this interface.\n"
6102"\n"
6103"If fix_imports is True and protocol is less than 3, pickle will try to\n"
6104"map the new Python 3.x names to the old module names used in Python 2.x,\n"
6105"so that the pickle data stream is readable with Python 2.x.\n");
6106
6107static PyObject *
6108pickle_dump(PyObject *self, PyObject *args, PyObject *kwds)
6109{
6110 static char *kwlist[] = {"obj", "file", "protocol", "fix_imports", 0};
6111 PyObject *obj;
6112 PyObject *file;
6113 PyObject *proto = NULL;
6114 PyObject *fix_imports = Py_True;
6115 PicklerObject *pickler;
6116
6117 /* fix_imports is a keyword-only argument. */
6118 if (Py_SIZE(args) > 3) {
6119 PyErr_Format(PyExc_TypeError,
6120 "pickle.dump() takes at most 3 positional "
6121 "argument (%zd given)", Py_SIZE(args));
6122 return NULL;
6123 }
6124
6125 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|OO:dump", kwlist,
6126 &obj, &file, &proto, &fix_imports))
6127 return NULL;
6128
6129 pickler = _Pickler_New();
6130 if (pickler == NULL)
6131 return NULL;
6132
6133 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6134 goto error;
6135
6136 if (_Pickler_SetOutputStream(pickler, file) < 0)
6137 goto error;
6138
6139 if (dump(pickler, obj) < 0)
6140 goto error;
6141
6142 if (_Pickler_FlushToFile(pickler) < 0)
6143 goto error;
6144
6145 Py_DECREF(pickler);
6146 Py_RETURN_NONE;
6147
6148 error:
6149 Py_XDECREF(pickler);
6150 return NULL;
6151}
6152
6153PyDoc_STRVAR(pickle_dumps_doc,
6154"dumps(obj, protocol=None, *, fix_imports=True) -> bytes\n"
6155"\n"
6156"Return the pickled representation of the object as a bytes\n"
6157"object, instead of writing it to a file.\n"
6158"\n"
6159"The optional protocol argument tells the pickler to use the given protocol;\n"
6160"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
6161"backward-incompatible protocol designed for Python 3.0.\n"
6162"\n"
6163"Specifying a negative protocol version selects the highest protocol version\n"
6164"supported. The higher the protocol used, the more recent the version of\n"
6165"Python needed to read the pickle produced.\n"
6166"\n"
6167"If fix_imports is True and *protocol* is less than 3, pickle will try to\n"
6168"map the new Python 3.x names to the old module names used in Python 2.x,\n"
6169"so that the pickle data stream is readable with Python 2.x.\n");
6170
6171static PyObject *
6172pickle_dumps(PyObject *self, PyObject *args, PyObject *kwds)
6173{
6174 static char *kwlist[] = {"obj", "protocol", "fix_imports", 0};
6175 PyObject *obj;
6176 PyObject *proto = NULL;
6177 PyObject *result;
6178 PyObject *fix_imports = Py_True;
6179 PicklerObject *pickler;
6180
6181 /* fix_imports is a keyword-only argument. */
6182 if (Py_SIZE(args) > 2) {
6183 PyErr_Format(PyExc_TypeError,
6184 "pickle.dumps() takes at most 2 positional "
6185 "argument (%zd given)", Py_SIZE(args));
6186 return NULL;
6187 }
6188
6189 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:dumps", kwlist,
6190 &obj, &proto, &fix_imports))
6191 return NULL;
6192
6193 pickler = _Pickler_New();
6194 if (pickler == NULL)
6195 return NULL;
6196
6197 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6198 goto error;
6199
6200 if (dump(pickler, obj) < 0)
6201 goto error;
6202
6203 result = _Pickler_GetString(pickler);
6204 Py_DECREF(pickler);
6205 return result;
6206
6207 error:
6208 Py_XDECREF(pickler);
6209 return NULL;
6210}
6211
6212PyDoc_STRVAR(pickle_load_doc,
6213"load(file, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6214"\n"
6215"Read a pickled object representation from the open file object file and\n"
6216"return the reconstituted object hierarchy specified therein. This is\n"
6217"equivalent to ``Unpickler(file).load()``, but may be more efficient.\n"
6218"\n"
6219"The protocol version of the pickle is detected automatically, so no protocol\n"
6220"argument is needed. Bytes past the pickled object's representation are\n"
6221"ignored.\n"
6222"\n"
6223"The argument file must have two methods, a read() method that takes an\n"
6224"integer argument, and a readline() method that requires no arguments. Both\n"
6225"methods should return bytes. Thus *file* can be a binary file object opened\n"
6226"for reading, a BytesIO object, or any other custom object that meets this\n"
6227"interface.\n"
6228"\n"
6229"Optional keyword arguments are fix_imports, encoding and errors,\n"
6230"which are used to control compatiblity support for pickle stream generated\n"
6231"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6232"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6233"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6234"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6235
6236static PyObject *
6237pickle_load(PyObject *self, PyObject *args, PyObject *kwds)
6238{
6239 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
6240 PyObject *file;
6241 PyObject *fix_imports = Py_True;
6242 PyObject *result;
6243 char *encoding = NULL;
6244 char *errors = NULL;
6245 UnpicklerObject *unpickler;
6246
6247 /* fix_imports, encoding and errors are a keyword-only argument. */
6248 if (Py_SIZE(args) != 1) {
6249 PyErr_Format(PyExc_TypeError,
6250 "pickle.load() takes exactly one positional "
6251 "argument (%zd given)", Py_SIZE(args));
6252 return NULL;
6253 }
6254
6255 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:load", kwlist,
6256 &file, &fix_imports, &encoding, &errors))
6257 return NULL;
6258
6259 unpickler = _Unpickler_New();
6260 if (unpickler == NULL)
6261 return NULL;
6262
6263 if (_Unpickler_SetInputStream(unpickler, file) < 0)
6264 goto error;
6265
6266 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6267 goto error;
6268
6269 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6270 if (unpickler->fix_imports == -1)
6271 goto error;
6272
6273 result = load(unpickler);
6274 Py_DECREF(unpickler);
6275 return result;
6276
6277 error:
6278 Py_XDECREF(unpickler);
6279 return NULL;
6280}
6281
6282PyDoc_STRVAR(pickle_loads_doc,
6283"loads(input, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6284"\n"
6285"Read a pickled object hierarchy from a bytes object and return the\n"
6286"reconstituted object hierarchy specified therein\n"
6287"\n"
6288"The protocol version of the pickle is detected automatically, so no protocol\n"
6289"argument is needed. Bytes past the pickled object's representation are\n"
6290"ignored.\n"
6291"\n"
6292"Optional keyword arguments are fix_imports, encoding and errors, which\n"
6293"are used to control compatiblity support for pickle stream generated\n"
6294"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6295"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6296"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6297"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6298
6299static PyObject *
6300pickle_loads(PyObject *self, PyObject *args, PyObject *kwds)
6301{
6302 static char *kwlist[] = {"input", "fix_imports", "encoding", "errors", 0};
6303 PyObject *input;
6304 PyObject *fix_imports = Py_True;
6305 PyObject *result;
6306 char *encoding = NULL;
6307 char *errors = NULL;
6308 UnpicklerObject *unpickler;
6309
6310 /* fix_imports, encoding and errors are a keyword-only argument. */
6311 if (Py_SIZE(args) != 1) {
6312 PyErr_Format(PyExc_TypeError,
6313 "pickle.loads() takes exactly one positional "
6314 "argument (%zd given)", Py_SIZE(args));
6315 return NULL;
6316 }
6317
6318 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:loads", kwlist,
6319 &input, &fix_imports, &encoding, &errors))
6320 return NULL;
6321
6322 unpickler = _Unpickler_New();
6323 if (unpickler == NULL)
6324 return NULL;
6325
6326 if (_Unpickler_SetStringInput(unpickler, input) < 0)
6327 goto error;
6328
6329 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6330 goto error;
6331
6332 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6333 if (unpickler->fix_imports == -1)
6334 goto error;
6335
6336 result = load(unpickler);
6337 Py_DECREF(unpickler);
6338 return result;
6339
6340 error:
6341 Py_XDECREF(unpickler);
6342 return NULL;
6343}
6344
6345
6346static struct PyMethodDef pickle_methods[] = {
6347 {"dump", (PyCFunction)pickle_dump, METH_VARARGS|METH_KEYWORDS,
6348 pickle_dump_doc},
6349 {"dumps", (PyCFunction)pickle_dumps, METH_VARARGS|METH_KEYWORDS,
6350 pickle_dumps_doc},
6351 {"load", (PyCFunction)pickle_load, METH_VARARGS|METH_KEYWORDS,
6352 pickle_load_doc},
6353 {"loads", (PyCFunction)pickle_loads, METH_VARARGS|METH_KEYWORDS,
6354 pickle_loads_doc},
6355 {NULL, NULL} /* sentinel */
6356};
6357
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006358static int
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006359initmodule(void)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006360{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006361 PyObject *copyreg = NULL;
6362 PyObject *compat_pickle = NULL;
6363
6364 /* XXX: We should ensure that the types of the dictionaries imported are
6365 exactly PyDict objects. Otherwise, it is possible to crash the pickle
6366 since we use the PyDict API directly to access these dictionaries. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006367
6368 copyreg = PyImport_ImportModule("copyreg");
6369 if (!copyreg)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006370 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006371 dispatch_table = PyObject_GetAttrString(copyreg, "dispatch_table");
6372 if (!dispatch_table)
6373 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006374 extension_registry = \
6375 PyObject_GetAttrString(copyreg, "_extension_registry");
6376 if (!extension_registry)
6377 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006378 inverted_registry = PyObject_GetAttrString(copyreg, "_inverted_registry");
6379 if (!inverted_registry)
6380 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006381 extension_cache = PyObject_GetAttrString(copyreg, "_extension_cache");
6382 if (!extension_cache)
6383 goto error;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006384 Py_CLEAR(copyreg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006385
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006386 /* Load the 2.x -> 3.x stdlib module mapping tables */
6387 compat_pickle = PyImport_ImportModule("_compat_pickle");
6388 if (!compat_pickle)
6389 goto error;
6390 name_mapping_2to3 = PyObject_GetAttrString(compat_pickle, "NAME_MAPPING");
6391 if (!name_mapping_2to3)
6392 goto error;
6393 if (!PyDict_CheckExact(name_mapping_2to3)) {
6394 PyErr_Format(PyExc_RuntimeError,
6395 "_compat_pickle.NAME_MAPPING should be a dict, not %.200s",
6396 Py_TYPE(name_mapping_2to3)->tp_name);
6397 goto error;
6398 }
6399 import_mapping_2to3 = PyObject_GetAttrString(compat_pickle,
6400 "IMPORT_MAPPING");
6401 if (!import_mapping_2to3)
6402 goto error;
6403 if (!PyDict_CheckExact(import_mapping_2to3)) {
6404 PyErr_Format(PyExc_RuntimeError,
6405 "_compat_pickle.IMPORT_MAPPING should be a dict, "
6406 "not %.200s", Py_TYPE(import_mapping_2to3)->tp_name);
6407 goto error;
6408 }
6409 /* ... and the 3.x -> 2.x mapping tables */
6410 name_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6411 "REVERSE_NAME_MAPPING");
6412 if (!name_mapping_3to2)
6413 goto error;
6414 if (!PyDict_CheckExact(name_mapping_3to2)) {
6415 PyErr_Format(PyExc_RuntimeError,
Ezio Melotti13925002011-03-16 11:05:33 +02006416 "_compat_pickle.REVERSE_NAME_MAPPING should be a dict, "
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006417 "not %.200s", Py_TYPE(name_mapping_3to2)->tp_name);
6418 goto error;
6419 }
6420 import_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6421 "REVERSE_IMPORT_MAPPING");
6422 if (!import_mapping_3to2)
6423 goto error;
6424 if (!PyDict_CheckExact(import_mapping_3to2)) {
6425 PyErr_Format(PyExc_RuntimeError,
6426 "_compat_pickle.REVERSE_IMPORT_MAPPING should be a dict, "
6427 "not %.200s", Py_TYPE(import_mapping_3to2)->tp_name);
6428 goto error;
6429 }
6430 Py_CLEAR(compat_pickle);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006431
6432 empty_tuple = PyTuple_New(0);
6433 if (empty_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006434 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006435 two_tuple = PyTuple_New(2);
6436 if (two_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006437 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006438 /* We use this temp container with no regard to refcounts, or to
6439 * keeping containees alive. Exempt from GC, because we don't
6440 * want anything looking at two_tuple() by magic.
6441 */
6442 PyObject_GC_UnTrack(two_tuple);
6443
6444 return 0;
6445
6446 error:
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006447 Py_CLEAR(copyreg);
6448 Py_CLEAR(dispatch_table);
6449 Py_CLEAR(extension_registry);
6450 Py_CLEAR(inverted_registry);
6451 Py_CLEAR(extension_cache);
6452 Py_CLEAR(compat_pickle);
6453 Py_CLEAR(name_mapping_2to3);
6454 Py_CLEAR(import_mapping_2to3);
6455 Py_CLEAR(name_mapping_3to2);
6456 Py_CLEAR(import_mapping_3to2);
6457 Py_CLEAR(empty_tuple);
6458 Py_CLEAR(two_tuple);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006459 return -1;
6460}
6461
6462static struct PyModuleDef _picklemodule = {
6463 PyModuleDef_HEAD_INIT,
6464 "_pickle",
6465 pickle_module_doc,
6466 -1,
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006467 pickle_methods,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006468 NULL,
6469 NULL,
6470 NULL,
6471 NULL
6472};
6473
6474PyMODINIT_FUNC
6475PyInit__pickle(void)
6476{
6477 PyObject *m;
6478
6479 if (PyType_Ready(&Unpickler_Type) < 0)
6480 return NULL;
6481 if (PyType_Ready(&Pickler_Type) < 0)
6482 return NULL;
6483 if (PyType_Ready(&Pdata_Type) < 0)
6484 return NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006485 if (PyType_Ready(&PicklerMemoProxyType) < 0)
6486 return NULL;
6487 if (PyType_Ready(&UnpicklerMemoProxyType) < 0)
6488 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006489
6490 /* Create the module and add the functions. */
6491 m = PyModule_Create(&_picklemodule);
6492 if (m == NULL)
6493 return NULL;
6494
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006495 Py_INCREF(&Pickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006496 if (PyModule_AddObject(m, "Pickler", (PyObject *)&Pickler_Type) < 0)
6497 return NULL;
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006498 Py_INCREF(&Unpickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006499 if (PyModule_AddObject(m, "Unpickler", (PyObject *)&Unpickler_Type) < 0)
6500 return NULL;
6501
6502 /* Initialize the exceptions. */
6503 PickleError = PyErr_NewException("_pickle.PickleError", NULL, NULL);
6504 if (PickleError == NULL)
6505 return NULL;
6506 PicklingError = \
6507 PyErr_NewException("_pickle.PicklingError", PickleError, NULL);
6508 if (PicklingError == NULL)
6509 return NULL;
6510 UnpicklingError = \
6511 PyErr_NewException("_pickle.UnpicklingError", PickleError, NULL);
6512 if (UnpicklingError == NULL)
6513 return NULL;
6514
6515 if (PyModule_AddObject(m, "PickleError", PickleError) < 0)
6516 return NULL;
6517 if (PyModule_AddObject(m, "PicklingError", PicklingError) < 0)
6518 return NULL;
6519 if (PyModule_AddObject(m, "UnpicklingError", UnpicklingError) < 0)
6520 return NULL;
6521
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006522 if (initmodule() < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006523 return NULL;
6524
6525 return m;
6526}