blob: adb642e2f6248fdd7a13111159ad1c1c0330ff7d [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/*************************************************************************
390 A custom hashtable mapping void* to longs. This is used by the pickler for
391 memoization. Using a custom hashtable rather than PyDict allows us to skip
392 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;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000777 self->max_output_len = WRITE_BUF_SIZE;
778 self->output_len = 0;
Victor Stinner68c8ea22013-07-11 22:56:25 +0200779
780 self->memo = PyMemoTable_New();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000781 self->output_buffer = PyBytes_FromStringAndSize(NULL,
782 self->max_output_len);
Victor Stinner68c8ea22013-07-11 22:56:25 +0200783
784 if (self->memo == NULL || self->output_buffer == NULL) {
785 PyObject_GC_Del(self);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000786 return NULL;
787 }
788 return self;
789}
790
791static int
792_Pickler_SetProtocol(PicklerObject *self, PyObject *proto_obj,
793 PyObject *fix_imports_obj)
794{
795 long proto = 0;
796 int fix_imports;
797
798 if (proto_obj == NULL || proto_obj == Py_None)
799 proto = DEFAULT_PROTOCOL;
800 else {
801 proto = PyLong_AsLong(proto_obj);
802 if (proto == -1 && PyErr_Occurred())
803 return -1;
804 }
805 if (proto < 0)
806 proto = HIGHEST_PROTOCOL;
807 if (proto > HIGHEST_PROTOCOL) {
808 PyErr_Format(PyExc_ValueError, "pickle protocol must be <= %d",
809 HIGHEST_PROTOCOL);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000810 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000811 }
812 fix_imports = PyObject_IsTrue(fix_imports_obj);
813 if (fix_imports == -1)
814 return -1;
Victor Stinner121aab42011-09-29 23:40:53 +0200815
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000816 self->proto = proto;
817 self->bin = proto > 0;
818 self->fix_imports = fix_imports && proto < 3;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000819
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000820 return 0;
821}
822
823/* Returns -1 (with an exception set) on failure, 0 on success. This may
824 be called once on a freshly created Pickler. */
825static int
826_Pickler_SetOutputStream(PicklerObject *self, PyObject *file)
827{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +0200828 _Py_IDENTIFIER(write);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000829 assert(file != NULL);
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +0200830 self->write = _PyObject_GetAttrId(file, &PyId_write);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000831 if (self->write == NULL) {
832 if (PyErr_ExceptionMatches(PyExc_AttributeError))
833 PyErr_SetString(PyExc_TypeError,
834 "file must have a 'write' attribute");
835 return -1;
836 }
837
838 return 0;
839}
840
841/* See documentation for _Pickler_FastCall(). */
842static PyObject *
843_Unpickler_FastCall(UnpicklerObject *self, PyObject *func, PyObject *arg)
844{
845 PyObject *result = NULL;
846
847 ARG_TUP(self, arg);
848 if (self->arg) {
849 result = PyObject_Call(func, self->arg, NULL);
850 FREE_ARG_TUP(self);
851 }
852 return result;
853}
854
855/* Returns the size of the input on success, -1 on failure. This takes its
856 own reference to `input`. */
857static Py_ssize_t
858_Unpickler_SetStringInput(UnpicklerObject *self, PyObject *input)
859{
860 if (self->buffer.buf != NULL)
861 PyBuffer_Release(&self->buffer);
862 if (PyObject_GetBuffer(input, &self->buffer, PyBUF_CONTIG_RO) < 0)
863 return -1;
864 self->input_buffer = self->buffer.buf;
865 self->input_len = self->buffer.len;
866 self->next_read_idx = 0;
Antoine Pitrou04248a82010-10-12 20:51:21 +0000867 self->prefetched_idx = self->input_len;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000868 return self->input_len;
869}
870
Antoine Pitrou04248a82010-10-12 20:51:21 +0000871static int
872_Unpickler_SkipConsumed(UnpicklerObject *self)
873{
874 Py_ssize_t consumed = self->next_read_idx - self->prefetched_idx;
875
876 if (consumed > 0) {
877 PyObject *r;
878 assert(self->peek); /* otherwise we did something wrong */
879 /* This makes an useless copy... */
880 r = PyObject_CallFunction(self->read, "n", consumed);
881 if (r == NULL)
882 return -1;
883 Py_DECREF(r);
884 self->prefetched_idx = self->next_read_idx;
885 }
886 return 0;
887}
888
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000889static const Py_ssize_t READ_WHOLE_LINE = -1;
890
891/* If reading from a file, we need to only pull the bytes we need, since there
892 may be multiple pickle objects arranged contiguously in the same input
893 buffer.
894
895 If `n` is READ_WHOLE_LINE, read a whole line. Otherwise, read up to `n`
896 bytes from the input stream/buffer.
897
898 Update the unpickler's input buffer with the newly-read data. Returns -1 on
899 failure; on success, returns the number of bytes read from the file.
900
901 On success, self->input_len will be 0; this is intentional so that when
902 unpickling from a file, the "we've run out of data" code paths will trigger,
903 causing the Unpickler to go back to the file for more data. Use the returned
904 size to tell you how much data you can process. */
905static Py_ssize_t
906_Unpickler_ReadFromFile(UnpicklerObject *self, Py_ssize_t n)
907{
908 PyObject *data;
Antoine Pitrou04248a82010-10-12 20:51:21 +0000909 Py_ssize_t read_size, prefetched_size = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000910
911 assert(self->read != NULL);
Victor Stinner121aab42011-09-29 23:40:53 +0200912
Antoine Pitrou04248a82010-10-12 20:51:21 +0000913 if (_Unpickler_SkipConsumed(self) < 0)
914 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000915
916 if (n == READ_WHOLE_LINE)
917 data = PyObject_Call(self->readline, empty_tuple, NULL);
918 else {
919 PyObject *len = PyLong_FromSsize_t(n);
920 if (len == NULL)
921 return -1;
922 data = _Unpickler_FastCall(self, self->read, len);
923 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000924 if (data == NULL)
925 return -1;
926
Antoine Pitrou04248a82010-10-12 20:51:21 +0000927 /* Prefetch some data without advancing the file pointer, if possible */
928 if (self->peek) {
929 PyObject *len, *prefetched;
930 len = PyLong_FromSsize_t(PREFETCH);
931 if (len == NULL) {
932 Py_DECREF(data);
933 return -1;
934 }
935 prefetched = _Unpickler_FastCall(self, self->peek, len);
936 if (prefetched == NULL) {
937 if (PyErr_ExceptionMatches(PyExc_NotImplementedError)) {
938 /* peek() is probably not supported by the given file object */
939 PyErr_Clear();
940 Py_CLEAR(self->peek);
941 }
942 else {
943 Py_DECREF(data);
944 return -1;
945 }
946 }
947 else {
948 assert(PyBytes_Check(prefetched));
949 prefetched_size = PyBytes_GET_SIZE(prefetched);
950 PyBytes_ConcatAndDel(&data, prefetched);
951 if (data == NULL)
952 return -1;
953 }
954 }
955
956 read_size = _Unpickler_SetStringInput(self, data) - prefetched_size;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000957 Py_DECREF(data);
Antoine Pitrou04248a82010-10-12 20:51:21 +0000958 self->prefetched_idx = read_size;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000959 return read_size;
960}
961
962/* Read `n` bytes from the unpickler's data source, storing the result in `*s`.
963
964 This should be used for all data reads, rather than accessing the unpickler's
965 input buffer directly. This method deals correctly with reading from input
966 streams, which the input buffer doesn't deal with.
967
968 Note that when reading from a file-like object, self->next_read_idx won't
969 be updated (it should remain at 0 for the entire unpickling process). You
970 should use this function's return value to know how many bytes you can
971 consume.
972
973 Returns -1 (with an exception set) on failure. On success, return the
974 number of chars read. */
975static Py_ssize_t
976_Unpickler_Read(UnpicklerObject *self, char **s, Py_ssize_t n)
977{
Antoine Pitrou04248a82010-10-12 20:51:21 +0000978 Py_ssize_t num_read;
979
Antoine Pitrou04248a82010-10-12 20:51:21 +0000980 if (self->next_read_idx + n <= self->input_len) {
981 *s = self->input_buffer + self->next_read_idx;
982 self->next_read_idx += n;
983 return n;
984 }
985 if (!self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000986 PyErr_Format(PyExc_EOFError, "Ran out of input");
Amaury Forgeot d'Arc3e4e72f2008-11-11 20:05:06 +0000987 return -1;
988 }
Antoine Pitrou04248a82010-10-12 20:51:21 +0000989 num_read = _Unpickler_ReadFromFile(self, n);
990 if (num_read < 0)
991 return -1;
992 if (num_read < n) {
993 PyErr_Format(PyExc_EOFError, "Ran out of input");
994 return -1;
995 }
996 *s = self->input_buffer;
997 self->next_read_idx = n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000998 return n;
999}
1000
1001static Py_ssize_t
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001002_Unpickler_CopyLine(UnpicklerObject *self, char *line, Py_ssize_t len,
1003 char **result)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001004{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001005 char *input_line = PyMem_Realloc(self->input_line, len + 1);
1006 if (input_line == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001007 return -1;
1008
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001009 memcpy(input_line, line, len);
1010 input_line[len] = '\0';
1011 self->input_line = input_line;
1012 *result = self->input_line;
1013 return len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001014}
1015
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001016/* Read a line from the input stream/buffer. If we run off the end of the input
1017 before hitting \n, return the data we found.
1018
1019 Returns the number of chars read, or -1 on failure. */
1020static Py_ssize_t
1021_Unpickler_Readline(UnpicklerObject *self, char **result)
1022{
1023 Py_ssize_t i, num_read;
1024
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001025 for (i = self->next_read_idx; i < self->input_len; i++) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001026 if (self->input_buffer[i] == '\n') {
1027 char *line_start = self->input_buffer + self->next_read_idx;
1028 num_read = i - self->next_read_idx + 1;
1029 self->next_read_idx = i + 1;
1030 return _Unpickler_CopyLine(self, line_start, num_read, result);
1031 }
1032 }
1033 if (self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001034 num_read = _Unpickler_ReadFromFile(self, READ_WHOLE_LINE);
1035 if (num_read < 0)
1036 return -1;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001037 self->next_read_idx = num_read;
Antoine Pitrouf6c7a852011-08-11 21:04:02 +02001038 return _Unpickler_CopyLine(self, self->input_buffer, num_read, result);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001039 }
Victor Stinner121aab42011-09-29 23:40:53 +02001040
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001041 /* If we get here, we've run off the end of the input string. Return the
1042 remaining string and let the caller figure it out. */
1043 *result = self->input_buffer + self->next_read_idx;
1044 num_read = i - self->next_read_idx;
1045 self->next_read_idx = i;
1046 return num_read;
1047}
1048
1049/* Returns -1 (with an exception set) on failure, 0 on success. The memo array
1050 will be modified in place. */
1051static int
1052_Unpickler_ResizeMemoList(UnpicklerObject *self, Py_ssize_t new_size)
1053{
1054 Py_ssize_t i;
1055 PyObject **memo;
1056
1057 assert(new_size > self->memo_size);
1058
1059 memo = PyMem_REALLOC(self->memo, new_size * sizeof(PyObject *));
1060 if (memo == NULL) {
1061 PyErr_NoMemory();
1062 return -1;
1063 }
1064 self->memo = memo;
1065 for (i = self->memo_size; i < new_size; i++)
1066 self->memo[i] = NULL;
1067 self->memo_size = new_size;
1068 return 0;
1069}
1070
1071/* Returns NULL if idx is out of bounds. */
1072static PyObject *
1073_Unpickler_MemoGet(UnpicklerObject *self, Py_ssize_t idx)
1074{
1075 if (idx < 0 || idx >= self->memo_size)
1076 return NULL;
1077
1078 return self->memo[idx];
1079}
1080
1081/* Returns -1 (with an exception set) on failure, 0 on success.
1082 This takes its own reference to `value`. */
1083static int
1084_Unpickler_MemoPut(UnpicklerObject *self, Py_ssize_t idx, PyObject *value)
1085{
1086 PyObject *old_item;
1087
1088 if (idx >= self->memo_size) {
1089 if (_Unpickler_ResizeMemoList(self, idx * 2) < 0)
1090 return -1;
1091 assert(idx < self->memo_size);
1092 }
1093 Py_INCREF(value);
1094 old_item = self->memo[idx];
1095 self->memo[idx] = value;
1096 Py_XDECREF(old_item);
1097 return 0;
1098}
1099
1100static PyObject **
1101_Unpickler_NewMemo(Py_ssize_t new_size)
1102{
1103 PyObject **memo = PyMem_MALLOC(new_size * sizeof(PyObject *));
1104 if (memo == NULL)
1105 return NULL;
1106 memset(memo, 0, new_size * sizeof(PyObject *));
1107 return memo;
1108}
1109
1110/* Free the unpickler's memo, taking care to decref any items left in it. */
1111static void
1112_Unpickler_MemoCleanup(UnpicklerObject *self)
1113{
1114 Py_ssize_t i;
1115 PyObject **memo = self->memo;
1116
1117 if (self->memo == NULL)
1118 return;
1119 self->memo = NULL;
1120 i = self->memo_size;
1121 while (--i >= 0) {
1122 Py_XDECREF(memo[i]);
1123 }
1124 PyMem_FREE(memo);
1125}
1126
1127static UnpicklerObject *
1128_Unpickler_New(void)
1129{
1130 UnpicklerObject *self;
1131
1132 self = PyObject_GC_New(UnpicklerObject, &Unpickler_Type);
1133 if (self == NULL)
1134 return NULL;
1135
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001136 self->arg = NULL;
1137 self->pers_func = NULL;
1138 self->input_buffer = NULL;
1139 self->input_line = NULL;
1140 self->input_len = 0;
1141 self->next_read_idx = 0;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001142 self->prefetched_idx = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001143 self->read = NULL;
1144 self->readline = NULL;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001145 self->peek = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001146 self->encoding = NULL;
1147 self->errors = NULL;
1148 self->marks = NULL;
1149 self->num_marks = 0;
1150 self->marks_size = 0;
1151 self->proto = 0;
1152 self->fix_imports = 0;
Victor Stinner68c8ea22013-07-11 22:56:25 +02001153 memset(&self->buffer, 0, sizeof(Py_buffer));
1154 self->memo_size = 32;
1155 self->memo = _Unpickler_NewMemo(self->memo_size);
1156 self->stack = (Pdata *)Pdata_New();
1157
1158 if (self->memo == NULL || self->stack == NULL) {
1159 Py_DECREF(self);
1160 return NULL;
1161 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001162
1163 return self;
1164}
1165
1166/* Returns -1 (with an exception set) on failure, 0 on success. This may
1167 be called once on a freshly created Pickler. */
1168static int
1169_Unpickler_SetInputStream(UnpicklerObject *self, PyObject *file)
1170{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02001171 _Py_IDENTIFIER(peek);
1172 _Py_IDENTIFIER(read);
1173 _Py_IDENTIFIER(readline);
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02001174
1175 self->peek = _PyObject_GetAttrId(file, &PyId_peek);
Antoine Pitrou04248a82010-10-12 20:51:21 +00001176 if (self->peek == NULL) {
1177 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1178 PyErr_Clear();
1179 else
1180 return -1;
1181 }
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02001182 self->read = _PyObject_GetAttrId(file, &PyId_read);
1183 self->readline = _PyObject_GetAttrId(file, &PyId_readline);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001184 if (self->readline == NULL || self->read == NULL) {
1185 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1186 PyErr_SetString(PyExc_TypeError,
1187 "file must have 'read' and 'readline' attributes");
1188 Py_CLEAR(self->read);
1189 Py_CLEAR(self->readline);
Antoine Pitrou04248a82010-10-12 20:51:21 +00001190 Py_CLEAR(self->peek);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001191 return -1;
1192 }
1193 return 0;
1194}
1195
1196/* Returns -1 (with an exception set) on failure, 0 on success. This may
1197 be called once on a freshly created Pickler. */
1198static int
1199_Unpickler_SetInputEncoding(UnpicklerObject *self,
1200 const char *encoding,
1201 const char *errors)
1202{
1203 if (encoding == NULL)
1204 encoding = "ASCII";
1205 if (errors == NULL)
1206 errors = "strict";
1207
Victor Stinner49fc8ec2013-07-07 23:30:24 +02001208 self->encoding = _PyMem_Strdup(encoding);
1209 self->errors = _PyMem_Strdup(errors);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001210 if (self->encoding == NULL || self->errors == NULL) {
1211 PyErr_NoMemory();
1212 return -1;
1213 }
1214 return 0;
1215}
1216
1217/* Generate a GET opcode for an object stored in the memo. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001218static int
1219memo_get(PicklerObject *self, PyObject *key)
1220{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001221 Py_ssize_t *value;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001222 char pdata[30];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001223 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001224
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001225 value = PyMemoTable_Get(self->memo, key);
1226 if (value == NULL) {
1227 PyErr_SetObject(PyExc_KeyError, key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001228 return -1;
1229 }
1230
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001231 if (!self->bin) {
1232 pdata[0] = GET;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001233 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1,
1234 "%" PY_FORMAT_SIZE_T "d\n", *value);
1235 len = strlen(pdata);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001236 }
1237 else {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001238 if (*value < 256) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001239 pdata[0] = BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001240 pdata[1] = (unsigned char)(*value & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001241 len = 2;
1242 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001243 else if (*value <= 0xffffffffL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001244 pdata[0] = LONG_BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001245 pdata[1] = (unsigned char)(*value & 0xff);
1246 pdata[2] = (unsigned char)((*value >> 8) & 0xff);
1247 pdata[3] = (unsigned char)((*value >> 16) & 0xff);
1248 pdata[4] = (unsigned char)((*value >> 24) & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001249 len = 5;
1250 }
1251 else { /* unlikely */
1252 PyErr_SetString(PicklingError,
1253 "memo id too large for LONG_BINGET");
1254 return -1;
1255 }
1256 }
1257
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001258 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001259 return -1;
1260
1261 return 0;
1262}
1263
1264/* Store an object in the memo, assign it a new unique ID based on the number
1265 of objects currently stored in the memo and generate a PUT opcode. */
1266static int
1267memo_put(PicklerObject *self, PyObject *obj)
1268{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001269 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001270 char pdata[30];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001271 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001272 int status = 0;
1273
1274 if (self->fast)
1275 return 0;
1276
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001277 x = PyMemoTable_Size(self->memo);
1278 if (PyMemoTable_Set(self->memo, obj, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001279 goto error;
1280
1281 if (!self->bin) {
1282 pdata[0] = PUT;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001283 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1,
1284 "%" PY_FORMAT_SIZE_T "d\n", x);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001285 len = strlen(pdata);
1286 }
1287 else {
1288 if (x < 256) {
1289 pdata[0] = BINPUT;
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00001290 pdata[1] = (unsigned char)x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001291 len = 2;
1292 }
1293 else if (x <= 0xffffffffL) {
1294 pdata[0] = LONG_BINPUT;
1295 pdata[1] = (unsigned char)(x & 0xff);
1296 pdata[2] = (unsigned char)((x >> 8) & 0xff);
1297 pdata[3] = (unsigned char)((x >> 16) & 0xff);
1298 pdata[4] = (unsigned char)((x >> 24) & 0xff);
1299 len = 5;
1300 }
1301 else { /* unlikely */
1302 PyErr_SetString(PicklingError,
1303 "memo id too large for LONG_BINPUT");
1304 return -1;
1305 }
1306 }
1307
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001308 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001309 goto error;
1310
1311 if (0) {
1312 error:
1313 status = -1;
1314 }
1315
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001316 return status;
1317}
1318
1319static PyObject *
1320whichmodule(PyObject *global, PyObject *global_name)
1321{
1322 Py_ssize_t i, j;
1323 static PyObject *module_str = NULL;
1324 static PyObject *main_str = NULL;
1325 PyObject *module_name;
1326 PyObject *modules_dict;
1327 PyObject *module;
1328 PyObject *obj;
1329
1330 if (module_str == NULL) {
1331 module_str = PyUnicode_InternFromString("__module__");
1332 if (module_str == NULL)
1333 return NULL;
1334 main_str = PyUnicode_InternFromString("__main__");
1335 if (main_str == NULL)
1336 return NULL;
1337 }
1338
1339 module_name = PyObject_GetAttr(global, module_str);
1340
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00001341 /* In some rare cases (e.g., bound methods of extension types),
1342 __module__ can be None. If it is so, then search sys.modules
1343 for the module of global. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001344 if (module_name == Py_None) {
1345 Py_DECREF(module_name);
1346 goto search;
1347 }
1348
1349 if (module_name) {
1350 return module_name;
1351 }
1352 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1353 PyErr_Clear();
1354 else
1355 return NULL;
1356
1357 search:
1358 modules_dict = PySys_GetObject("modules");
1359 if (modules_dict == NULL)
1360 return NULL;
1361
1362 i = 0;
1363 module_name = NULL;
1364 while ((j = PyDict_Next(modules_dict, &i, &module_name, &module))) {
Mark Dickinson211c6252009-02-01 10:28:51 +00001365 if (PyObject_RichCompareBool(module_name, main_str, Py_EQ) == 1)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001366 continue;
1367
1368 obj = PyObject_GetAttr(module, global_name);
1369 if (obj == NULL) {
1370 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1371 PyErr_Clear();
1372 else
1373 return NULL;
1374 continue;
1375 }
1376
1377 if (obj != global) {
1378 Py_DECREF(obj);
1379 continue;
1380 }
1381
1382 Py_DECREF(obj);
1383 break;
1384 }
1385
1386 /* If no module is found, use __main__. */
1387 if (!j) {
1388 module_name = main_str;
1389 }
1390
1391 Py_INCREF(module_name);
1392 return module_name;
1393}
1394
1395/* fast_save_enter() and fast_save_leave() are guards against recursive
1396 objects when Pickler is used with the "fast mode" (i.e., with object
1397 memoization disabled). If the nesting of a list or dict object exceed
1398 FAST_NESTING_LIMIT, these guards will start keeping an internal
1399 reference to the seen list or dict objects and check whether these objects
1400 are recursive. These are not strictly necessary, since save() has a
1401 hard-coded recursion limit, but they give a nicer error message than the
1402 typical RuntimeError. */
1403static int
1404fast_save_enter(PicklerObject *self, PyObject *obj)
1405{
1406 /* if fast_nesting < 0, we're doing an error exit. */
1407 if (++self->fast_nesting >= FAST_NESTING_LIMIT) {
1408 PyObject *key = NULL;
1409 if (self->fast_memo == NULL) {
1410 self->fast_memo = PyDict_New();
1411 if (self->fast_memo == NULL) {
1412 self->fast_nesting = -1;
1413 return 0;
1414 }
1415 }
1416 key = PyLong_FromVoidPtr(obj);
1417 if (key == NULL)
1418 return 0;
1419 if (PyDict_GetItem(self->fast_memo, key)) {
1420 Py_DECREF(key);
1421 PyErr_Format(PyExc_ValueError,
1422 "fast mode: can't pickle cyclic objects "
1423 "including object type %.200s at %p",
1424 obj->ob_type->tp_name, obj);
1425 self->fast_nesting = -1;
1426 return 0;
1427 }
1428 if (PyDict_SetItem(self->fast_memo, key, Py_None) < 0) {
1429 Py_DECREF(key);
1430 self->fast_nesting = -1;
1431 return 0;
1432 }
1433 Py_DECREF(key);
1434 }
1435 return 1;
1436}
1437
1438static int
1439fast_save_leave(PicklerObject *self, PyObject *obj)
1440{
1441 if (self->fast_nesting-- >= FAST_NESTING_LIMIT) {
1442 PyObject *key = PyLong_FromVoidPtr(obj);
1443 if (key == NULL)
1444 return 0;
1445 if (PyDict_DelItem(self->fast_memo, key) < 0) {
1446 Py_DECREF(key);
1447 return 0;
1448 }
1449 Py_DECREF(key);
1450 }
1451 return 1;
1452}
1453
1454static int
1455save_none(PicklerObject *self, PyObject *obj)
1456{
1457 const char none_op = NONE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001458 if (_Pickler_Write(self, &none_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001459 return -1;
1460
1461 return 0;
1462}
1463
1464static int
1465save_bool(PicklerObject *self, PyObject *obj)
1466{
1467 static const char *buf[2] = { FALSE, TRUE };
1468 const char len[2] = {sizeof(FALSE) - 1, sizeof(TRUE) - 1};
1469 int p = (obj == Py_True);
1470
1471 if (self->proto >= 2) {
1472 const char bool_op = p ? NEWTRUE : NEWFALSE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001473 if (_Pickler_Write(self, &bool_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001474 return -1;
1475 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001476 else if (_Pickler_Write(self, buf[p], len[p]) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001477 return -1;
1478
1479 return 0;
1480}
1481
1482static int
1483save_int(PicklerObject *self, long x)
1484{
1485 char pdata[32];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001486 Py_ssize_t len = 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001487
1488 if (!self->bin
1489#if SIZEOF_LONG > 4
1490 || x > 0x7fffffffL || x < -0x80000000L
1491#endif
1492 ) {
1493 /* Text-mode pickle, or long too big to fit in the 4-byte
1494 * signed BININT format: store as a string.
1495 */
Mark Dickinson8dd05142009-01-20 20:43:58 +00001496 pdata[0] = LONG; /* use LONG for consistency with pickle.py */
1497 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1, "%ldL\n", x);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001498 if (_Pickler_Write(self, pdata, strlen(pdata)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001499 return -1;
1500 }
1501 else {
1502 /* Binary pickle and x fits in a signed 4-byte int. */
1503 pdata[1] = (unsigned char)(x & 0xff);
1504 pdata[2] = (unsigned char)((x >> 8) & 0xff);
1505 pdata[3] = (unsigned char)((x >> 16) & 0xff);
1506 pdata[4] = (unsigned char)((x >> 24) & 0xff);
1507
1508 if ((pdata[4] == 0) && (pdata[3] == 0)) {
1509 if (pdata[2] == 0) {
1510 pdata[0] = BININT1;
1511 len = 2;
1512 }
1513 else {
1514 pdata[0] = BININT2;
1515 len = 3;
1516 }
1517 }
1518 else {
1519 pdata[0] = BININT;
1520 len = 5;
1521 }
1522
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001523 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001524 return -1;
1525 }
1526
1527 return 0;
1528}
1529
1530static int
1531save_long(PicklerObject *self, PyObject *obj)
1532{
1533 PyObject *repr = NULL;
1534 Py_ssize_t size;
1535 long val = PyLong_AsLong(obj);
1536 int status = 0;
1537
1538 const char long_op = LONG;
1539
1540 if (val == -1 && PyErr_Occurred()) {
1541 /* out of range for int pickling */
1542 PyErr_Clear();
1543 }
Antoine Pitroue58bffb2011-08-13 20:40:32 +02001544 else
1545#if SIZEOF_LONG > 4
1546 if (val <= 0x7fffffffL && val >= -0x80000000L)
1547#endif
1548 return save_int(self, val);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001549
1550 if (self->proto >= 2) {
1551 /* Linear-time pickling. */
1552 size_t nbits;
1553 size_t nbytes;
1554 unsigned char *pdata;
1555 char header[5];
1556 int i;
1557 int sign = _PyLong_Sign(obj);
1558
1559 if (sign == 0) {
1560 header[0] = LONG1;
1561 header[1] = 0; /* It's 0 -- an empty bytestring. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001562 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001563 goto error;
1564 return 0;
1565 }
1566 nbits = _PyLong_NumBits(obj);
1567 if (nbits == (size_t)-1 && PyErr_Occurred())
1568 goto error;
1569 /* How many bytes do we need? There are nbits >> 3 full
1570 * bytes of data, and nbits & 7 leftover bits. If there
1571 * are any leftover bits, then we clearly need another
1572 * byte. Wnat's not so obvious is that we *probably*
1573 * need another byte even if there aren't any leftovers:
1574 * the most-significant bit of the most-significant byte
1575 * acts like a sign bit, and it's usually got a sense
1576 * opposite of the one we need. The exception is longs
1577 * of the form -(2**(8*j-1)) for j > 0. Such a long is
1578 * its own 256's-complement, so has the right sign bit
1579 * even without the extra byte. That's a pain to check
1580 * for in advance, though, so we always grab an extra
1581 * byte at the start, and cut it back later if possible.
1582 */
1583 nbytes = (nbits >> 3) + 1;
Antoine Pitroubf6ecf92012-11-24 20:40:21 +01001584 if (nbytes > 0x7fffffffL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001585 PyErr_SetString(PyExc_OverflowError,
1586 "long too large to pickle");
1587 goto error;
1588 }
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001589 repr = PyBytes_FromStringAndSize(NULL, (Py_ssize_t)nbytes);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001590 if (repr == NULL)
1591 goto error;
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001592 pdata = (unsigned char *)PyBytes_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001593 i = _PyLong_AsByteArray((PyLongObject *)obj,
1594 pdata, nbytes,
1595 1 /* little endian */ , 1 /* signed */ );
1596 if (i < 0)
1597 goto error;
1598 /* If the long is negative, this may be a byte more than
1599 * needed. This is so iff the MSB is all redundant sign
1600 * bits.
1601 */
1602 if (sign < 0 &&
Victor Stinner121aab42011-09-29 23:40:53 +02001603 nbytes > 1 &&
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001604 pdata[nbytes - 1] == 0xff &&
1605 (pdata[nbytes - 2] & 0x80) != 0) {
1606 nbytes--;
1607 }
1608
1609 if (nbytes < 256) {
1610 header[0] = LONG1;
1611 header[1] = (unsigned char)nbytes;
1612 size = 2;
1613 }
1614 else {
1615 header[0] = LONG4;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001616 size = (Py_ssize_t) nbytes;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001617 for (i = 1; i < 5; i++) {
1618 header[i] = (unsigned char)(size & 0xff);
1619 size >>= 8;
1620 }
1621 size = 5;
1622 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001623 if (_Pickler_Write(self, header, size) < 0 ||
1624 _Pickler_Write(self, (char *)pdata, (int)nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001625 goto error;
1626 }
1627 else {
1628 char *string;
1629
Mark Dickinson8dd05142009-01-20 20:43:58 +00001630 /* proto < 2: write the repr and newline. This is quadratic-time (in
1631 the number of digits), in both directions. We add a trailing 'L'
1632 to the repr, for compatibility with Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001633
1634 repr = PyObject_Repr(obj);
1635 if (repr == NULL)
1636 goto error;
1637
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00001638 string = _PyUnicode_AsStringAndSize(repr, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001639 if (string == NULL)
1640 goto error;
1641
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001642 if (_Pickler_Write(self, &long_op, 1) < 0 ||
1643 _Pickler_Write(self, string, size) < 0 ||
1644 _Pickler_Write(self, "L\n", 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001645 goto error;
1646 }
1647
1648 if (0) {
1649 error:
1650 status = -1;
1651 }
1652 Py_XDECREF(repr);
1653
1654 return status;
1655}
1656
1657static int
1658save_float(PicklerObject *self, PyObject *obj)
1659{
1660 double x = PyFloat_AS_DOUBLE((PyFloatObject *)obj);
1661
1662 if (self->bin) {
1663 char pdata[9];
1664 pdata[0] = BINFLOAT;
1665 if (_PyFloat_Pack8(x, (unsigned char *)&pdata[1], 0) < 0)
1666 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001667 if (_Pickler_Write(self, pdata, 9) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001668 return -1;
Victor Stinner121aab42011-09-29 23:40:53 +02001669 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001670 else {
Eric Smith0923d1d2009-04-16 20:16:10 +00001671 int result = -1;
1672 char *buf = NULL;
1673 char op = FLOAT;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001674
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001675 if (_Pickler_Write(self, &op, 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001676 goto done;
1677
Mark Dickinson3e09f432009-04-17 08:41:23 +00001678 buf = PyOS_double_to_string(x, 'g', 17, 0, NULL);
Eric Smith0923d1d2009-04-16 20:16:10 +00001679 if (!buf) {
1680 PyErr_NoMemory();
1681 goto done;
1682 }
1683
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001684 if (_Pickler_Write(self, buf, strlen(buf)) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001685 goto done;
1686
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001687 if (_Pickler_Write(self, "\n", 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001688 goto done;
1689
1690 result = 0;
1691done:
1692 PyMem_Free(buf);
1693 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001694 }
1695
1696 return 0;
1697}
1698
1699static int
1700save_bytes(PicklerObject *self, PyObject *obj)
1701{
1702 if (self->proto < 3) {
1703 /* Older pickle protocols do not have an opcode for pickling bytes
1704 objects. Therefore, we need to fake the copy protocol (i.e.,
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001705 the __reduce__ method) to permit bytes object unpickling.
1706
1707 Here we use a hack to be compatible with Python 2. Since in Python
1708 2 'bytes' is just an alias for 'str' (which has different
1709 parameters than the actual bytes object), we use codecs.encode
1710 to create the appropriate 'str' object when unpickled using
1711 Python 2 *and* the appropriate 'bytes' object when unpickled
1712 using Python 3. Again this is a hack and we don't need to do this
1713 with newer protocols. */
1714 static PyObject *codecs_encode = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001715 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001716 int status;
1717
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001718 if (codecs_encode == NULL) {
1719 PyObject *codecs_module = PyImport_ImportModule("codecs");
1720 if (codecs_module == NULL) {
1721 return -1;
1722 }
1723 codecs_encode = PyObject_GetAttrString(codecs_module, "encode");
1724 Py_DECREF(codecs_module);
1725 if (codecs_encode == NULL) {
1726 return -1;
1727 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001728 }
1729
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001730 if (PyBytes_GET_SIZE(obj) == 0) {
1731 reduce_value = Py_BuildValue("(O())", (PyObject*)&PyBytes_Type);
1732 }
1733 else {
1734 static PyObject *latin1 = NULL;
1735 PyObject *unicode_str =
1736 PyUnicode_DecodeLatin1(PyBytes_AS_STRING(obj),
1737 PyBytes_GET_SIZE(obj),
1738 "strict");
1739 if (unicode_str == NULL)
1740 return -1;
1741 if (latin1 == NULL) {
1742 latin1 = PyUnicode_InternFromString("latin1");
Christian Heimes82e6b942013-06-29 21:37:34 +02001743 if (latin1 == NULL) {
1744 Py_DECREF(unicode_str);
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001745 return -1;
Christian Heimes82e6b942013-06-29 21:37:34 +02001746 }
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001747 }
1748 reduce_value = Py_BuildValue("(O(OO))",
1749 codecs_encode, unicode_str, latin1);
1750 Py_DECREF(unicode_str);
1751 }
1752
1753 if (reduce_value == NULL)
1754 return -1;
1755
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001756 /* save_reduce() will memoize the object automatically. */
1757 status = save_reduce(self, reduce_value, obj);
1758 Py_DECREF(reduce_value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001759 return status;
1760 }
1761 else {
1762 Py_ssize_t size;
1763 char header[5];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001764 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001765
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001766 size = PyBytes_GET_SIZE(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001767 if (size < 0)
1768 return -1;
1769
1770 if (size < 256) {
1771 header[0] = SHORT_BINBYTES;
1772 header[1] = (unsigned char)size;
1773 len = 2;
1774 }
1775 else if (size <= 0xffffffffL) {
1776 header[0] = BINBYTES;
1777 header[1] = (unsigned char)(size & 0xff);
1778 header[2] = (unsigned char)((size >> 8) & 0xff);
1779 header[3] = (unsigned char)((size >> 16) & 0xff);
1780 header[4] = (unsigned char)((size >> 24) & 0xff);
1781 len = 5;
1782 }
1783 else {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001784 PyErr_SetString(PyExc_OverflowError,
Serhiy Storchakaf8def282013-02-16 17:29:56 +02001785 "cannot serialize a bytes object larger than 4 GiB");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001786 return -1; /* string too large */
1787 }
1788
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001789 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001790 return -1;
1791
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001792 if (_Pickler_Write(self, PyBytes_AS_STRING(obj), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001793 return -1;
1794
1795 if (memo_put(self, obj) < 0)
1796 return -1;
1797
1798 return 0;
1799 }
1800}
1801
1802/* A copy of PyUnicode_EncodeRawUnicodeEscape() that also translates
1803 backslash and newline characters to \uXXXX escapes. */
1804static PyObject *
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001805raw_unicode_escape(PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001806{
1807 PyObject *repr, *result;
1808 char *p;
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001809 Py_ssize_t i, size, expandsize;
1810 void *data;
1811 unsigned int kind;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001812
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001813 if (PyUnicode_READY(obj))
1814 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001815
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001816 size = PyUnicode_GET_LENGTH(obj);
1817 data = PyUnicode_DATA(obj);
1818 kind = PyUnicode_KIND(obj);
1819 if (kind == PyUnicode_4BYTE_KIND)
1820 expandsize = 10;
1821 else
1822 expandsize = 6;
Victor Stinner121aab42011-09-29 23:40:53 +02001823
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001824 if (size > PY_SSIZE_T_MAX / expandsize)
1825 return PyErr_NoMemory();
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001826 repr = PyByteArray_FromStringAndSize(NULL, expandsize * size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001827 if (repr == NULL)
1828 return NULL;
1829 if (size == 0)
1830 goto done;
1831
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001832 p = PyByteArray_AS_STRING(repr);
1833 for (i=0; i < size; i++) {
1834 Py_UCS4 ch = PyUnicode_READ(kind, data, i);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001835 /* Map 32-bit characters to '\Uxxxxxxxx' */
1836 if (ch >= 0x10000) {
1837 *p++ = '\\';
1838 *p++ = 'U';
Victor Stinnerf5cff562011-10-14 02:13:11 +02001839 *p++ = Py_hexdigits[(ch >> 28) & 0xf];
1840 *p++ = Py_hexdigits[(ch >> 24) & 0xf];
1841 *p++ = Py_hexdigits[(ch >> 20) & 0xf];
1842 *p++ = Py_hexdigits[(ch >> 16) & 0xf];
1843 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
1844 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
1845 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
1846 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001847 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001848 /* Map 16-bit characters to '\uxxxx' */
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001849 else if (ch >= 256 || ch == '\\' || ch == '\n') {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001850 *p++ = '\\';
1851 *p++ = 'u';
Victor Stinnerf5cff562011-10-14 02:13:11 +02001852 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
1853 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
1854 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
1855 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001856 }
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001857 /* Copy everything else as-is */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001858 else
1859 *p++ = (char) ch;
1860 }
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001861 size = p - PyByteArray_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001862
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001863done:
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001864 result = PyBytes_FromStringAndSize(PyByteArray_AS_STRING(repr), size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001865 Py_DECREF(repr);
1866 return result;
1867}
1868
1869static int
Antoine Pitrou299978d2013-04-07 17:38:11 +02001870write_utf8(PicklerObject *self, char *data, Py_ssize_t size)
1871{
1872 char pdata[5];
1873
1874#if SIZEOF_SIZE_T > 4
1875 if (size > 0xffffffffUL) {
1876 /* string too large */
1877 PyErr_SetString(PyExc_OverflowError,
Antoine Pitrou4b7b0f02013-04-07 23:46:52 +02001878 "cannot serialize a string larger than 4GiB");
Antoine Pitrou299978d2013-04-07 17:38:11 +02001879 return -1;
1880 }
1881#endif
1882
1883 pdata[0] = BINUNICODE;
1884 pdata[1] = (unsigned char)(size & 0xff);
1885 pdata[2] = (unsigned char)((size >> 8) & 0xff);
1886 pdata[3] = (unsigned char)((size >> 16) & 0xff);
1887 pdata[4] = (unsigned char)((size >> 24) & 0xff);
1888
1889 if (_Pickler_Write(self, pdata, sizeof(pdata)) < 0)
1890 return -1;
1891
1892 if (_Pickler_Write(self, data, size) < 0)
1893 return -1;
1894
1895 return 0;
1896}
1897
1898static int
1899write_unicode_binary(PicklerObject *self, PyObject *obj)
1900{
1901 PyObject *encoded = NULL;
1902 Py_ssize_t size;
1903 char *data;
1904 int r;
1905
1906 if (PyUnicode_READY(obj))
1907 return -1;
1908
1909 data = PyUnicode_AsUTF8AndSize(obj, &size);
1910 if (data != NULL)
1911 return write_utf8(self, data, size);
1912
1913 /* Issue #8383: for strings with lone surrogates, fallback on the
1914 "surrogatepass" error handler. */
1915 PyErr_Clear();
1916 encoded = PyUnicode_AsEncodedString(obj, "utf-8", "surrogatepass");
1917 if (encoded == NULL)
1918 return -1;
1919
1920 r = write_utf8(self, PyBytes_AS_STRING(encoded),
1921 PyBytes_GET_SIZE(encoded));
1922 Py_DECREF(encoded);
1923 return r;
1924}
1925
1926static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001927save_unicode(PicklerObject *self, PyObject *obj)
1928{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001929 if (self->bin) {
Antoine Pitrou299978d2013-04-07 17:38:11 +02001930 if (write_unicode_binary(self, obj) < 0)
1931 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001932 }
1933 else {
Antoine Pitrou299978d2013-04-07 17:38:11 +02001934 PyObject *encoded;
1935 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001936 const char unicode_op = UNICODE;
1937
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001938 encoded = raw_unicode_escape(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001939 if (encoded == NULL)
Antoine Pitrou299978d2013-04-07 17:38:11 +02001940 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001941
Antoine Pitrou299978d2013-04-07 17:38:11 +02001942 if (_Pickler_Write(self, &unicode_op, 1) < 0) {
1943 Py_DECREF(encoded);
1944 return -1;
1945 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001946
1947 size = PyBytes_GET_SIZE(encoded);
Antoine Pitrou299978d2013-04-07 17:38:11 +02001948 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0) {
1949 Py_DECREF(encoded);
1950 return -1;
1951 }
1952 Py_DECREF(encoded);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001953
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001954 if (_Pickler_Write(self, "\n", 1) < 0)
Antoine Pitrou299978d2013-04-07 17:38:11 +02001955 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001956 }
1957 if (memo_put(self, obj) < 0)
Antoine Pitrou299978d2013-04-07 17:38:11 +02001958 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001959
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001960 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001961}
1962
1963/* A helper for save_tuple. Push the len elements in tuple t on the stack. */
1964static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001965store_tuple_elements(PicklerObject *self, PyObject *t, Py_ssize_t len)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001966{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001967 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001968
1969 assert(PyTuple_Size(t) == len);
1970
1971 for (i = 0; i < len; i++) {
1972 PyObject *element = PyTuple_GET_ITEM(t, i);
1973
1974 if (element == NULL)
1975 return -1;
1976 if (save(self, element, 0) < 0)
1977 return -1;
1978 }
1979
1980 return 0;
1981}
1982
1983/* Tuples are ubiquitous in the pickle protocols, so many techniques are
1984 * used across protocols to minimize the space needed to pickle them.
1985 * Tuples are also the only builtin immutable type that can be recursive
1986 * (a tuple can be reached from itself), and that requires some subtle
1987 * magic so that it works in all cases. IOW, this is a long routine.
1988 */
1989static int
1990save_tuple(PicklerObject *self, PyObject *obj)
1991{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001992 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001993
1994 const char mark_op = MARK;
1995 const char tuple_op = TUPLE;
1996 const char pop_op = POP;
1997 const char pop_mark_op = POP_MARK;
1998 const char len2opcode[] = {EMPTY_TUPLE, TUPLE1, TUPLE2, TUPLE3};
1999
2000 if ((len = PyTuple_Size(obj)) < 0)
2001 return -1;
2002
2003 if (len == 0) {
2004 char pdata[2];
2005
2006 if (self->proto) {
2007 pdata[0] = EMPTY_TUPLE;
2008 len = 1;
2009 }
2010 else {
2011 pdata[0] = MARK;
2012 pdata[1] = TUPLE;
2013 len = 2;
2014 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002015 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002016 return -1;
2017 return 0;
2018 }
2019
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002020 /* The tuple isn't in the memo now. If it shows up there after
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002021 * saving the tuple elements, the tuple must be recursive, in
2022 * which case we'll pop everything we put on the stack, and fetch
2023 * its value from the memo.
2024 */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002025 if (len <= 3 && self->proto >= 2) {
2026 /* Use TUPLE{1,2,3} opcodes. */
2027 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002028 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002029
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002030 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002031 /* pop the len elements */
2032 for (i = 0; i < len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002033 if (_Pickler_Write(self, &pop_op, 1) < 0)
2034 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002035 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002036 if (memo_get(self, obj) < 0)
2037 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002038
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002039 return 0;
2040 }
2041 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002042 if (_Pickler_Write(self, len2opcode + len, 1) < 0)
2043 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002044 }
2045 goto memoize;
2046 }
2047
2048 /* proto < 2 and len > 0, or proto >= 2 and len > 3.
2049 * Generate MARK e1 e2 ... TUPLE
2050 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002051 if (_Pickler_Write(self, &mark_op, 1) < 0)
2052 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002053
2054 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002055 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002056
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002057 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002058 /* pop the stack stuff we pushed */
2059 if (self->bin) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002060 if (_Pickler_Write(self, &pop_mark_op, 1) < 0)
2061 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002062 }
2063 else {
2064 /* Note that we pop one more than len, to remove
2065 * the MARK too.
2066 */
2067 for (i = 0; i <= len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002068 if (_Pickler_Write(self, &pop_op, 1) < 0)
2069 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002070 }
2071 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002072 if (memo_get(self, obj) < 0)
2073 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002074
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002075 return 0;
2076 }
2077 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002078 if (_Pickler_Write(self, &tuple_op, 1) < 0)
2079 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002080 }
2081
2082 memoize:
2083 if (memo_put(self, obj) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002084 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002085
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002086 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002087}
2088
2089/* iter is an iterator giving items, and we batch up chunks of
2090 * MARK item item ... item APPENDS
2091 * opcode sequences. Calling code should have arranged to first create an
2092 * empty list, or list-like object, for the APPENDS to operate on.
2093 * Returns 0 on success, <0 on error.
2094 */
2095static int
2096batch_list(PicklerObject *self, PyObject *iter)
2097{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002098 PyObject *obj = NULL;
2099 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002100 int i, n;
2101
2102 const char mark_op = MARK;
2103 const char append_op = APPEND;
2104 const char appends_op = APPENDS;
2105
2106 assert(iter != NULL);
2107
2108 /* XXX: I think this function could be made faster by avoiding the
2109 iterator interface and fetching objects directly from list using
2110 PyList_GET_ITEM.
2111 */
2112
2113 if (self->proto == 0) {
2114 /* APPENDS isn't available; do one at a time. */
2115 for (;;) {
2116 obj = PyIter_Next(iter);
2117 if (obj == NULL) {
2118 if (PyErr_Occurred())
2119 return -1;
2120 break;
2121 }
2122 i = save(self, obj, 0);
2123 Py_DECREF(obj);
2124 if (i < 0)
2125 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002126 if (_Pickler_Write(self, &append_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002127 return -1;
2128 }
2129 return 0;
2130 }
2131
2132 /* proto > 0: write in batches of BATCHSIZE. */
2133 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002134 /* Get first item */
2135 firstitem = PyIter_Next(iter);
2136 if (firstitem == NULL) {
2137 if (PyErr_Occurred())
2138 goto error;
2139
2140 /* nothing more to add */
2141 break;
2142 }
2143
2144 /* Try to get a second item */
2145 obj = PyIter_Next(iter);
2146 if (obj == NULL) {
2147 if (PyErr_Occurred())
2148 goto error;
2149
2150 /* Only one item to write */
2151 if (save(self, firstitem, 0) < 0)
2152 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002153 if (_Pickler_Write(self, &append_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002154 goto error;
2155 Py_CLEAR(firstitem);
2156 break;
2157 }
2158
2159 /* More than one item to write */
2160
2161 /* Pump out MARK, items, APPENDS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002162 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002163 goto error;
2164
2165 if (save(self, firstitem, 0) < 0)
2166 goto error;
2167 Py_CLEAR(firstitem);
2168 n = 1;
2169
2170 /* Fetch and save up to BATCHSIZE items */
2171 while (obj) {
2172 if (save(self, obj, 0) < 0)
2173 goto error;
2174 Py_CLEAR(obj);
2175 n += 1;
2176
2177 if (n == BATCHSIZE)
2178 break;
2179
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002180 obj = PyIter_Next(iter);
2181 if (obj == NULL) {
2182 if (PyErr_Occurred())
2183 goto error;
2184 break;
2185 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002186 }
2187
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002188 if (_Pickler_Write(self, &appends_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002189 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002190
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002191 } while (n == BATCHSIZE);
2192 return 0;
2193
2194 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002195 Py_XDECREF(firstitem);
2196 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002197 return -1;
2198}
2199
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002200/* This is a variant of batch_list() above, specialized for lists (with no
2201 * support for list subclasses). Like batch_list(), we batch up chunks of
2202 * MARK item item ... item APPENDS
2203 * opcode sequences. Calling code should have arranged to first create an
2204 * empty list, or list-like object, for the APPENDS to operate on.
2205 * Returns 0 on success, -1 on error.
2206 *
2207 * This version is considerably faster than batch_list(), if less general.
2208 *
2209 * Note that this only works for protocols > 0.
2210 */
2211static int
2212batch_list_exact(PicklerObject *self, PyObject *obj)
2213{
2214 PyObject *item = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002215 Py_ssize_t this_batch, total;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002216
2217 const char append_op = APPEND;
2218 const char appends_op = APPENDS;
2219 const char mark_op = MARK;
2220
2221 assert(obj != NULL);
2222 assert(self->proto > 0);
2223 assert(PyList_CheckExact(obj));
2224
2225 if (PyList_GET_SIZE(obj) == 1) {
2226 item = PyList_GET_ITEM(obj, 0);
2227 if (save(self, item, 0) < 0)
2228 return -1;
2229 if (_Pickler_Write(self, &append_op, 1) < 0)
2230 return -1;
2231 return 0;
2232 }
2233
2234 /* Write in batches of BATCHSIZE. */
2235 total = 0;
2236 do {
2237 this_batch = 0;
2238 if (_Pickler_Write(self, &mark_op, 1) < 0)
2239 return -1;
2240 while (total < PyList_GET_SIZE(obj)) {
2241 item = PyList_GET_ITEM(obj, total);
2242 if (save(self, item, 0) < 0)
2243 return -1;
2244 total++;
2245 if (++this_batch == BATCHSIZE)
2246 break;
2247 }
2248 if (_Pickler_Write(self, &appends_op, 1) < 0)
2249 return -1;
2250
2251 } while (total < PyList_GET_SIZE(obj));
2252
2253 return 0;
2254}
2255
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002256static int
2257save_list(PicklerObject *self, PyObject *obj)
2258{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002259 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002260 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002261 int status = 0;
2262
2263 if (self->fast && !fast_save_enter(self, obj))
2264 goto error;
2265
2266 /* Create an empty list. */
2267 if (self->bin) {
2268 header[0] = EMPTY_LIST;
2269 len = 1;
2270 }
2271 else {
2272 header[0] = MARK;
2273 header[1] = LIST;
2274 len = 2;
2275 }
2276
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002277 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002278 goto error;
2279
2280 /* Get list length, and bow out early if empty. */
2281 if ((len = PyList_Size(obj)) < 0)
2282 goto error;
2283
2284 if (memo_put(self, obj) < 0)
2285 goto error;
2286
2287 if (len != 0) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002288 /* Materialize the list elements. */
2289 if (PyList_CheckExact(obj) && self->proto > 0) {
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002290 if (Py_EnterRecursiveCall(" while pickling an object"))
2291 goto error;
2292 status = batch_list_exact(self, obj);
2293 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002294 } else {
2295 PyObject *iter = PyObject_GetIter(obj);
2296 if (iter == NULL)
2297 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002298
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002299 if (Py_EnterRecursiveCall(" while pickling an object")) {
2300 Py_DECREF(iter);
2301 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002302 }
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002303 status = batch_list(self, iter);
2304 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002305 Py_DECREF(iter);
2306 }
2307 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002308 if (0) {
2309 error:
2310 status = -1;
2311 }
2312
2313 if (self->fast && !fast_save_leave(self, obj))
2314 status = -1;
2315
2316 return status;
2317}
2318
2319/* iter is an iterator giving (key, value) pairs, and we batch up chunks of
2320 * MARK key value ... key value SETITEMS
2321 * opcode sequences. Calling code should have arranged to first create an
2322 * empty dict, or dict-like object, for the SETITEMS to operate on.
2323 * Returns 0 on success, <0 on error.
2324 *
2325 * This is very much like batch_list(). The difference between saving
2326 * elements directly, and picking apart two-tuples, is so long-winded at
2327 * the C level, though, that attempts to combine these routines were too
2328 * ugly to bear.
2329 */
2330static int
2331batch_dict(PicklerObject *self, PyObject *iter)
2332{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002333 PyObject *obj = NULL;
2334 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002335 int i, n;
2336
2337 const char mark_op = MARK;
2338 const char setitem_op = SETITEM;
2339 const char setitems_op = SETITEMS;
2340
2341 assert(iter != NULL);
2342
2343 if (self->proto == 0) {
2344 /* SETITEMS isn't available; do one at a time. */
2345 for (;;) {
2346 obj = PyIter_Next(iter);
2347 if (obj == NULL) {
2348 if (PyErr_Occurred())
2349 return -1;
2350 break;
2351 }
2352 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2353 PyErr_SetString(PyExc_TypeError, "dict items "
2354 "iterator must return 2-tuples");
2355 return -1;
2356 }
2357 i = save(self, PyTuple_GET_ITEM(obj, 0), 0);
2358 if (i >= 0)
2359 i = save(self, PyTuple_GET_ITEM(obj, 1), 0);
2360 Py_DECREF(obj);
2361 if (i < 0)
2362 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002363 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002364 return -1;
2365 }
2366 return 0;
2367 }
2368
2369 /* proto > 0: write in batches of BATCHSIZE. */
2370 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002371 /* Get first item */
2372 firstitem = PyIter_Next(iter);
2373 if (firstitem == NULL) {
2374 if (PyErr_Occurred())
2375 goto error;
2376
2377 /* nothing more to add */
2378 break;
2379 }
2380 if (!PyTuple_Check(firstitem) || PyTuple_Size(firstitem) != 2) {
2381 PyErr_SetString(PyExc_TypeError, "dict items "
2382 "iterator must return 2-tuples");
2383 goto error;
2384 }
2385
2386 /* Try to get a second item */
2387 obj = PyIter_Next(iter);
2388 if (obj == NULL) {
2389 if (PyErr_Occurred())
2390 goto error;
2391
2392 /* Only one item to write */
2393 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2394 goto error;
2395 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2396 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002397 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002398 goto error;
2399 Py_CLEAR(firstitem);
2400 break;
2401 }
2402
2403 /* More than one item to write */
2404
2405 /* Pump out MARK, items, SETITEMS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002406 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002407 goto error;
2408
2409 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2410 goto error;
2411 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2412 goto error;
2413 Py_CLEAR(firstitem);
2414 n = 1;
2415
2416 /* Fetch and save up to BATCHSIZE items */
2417 while (obj) {
2418 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2419 PyErr_SetString(PyExc_TypeError, "dict items "
2420 "iterator must return 2-tuples");
2421 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002422 }
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002423 if (save(self, PyTuple_GET_ITEM(obj, 0), 0) < 0 ||
2424 save(self, PyTuple_GET_ITEM(obj, 1), 0) < 0)
2425 goto error;
2426 Py_CLEAR(obj);
2427 n += 1;
2428
2429 if (n == BATCHSIZE)
2430 break;
2431
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002432 obj = PyIter_Next(iter);
2433 if (obj == NULL) {
2434 if (PyErr_Occurred())
2435 goto error;
2436 break;
2437 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002438 }
2439
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002440 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002441 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002442
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002443 } while (n == BATCHSIZE);
2444 return 0;
2445
2446 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002447 Py_XDECREF(firstitem);
2448 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002449 return -1;
2450}
2451
Collin Winter5c9b02d2009-05-25 05:43:30 +00002452/* This is a variant of batch_dict() above that specializes for dicts, with no
2453 * support for dict subclasses. Like batch_dict(), we batch up chunks of
2454 * MARK key value ... key value SETITEMS
2455 * opcode sequences. Calling code should have arranged to first create an
2456 * empty dict, or dict-like object, for the SETITEMS to operate on.
2457 * Returns 0 on success, -1 on error.
2458 *
2459 * Note that this currently doesn't work for protocol 0.
2460 */
2461static int
2462batch_dict_exact(PicklerObject *self, PyObject *obj)
2463{
2464 PyObject *key = NULL, *value = NULL;
2465 int i;
2466 Py_ssize_t dict_size, ppos = 0;
2467
Alexandre Vassalottif70b1292009-05-25 18:00:52 +00002468 const char mark_op = MARK;
2469 const char setitem_op = SETITEM;
2470 const char setitems_op = SETITEMS;
Collin Winter5c9b02d2009-05-25 05:43:30 +00002471
2472 assert(obj != NULL);
2473 assert(self->proto > 0);
2474
2475 dict_size = PyDict_Size(obj);
2476
2477 /* Special-case len(d) == 1 to save space. */
2478 if (dict_size == 1) {
2479 PyDict_Next(obj, &ppos, &key, &value);
2480 if (save(self, key, 0) < 0)
2481 return -1;
2482 if (save(self, value, 0) < 0)
2483 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002484 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002485 return -1;
2486 return 0;
2487 }
2488
2489 /* Write in batches of BATCHSIZE. */
2490 do {
2491 i = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002492 if (_Pickler_Write(self, &mark_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002493 return -1;
2494 while (PyDict_Next(obj, &ppos, &key, &value)) {
2495 if (save(self, key, 0) < 0)
2496 return -1;
2497 if (save(self, value, 0) < 0)
2498 return -1;
2499 if (++i == BATCHSIZE)
2500 break;
2501 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002502 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002503 return -1;
2504 if (PyDict_Size(obj) != dict_size) {
2505 PyErr_Format(
2506 PyExc_RuntimeError,
2507 "dictionary changed size during iteration");
2508 return -1;
2509 }
2510
2511 } while (i == BATCHSIZE);
2512 return 0;
2513}
2514
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002515static int
2516save_dict(PicklerObject *self, PyObject *obj)
2517{
2518 PyObject *items, *iter;
2519 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002520 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002521 int status = 0;
2522
2523 if (self->fast && !fast_save_enter(self, obj))
2524 goto error;
2525
2526 /* Create an empty dict. */
2527 if (self->bin) {
2528 header[0] = EMPTY_DICT;
2529 len = 1;
2530 }
2531 else {
2532 header[0] = MARK;
2533 header[1] = DICT;
2534 len = 2;
2535 }
2536
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002537 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002538 goto error;
2539
2540 /* Get dict size, and bow out early if empty. */
2541 if ((len = PyDict_Size(obj)) < 0)
2542 goto error;
2543
2544 if (memo_put(self, obj) < 0)
2545 goto error;
2546
2547 if (len != 0) {
2548 /* Save the dict items. */
Collin Winter5c9b02d2009-05-25 05:43:30 +00002549 if (PyDict_CheckExact(obj) && self->proto > 0) {
2550 /* We can take certain shortcuts if we know this is a dict and
2551 not a dict subclass. */
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002552 if (Py_EnterRecursiveCall(" while pickling an object"))
2553 goto error;
2554 status = batch_dict_exact(self, obj);
2555 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002556 } else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02002557 _Py_IDENTIFIER(items);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02002558
2559 items = _PyObject_CallMethodId(obj, &PyId_items, "()");
Collin Winter5c9b02d2009-05-25 05:43:30 +00002560 if (items == NULL)
2561 goto error;
2562 iter = PyObject_GetIter(items);
2563 Py_DECREF(items);
2564 if (iter == NULL)
2565 goto error;
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002566 if (Py_EnterRecursiveCall(" while pickling an object")) {
2567 Py_DECREF(iter);
2568 goto error;
2569 }
Collin Winter5c9b02d2009-05-25 05:43:30 +00002570 status = batch_dict(self, iter);
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002571 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002572 Py_DECREF(iter);
2573 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002574 }
2575
2576 if (0) {
2577 error:
2578 status = -1;
2579 }
2580
2581 if (self->fast && !fast_save_leave(self, obj))
2582 status = -1;
2583
2584 return status;
2585}
2586
2587static int
2588save_global(PicklerObject *self, PyObject *obj, PyObject *name)
2589{
2590 static PyObject *name_str = NULL;
2591 PyObject *global_name = NULL;
2592 PyObject *module_name = NULL;
2593 PyObject *module = NULL;
2594 PyObject *cls;
2595 int status = 0;
2596
2597 const char global_op = GLOBAL;
2598
2599 if (name_str == NULL) {
2600 name_str = PyUnicode_InternFromString("__name__");
2601 if (name_str == NULL)
2602 goto error;
2603 }
2604
2605 if (name) {
2606 global_name = name;
2607 Py_INCREF(global_name);
2608 }
2609 else {
2610 global_name = PyObject_GetAttr(obj, name_str);
2611 if (global_name == NULL)
2612 goto error;
2613 }
2614
2615 module_name = whichmodule(obj, global_name);
2616 if (module_name == NULL)
2617 goto error;
2618
2619 /* XXX: Change to use the import C API directly with level=0 to disallow
2620 relative imports.
2621
2622 XXX: PyImport_ImportModuleLevel could be used. However, this bypasses
2623 builtins.__import__. Therefore, _pickle, unlike pickle.py, will ignore
2624 custom import functions (IMHO, this would be a nice security
2625 feature). The import C API would need to be extended to support the
2626 extra parameters of __import__ to fix that. */
2627 module = PyImport_Import(module_name);
2628 if (module == NULL) {
2629 PyErr_Format(PicklingError,
2630 "Can't pickle %R: import of module %R failed",
2631 obj, module_name);
2632 goto error;
2633 }
2634 cls = PyObject_GetAttr(module, global_name);
2635 if (cls == NULL) {
2636 PyErr_Format(PicklingError,
2637 "Can't pickle %R: attribute lookup %S.%S failed",
2638 obj, module_name, global_name);
2639 goto error;
2640 }
2641 if (cls != obj) {
2642 Py_DECREF(cls);
2643 PyErr_Format(PicklingError,
2644 "Can't pickle %R: it's not the same object as %S.%S",
2645 obj, module_name, global_name);
2646 goto error;
2647 }
2648 Py_DECREF(cls);
2649
2650 if (self->proto >= 2) {
2651 /* See whether this is in the extension registry, and if
2652 * so generate an EXT opcode.
2653 */
2654 PyObject *code_obj; /* extension code as Python object */
2655 long code; /* extension code as C value */
2656 char pdata[5];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002657 Py_ssize_t n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002658
2659 PyTuple_SET_ITEM(two_tuple, 0, module_name);
2660 PyTuple_SET_ITEM(two_tuple, 1, global_name);
2661 code_obj = PyDict_GetItem(extension_registry, two_tuple);
2662 /* The object is not registered in the extension registry.
2663 This is the most likely code path. */
2664 if (code_obj == NULL)
2665 goto gen_global;
2666
2667 /* XXX: pickle.py doesn't check neither the type, nor the range
2668 of the value returned by the extension_registry. It should for
2669 consistency. */
2670
2671 /* Verify code_obj has the right type and value. */
2672 if (!PyLong_Check(code_obj)) {
2673 PyErr_Format(PicklingError,
2674 "Can't pickle %R: extension code %R isn't an integer",
2675 obj, code_obj);
2676 goto error;
2677 }
2678 code = PyLong_AS_LONG(code_obj);
2679 if (code <= 0 || code > 0x7fffffffL) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002680 if (!PyErr_Occurred())
2681 PyErr_Format(PicklingError,
2682 "Can't pickle %R: extension code %ld is out of range",
2683 obj, code);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002684 goto error;
2685 }
2686
2687 /* Generate an EXT opcode. */
2688 if (code <= 0xff) {
2689 pdata[0] = EXT1;
2690 pdata[1] = (unsigned char)code;
2691 n = 2;
2692 }
2693 else if (code <= 0xffff) {
2694 pdata[0] = EXT2;
2695 pdata[1] = (unsigned char)(code & 0xff);
2696 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2697 n = 3;
2698 }
2699 else {
2700 pdata[0] = EXT4;
2701 pdata[1] = (unsigned char)(code & 0xff);
2702 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2703 pdata[3] = (unsigned char)((code >> 16) & 0xff);
2704 pdata[4] = (unsigned char)((code >> 24) & 0xff);
2705 n = 5;
2706 }
2707
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002708 if (_Pickler_Write(self, pdata, n) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002709 goto error;
2710 }
2711 else {
2712 /* Generate a normal global opcode if we are using a pickle
2713 protocol <= 2, or if the object is not registered in the
2714 extension registry. */
2715 PyObject *encoded;
2716 PyObject *(*unicode_encoder)(PyObject *);
2717
2718 gen_global:
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002719 if (_Pickler_Write(self, &global_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002720 goto error;
2721
2722 /* Since Python 3.0 now supports non-ASCII identifiers, we encode both
2723 the module name and the global name using UTF-8. We do so only when
2724 we are using the pickle protocol newer than version 3. This is to
2725 ensure compatibility with older Unpickler running on Python 2.x. */
2726 if (self->proto >= 3) {
2727 unicode_encoder = PyUnicode_AsUTF8String;
2728 }
2729 else {
2730 unicode_encoder = PyUnicode_AsASCIIString;
2731 }
2732
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00002733 /* For protocol < 3 and if the user didn't request against doing so,
2734 we convert module names to the old 2.x module names. */
2735 if (self->fix_imports) {
2736 PyObject *key;
2737 PyObject *item;
2738
2739 key = PyTuple_Pack(2, module_name, global_name);
2740 if (key == NULL)
2741 goto error;
2742 item = PyDict_GetItemWithError(name_mapping_3to2, key);
2743 Py_DECREF(key);
2744 if (item) {
2745 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
2746 PyErr_Format(PyExc_RuntimeError,
2747 "_compat_pickle.REVERSE_NAME_MAPPING values "
2748 "should be 2-tuples, not %.200s",
2749 Py_TYPE(item)->tp_name);
2750 goto error;
2751 }
2752 Py_CLEAR(module_name);
2753 Py_CLEAR(global_name);
2754 module_name = PyTuple_GET_ITEM(item, 0);
2755 global_name = PyTuple_GET_ITEM(item, 1);
2756 if (!PyUnicode_Check(module_name) ||
2757 !PyUnicode_Check(global_name)) {
2758 PyErr_Format(PyExc_RuntimeError,
2759 "_compat_pickle.REVERSE_NAME_MAPPING values "
2760 "should be pairs of str, not (%.200s, %.200s)",
2761 Py_TYPE(module_name)->tp_name,
2762 Py_TYPE(global_name)->tp_name);
2763 goto error;
2764 }
2765 Py_INCREF(module_name);
2766 Py_INCREF(global_name);
2767 }
2768 else if (PyErr_Occurred()) {
2769 goto error;
2770 }
2771
2772 item = PyDict_GetItemWithError(import_mapping_3to2, module_name);
2773 if (item) {
2774 if (!PyUnicode_Check(item)) {
2775 PyErr_Format(PyExc_RuntimeError,
2776 "_compat_pickle.REVERSE_IMPORT_MAPPING values "
2777 "should be strings, not %.200s",
2778 Py_TYPE(item)->tp_name);
2779 goto error;
2780 }
2781 Py_CLEAR(module_name);
2782 module_name = item;
2783 Py_INCREF(module_name);
2784 }
2785 else if (PyErr_Occurred()) {
2786 goto error;
2787 }
2788 }
2789
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002790 /* Save the name of the module. */
2791 encoded = unicode_encoder(module_name);
2792 if (encoded == NULL) {
2793 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2794 PyErr_Format(PicklingError,
2795 "can't pickle module identifier '%S' using "
2796 "pickle protocol %i", module_name, self->proto);
2797 goto error;
2798 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002799 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002800 PyBytes_GET_SIZE(encoded)) < 0) {
2801 Py_DECREF(encoded);
2802 goto error;
2803 }
2804 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002805 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002806 goto error;
2807
2808 /* Save the name of the module. */
2809 encoded = unicode_encoder(global_name);
2810 if (encoded == NULL) {
2811 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2812 PyErr_Format(PicklingError,
2813 "can't pickle global identifier '%S' using "
2814 "pickle protocol %i", global_name, self->proto);
2815 goto error;
2816 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002817 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002818 PyBytes_GET_SIZE(encoded)) < 0) {
2819 Py_DECREF(encoded);
2820 goto error;
2821 }
2822 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002823 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002824 goto error;
2825
2826 /* Memoize the object. */
2827 if (memo_put(self, obj) < 0)
2828 goto error;
2829 }
2830
2831 if (0) {
2832 error:
2833 status = -1;
2834 }
2835 Py_XDECREF(module_name);
2836 Py_XDECREF(global_name);
2837 Py_XDECREF(module);
2838
2839 return status;
2840}
2841
2842static int
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002843save_ellipsis(PicklerObject *self, PyObject *obj)
2844{
Łukasz Langadbd78252012-03-12 22:59:11 +01002845 PyObject *str = PyUnicode_FromString("Ellipsis");
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002846 int res;
Łukasz Langadbd78252012-03-12 22:59:11 +01002847 if (str == NULL)
Łukasz Langacad1a072012-03-12 23:41:07 +01002848 return -1;
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002849 res = save_global(self, Py_Ellipsis, str);
2850 Py_DECREF(str);
2851 return res;
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002852}
2853
2854static int
2855save_notimplemented(PicklerObject *self, PyObject *obj)
2856{
Łukasz Langadbd78252012-03-12 22:59:11 +01002857 PyObject *str = PyUnicode_FromString("NotImplemented");
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002858 int res;
Łukasz Langadbd78252012-03-12 22:59:11 +01002859 if (str == NULL)
Łukasz Langacad1a072012-03-12 23:41:07 +01002860 return -1;
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002861 res = save_global(self, Py_NotImplemented, str);
2862 Py_DECREF(str);
2863 return res;
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002864}
2865
2866static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002867save_pers(PicklerObject *self, PyObject *obj, PyObject *func)
2868{
2869 PyObject *pid = NULL;
2870 int status = 0;
2871
2872 const char persid_op = PERSID;
2873 const char binpersid_op = BINPERSID;
2874
2875 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002876 pid = _Pickler_FastCall(self, func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002877 if (pid == NULL)
2878 return -1;
2879
2880 if (pid != Py_None) {
2881 if (self->bin) {
2882 if (save(self, pid, 1) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002883 _Pickler_Write(self, &binpersid_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002884 goto error;
2885 }
2886 else {
2887 PyObject *pid_str = NULL;
2888 char *pid_ascii_bytes;
2889 Py_ssize_t size;
2890
2891 pid_str = PyObject_Str(pid);
2892 if (pid_str == NULL)
2893 goto error;
2894
2895 /* XXX: Should it check whether the persistent id only contains
2896 ASCII characters? And what if the pid contains embedded
2897 newlines? */
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00002898 pid_ascii_bytes = _PyUnicode_AsStringAndSize(pid_str, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002899 Py_DECREF(pid_str);
2900 if (pid_ascii_bytes == NULL)
2901 goto error;
2902
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002903 if (_Pickler_Write(self, &persid_op, 1) < 0 ||
2904 _Pickler_Write(self, pid_ascii_bytes, size) < 0 ||
2905 _Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002906 goto error;
2907 }
2908 status = 1;
2909 }
2910
2911 if (0) {
2912 error:
2913 status = -1;
2914 }
2915 Py_XDECREF(pid);
2916
2917 return status;
2918}
2919
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002920static PyObject *
2921get_class(PyObject *obj)
2922{
2923 PyObject *cls;
2924 static PyObject *str_class;
2925
2926 if (str_class == NULL) {
2927 str_class = PyUnicode_InternFromString("__class__");
2928 if (str_class == NULL)
2929 return NULL;
2930 }
2931 cls = PyObject_GetAttr(obj, str_class);
2932 if (cls == NULL) {
2933 if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
2934 PyErr_Clear();
2935 cls = (PyObject *) Py_TYPE(obj);
2936 Py_INCREF(cls);
2937 }
2938 }
2939 return cls;
2940}
2941
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002942/* We're saving obj, and args is the 2-thru-5 tuple returned by the
2943 * appropriate __reduce__ method for obj.
2944 */
2945static int
2946save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
2947{
2948 PyObject *callable;
2949 PyObject *argtup;
2950 PyObject *state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002951 PyObject *listitems = Py_None;
2952 PyObject *dictitems = Py_None;
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002953 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002954
2955 int use_newobj = self->proto >= 2;
2956
2957 const char reduce_op = REDUCE;
2958 const char build_op = BUILD;
2959 const char newobj_op = NEWOBJ;
2960
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002961 size = PyTuple_Size(args);
2962 if (size < 2 || size > 5) {
2963 PyErr_SetString(PicklingError, "tuple returned by "
2964 "__reduce__ must contain 2 through 5 elements");
2965 return -1;
2966 }
2967
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002968 if (!PyArg_UnpackTuple(args, "save_reduce", 2, 5,
2969 &callable, &argtup, &state, &listitems, &dictitems))
2970 return -1;
2971
2972 if (!PyCallable_Check(callable)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002973 PyErr_SetString(PicklingError, "first item of the tuple "
2974 "returned by __reduce__ must be callable");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002975 return -1;
2976 }
2977 if (!PyTuple_Check(argtup)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002978 PyErr_SetString(PicklingError, "second item of the tuple "
2979 "returned by __reduce__ must be a tuple");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002980 return -1;
2981 }
2982
2983 if (state == Py_None)
2984 state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002985
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002986 if (listitems == Py_None)
2987 listitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002988 else if (!PyIter_Check(listitems)) {
Alexandre Vassalotti00d83f22013-04-14 01:28:01 -07002989 PyErr_Format(PicklingError, "fourth element of the tuple "
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002990 "returned by __reduce__ must be an iterator, not %s",
2991 Py_TYPE(listitems)->tp_name);
2992 return -1;
2993 }
2994
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002995 if (dictitems == Py_None)
2996 dictitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002997 else if (!PyIter_Check(dictitems)) {
Alexandre Vassalotti00d83f22013-04-14 01:28:01 -07002998 PyErr_Format(PicklingError, "fifth element of the tuple "
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002999 "returned by __reduce__ must be an iterator, not %s",
3000 Py_TYPE(dictitems)->tp_name);
3001 return -1;
3002 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003003
3004 /* Protocol 2 special case: if callable's name is __newobj__, use
3005 NEWOBJ. */
3006 if (use_newobj) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003007 static PyObject *newobj_str = NULL, *name_str = NULL;
3008 PyObject *name;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003009
3010 if (newobj_str == NULL) {
3011 newobj_str = PyUnicode_InternFromString("__newobj__");
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003012 name_str = PyUnicode_InternFromString("__name__");
3013 if (newobj_str == NULL || name_str == NULL)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003014 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003015 }
3016
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003017 name = PyObject_GetAttr(callable, name_str);
3018 if (name == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003019 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3020 PyErr_Clear();
3021 else
3022 return -1;
3023 use_newobj = 0;
3024 }
3025 else {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003026 use_newobj = PyUnicode_Check(name) &&
3027 PyUnicode_Compare(name, newobj_str) == 0;
3028 Py_DECREF(name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003029 }
3030 }
3031 if (use_newobj) {
3032 PyObject *cls;
3033 PyObject *newargtup;
3034 PyObject *obj_class;
3035 int p;
3036
3037 /* Sanity checks. */
3038 if (Py_SIZE(argtup) < 1) {
3039 PyErr_SetString(PicklingError, "__newobj__ arglist is empty");
3040 return -1;
3041 }
3042
3043 cls = PyTuple_GET_ITEM(argtup, 0);
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003044 if (!PyType_Check(cls)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003045 PyErr_SetString(PicklingError, "args[0] from "
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003046 "__newobj__ args is not a type");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003047 return -1;
3048 }
3049
3050 if (obj != NULL) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003051 obj_class = get_class(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003052 p = obj_class != cls; /* true iff a problem */
3053 Py_DECREF(obj_class);
3054 if (p) {
3055 PyErr_SetString(PicklingError, "args[0] from "
3056 "__newobj__ args has the wrong class");
3057 return -1;
3058 }
3059 }
3060 /* XXX: These calls save() are prone to infinite recursion. Imagine
3061 what happen if the value returned by the __reduce__() method of
3062 some extension type contains another object of the same type. Ouch!
3063
3064 Here is a quick example, that I ran into, to illustrate what I
3065 mean:
3066
3067 >>> import pickle, copyreg
3068 >>> copyreg.dispatch_table.pop(complex)
3069 >>> pickle.dumps(1+2j)
3070 Traceback (most recent call last):
3071 ...
3072 RuntimeError: maximum recursion depth exceeded
3073
3074 Removing the complex class from copyreg.dispatch_table made the
3075 __reduce_ex__() method emit another complex object:
3076
3077 >>> (1+1j).__reduce_ex__(2)
3078 (<function __newobj__ at 0xb7b71c3c>,
3079 (<class 'complex'>, (1+1j)), None, None, None)
3080
3081 Thus when save() was called on newargstup (the 2nd item) recursion
3082 ensued. Of course, the bug was in the complex class which had a
3083 broken __getnewargs__() that emitted another complex object. But,
3084 the point, here, is it is quite easy to end up with a broken reduce
3085 function. */
3086
3087 /* Save the class and its __new__ arguments. */
3088 if (save(self, cls, 0) < 0)
3089 return -1;
3090
3091 newargtup = PyTuple_GetSlice(argtup, 1, Py_SIZE(argtup));
3092 if (newargtup == NULL)
3093 return -1;
3094
3095 p = save(self, newargtup, 0);
3096 Py_DECREF(newargtup);
3097 if (p < 0)
3098 return -1;
3099
3100 /* Add NEWOBJ opcode. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003101 if (_Pickler_Write(self, &newobj_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003102 return -1;
3103 }
3104 else { /* Not using NEWOBJ. */
3105 if (save(self, callable, 0) < 0 ||
3106 save(self, argtup, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003107 _Pickler_Write(self, &reduce_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003108 return -1;
3109 }
3110
3111 /* obj can be NULL when save_reduce() is used directly. A NULL obj means
3112 the caller do not want to memoize the object. Not particularly useful,
3113 but that is to mimic the behavior save_reduce() in pickle.py when
3114 obj is None. */
3115 if (obj && memo_put(self, obj) < 0)
3116 return -1;
3117
3118 if (listitems && batch_list(self, listitems) < 0)
3119 return -1;
3120
3121 if (dictitems && batch_dict(self, dictitems) < 0)
3122 return -1;
3123
3124 if (state) {
Victor Stinner121aab42011-09-29 23:40:53 +02003125 if (save(self, state, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003126 _Pickler_Write(self, &build_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003127 return -1;
3128 }
3129
3130 return 0;
3131}
3132
3133static int
3134save(PicklerObject *self, PyObject *obj, int pers_save)
3135{
3136 PyTypeObject *type;
3137 PyObject *reduce_func = NULL;
3138 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003139 int status = 0;
3140
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00003141 if (Py_EnterRecursiveCall(" while pickling an object"))
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003142 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003143
3144 /* The extra pers_save argument is necessary to avoid calling save_pers()
3145 on its returned object. */
3146 if (!pers_save && self->pers_func) {
3147 /* save_pers() returns:
3148 -1 to signal an error;
3149 0 if it did nothing successfully;
3150 1 if a persistent id was saved.
3151 */
3152 if ((status = save_pers(self, obj, self->pers_func)) != 0)
3153 goto done;
3154 }
3155
3156 type = Py_TYPE(obj);
3157
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003158 /* The old cPickle had an optimization that used switch-case statement
3159 dispatching on the first letter of the type name. This has was removed
3160 since benchmarks shown that this optimization was actually slowing
3161 things down. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003162
3163 /* Atom types; these aren't memoized, so don't check the memo. */
3164
3165 if (obj == Py_None) {
3166 status = save_none(self, obj);
3167 goto done;
3168 }
Łukasz Langaf3078fb2012-03-12 19:46:12 +01003169 else if (obj == Py_Ellipsis) {
3170 status = save_ellipsis(self, obj);
3171 goto done;
3172 }
3173 else if (obj == Py_NotImplemented) {
3174 status = save_notimplemented(self, obj);
3175 goto done;
3176 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003177 else if (obj == Py_False || obj == Py_True) {
3178 status = save_bool(self, obj);
3179 goto done;
3180 }
3181 else if (type == &PyLong_Type) {
3182 status = save_long(self, obj);
3183 goto done;
3184 }
3185 else if (type == &PyFloat_Type) {
3186 status = save_float(self, obj);
3187 goto done;
3188 }
3189
3190 /* Check the memo to see if it has the object. If so, generate
3191 a GET (or BINGET) opcode, instead of pickling the object
3192 once again. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003193 if (PyMemoTable_Get(self->memo, obj)) {
3194 if (memo_get(self, obj) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003195 goto error;
3196 goto done;
3197 }
3198
3199 if (type == &PyBytes_Type) {
3200 status = save_bytes(self, obj);
3201 goto done;
3202 }
3203 else if (type == &PyUnicode_Type) {
3204 status = save_unicode(self, obj);
3205 goto done;
3206 }
3207 else if (type == &PyDict_Type) {
3208 status = save_dict(self, obj);
3209 goto done;
3210 }
3211 else if (type == &PyList_Type) {
3212 status = save_list(self, obj);
3213 goto done;
3214 }
3215 else if (type == &PyTuple_Type) {
3216 status = save_tuple(self, obj);
3217 goto done;
3218 }
3219 else if (type == &PyType_Type) {
3220 status = save_global(self, obj, NULL);
3221 goto done;
3222 }
3223 else if (type == &PyFunction_Type) {
3224 status = save_global(self, obj, NULL);
3225 if (status < 0 && PyErr_ExceptionMatches(PickleError)) {
3226 /* fall back to reduce */
3227 PyErr_Clear();
3228 }
3229 else {
3230 goto done;
3231 }
3232 }
3233 else if (type == &PyCFunction_Type) {
3234 status = save_global(self, obj, NULL);
3235 goto done;
3236 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003237
3238 /* XXX: This part needs some unit tests. */
3239
3240 /* Get a reduction callable, and call it. This may come from
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003241 * self.dispatch_table, copyreg.dispatch_table, the object's
3242 * __reduce_ex__ method, or the object's __reduce__ method.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003243 */
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003244 if (self->dispatch_table == NULL) {
3245 reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
3246 /* PyDict_GetItem() unlike PyObject_GetItem() and
3247 PyObject_GetAttr() returns a borrowed ref */
3248 Py_XINCREF(reduce_func);
3249 } else {
3250 reduce_func = PyObject_GetItem(self->dispatch_table, (PyObject *)type);
3251 if (reduce_func == NULL) {
3252 if (PyErr_ExceptionMatches(PyExc_KeyError))
3253 PyErr_Clear();
3254 else
3255 goto error;
3256 }
3257 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003258 if (reduce_func != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003259 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003260 reduce_value = _Pickler_FastCall(self, reduce_func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003261 }
Antoine Pitrouffd41d92011-10-04 09:23:04 +02003262 else if (PyType_IsSubtype(type, &PyType_Type)) {
3263 status = save_global(self, obj, NULL);
3264 goto done;
3265 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003266 else {
3267 static PyObject *reduce_str = NULL;
3268 static PyObject *reduce_ex_str = NULL;
3269
3270 /* Cache the name of the reduce methods. */
3271 if (reduce_str == NULL) {
3272 reduce_str = PyUnicode_InternFromString("__reduce__");
3273 if (reduce_str == NULL)
3274 goto error;
3275 reduce_ex_str = PyUnicode_InternFromString("__reduce_ex__");
3276 if (reduce_ex_str == NULL)
3277 goto error;
3278 }
3279
3280 /* XXX: If the __reduce__ method is defined, __reduce_ex__ is
3281 automatically defined as __reduce__. While this is convenient, this
3282 make it impossible to know which method was actually called. Of
3283 course, this is not a big deal. But still, it would be nice to let
3284 the user know which method was called when something go
3285 wrong. Incidentally, this means if __reduce_ex__ is not defined, we
3286 don't actually have to check for a __reduce__ method. */
3287
3288 /* Check for a __reduce_ex__ method. */
3289 reduce_func = PyObject_GetAttr(obj, reduce_ex_str);
3290 if (reduce_func != NULL) {
3291 PyObject *proto;
3292 proto = PyLong_FromLong(self->proto);
3293 if (proto != NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003294 reduce_value = _Pickler_FastCall(self, reduce_func, proto);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003295 }
3296 }
3297 else {
3298 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3299 PyErr_Clear();
3300 else
3301 goto error;
3302 /* Check for a __reduce__ method. */
3303 reduce_func = PyObject_GetAttr(obj, reduce_str);
3304 if (reduce_func != NULL) {
3305 reduce_value = PyObject_Call(reduce_func, empty_tuple, NULL);
3306 }
3307 else {
3308 PyErr_Format(PicklingError, "can't pickle '%.200s' object: %R",
3309 type->tp_name, obj);
3310 goto error;
3311 }
3312 }
3313 }
3314
3315 if (reduce_value == NULL)
3316 goto error;
3317
3318 if (PyUnicode_Check(reduce_value)) {
3319 status = save_global(self, obj, reduce_value);
3320 goto done;
3321 }
3322
3323 if (!PyTuple_Check(reduce_value)) {
3324 PyErr_SetString(PicklingError,
3325 "__reduce__ must return a string or tuple");
3326 goto error;
3327 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003328
3329 status = save_reduce(self, reduce_value, obj);
3330
3331 if (0) {
3332 error:
3333 status = -1;
3334 }
3335 done:
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003336 Py_LeaveRecursiveCall();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003337 Py_XDECREF(reduce_func);
3338 Py_XDECREF(reduce_value);
3339
3340 return status;
3341}
3342
3343static int
3344dump(PicklerObject *self, PyObject *obj)
3345{
3346 const char stop_op = STOP;
3347
3348 if (self->proto >= 2) {
3349 char header[2];
3350
3351 header[0] = PROTO;
3352 assert(self->proto >= 0 && self->proto < 256);
3353 header[1] = (unsigned char)self->proto;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003354 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003355 return -1;
3356 }
3357
3358 if (save(self, obj, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003359 _Pickler_Write(self, &stop_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003360 return -1;
3361
3362 return 0;
3363}
3364
3365PyDoc_STRVAR(Pickler_clear_memo_doc,
3366"clear_memo() -> None. Clears the pickler's \"memo\"."
3367"\n"
3368"The memo is the data structure that remembers which objects the\n"
3369"pickler has already seen, so that shared or recursive objects are\n"
3370"pickled by reference and not by value. This method is useful when\n"
3371"re-using picklers.");
3372
3373static PyObject *
3374Pickler_clear_memo(PicklerObject *self)
3375{
3376 if (self->memo)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003377 PyMemoTable_Clear(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003378
3379 Py_RETURN_NONE;
3380}
3381
3382PyDoc_STRVAR(Pickler_dump_doc,
3383"dump(obj) -> None. Write a pickled representation of obj to the open file.");
3384
3385static PyObject *
3386Pickler_dump(PicklerObject *self, PyObject *args)
3387{
3388 PyObject *obj;
3389
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003390 /* Check whether the Pickler was initialized correctly (issue3664).
3391 Developers often forget to call __init__() in their subclasses, which
3392 would trigger a segfault without this check. */
3393 if (self->write == NULL) {
Victor Stinner121aab42011-09-29 23:40:53 +02003394 PyErr_Format(PicklingError,
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003395 "Pickler.__init__() was not called by %s.__init__()",
3396 Py_TYPE(self)->tp_name);
3397 return NULL;
3398 }
3399
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003400 if (!PyArg_ParseTuple(args, "O:dump", &obj))
3401 return NULL;
3402
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003403 if (_Pickler_ClearBuffer(self) < 0)
3404 return NULL;
3405
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003406 if (dump(self, obj) < 0)
3407 return NULL;
3408
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003409 if (_Pickler_FlushToFile(self) < 0)
3410 return NULL;
3411
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003412 Py_RETURN_NONE;
3413}
3414
3415static struct PyMethodDef Pickler_methods[] = {
3416 {"dump", (PyCFunction)Pickler_dump, METH_VARARGS,
3417 Pickler_dump_doc},
3418 {"clear_memo", (PyCFunction)Pickler_clear_memo, METH_NOARGS,
3419 Pickler_clear_memo_doc},
3420 {NULL, NULL} /* sentinel */
3421};
3422
3423static void
3424Pickler_dealloc(PicklerObject *self)
3425{
3426 PyObject_GC_UnTrack(self);
3427
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003428 Py_XDECREF(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003429 Py_XDECREF(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003430 Py_XDECREF(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003431 Py_XDECREF(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003432 Py_XDECREF(self->arg);
3433 Py_XDECREF(self->fast_memo);
3434
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003435 PyMemoTable_Del(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003436
3437 Py_TYPE(self)->tp_free((PyObject *)self);
3438}
3439
3440static int
3441Pickler_traverse(PicklerObject *self, visitproc visit, void *arg)
3442{
3443 Py_VISIT(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003444 Py_VISIT(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003445 Py_VISIT(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003446 Py_VISIT(self->arg);
3447 Py_VISIT(self->fast_memo);
3448 return 0;
3449}
3450
3451static int
3452Pickler_clear(PicklerObject *self)
3453{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003454 Py_CLEAR(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003455 Py_CLEAR(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003456 Py_CLEAR(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003457 Py_CLEAR(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003458 Py_CLEAR(self->arg);
3459 Py_CLEAR(self->fast_memo);
3460
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003461 if (self->memo != NULL) {
3462 PyMemoTable *memo = self->memo;
3463 self->memo = NULL;
3464 PyMemoTable_Del(memo);
3465 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003466 return 0;
3467}
3468
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003469
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003470PyDoc_STRVAR(Pickler_doc,
3471"Pickler(file, protocol=None)"
3472"\n"
3473"This takes a binary file for writing a pickle data stream.\n"
3474"\n"
3475"The optional protocol argument tells the pickler to use the\n"
3476"given protocol; supported protocols are 0, 1, 2, 3. The default\n"
3477"protocol is 3; a backward-incompatible protocol designed for\n"
3478"Python 3.0.\n"
3479"\n"
3480"Specifying a negative protocol version selects the highest\n"
3481"protocol version supported. The higher the protocol used, the\n"
3482"more recent the version of Python needed to read the pickle\n"
3483"produced.\n"
3484"\n"
3485"The file argument must have a write() method that accepts a single\n"
3486"bytes argument. It can thus be a file object opened for binary\n"
3487"writing, a io.BytesIO instance, or any other custom object that\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003488"meets this interface.\n"
3489"\n"
3490"If fix_imports is True and protocol is less than 3, pickle will try to\n"
3491"map the new Python 3.x names to the old module names used in Python\n"
3492"2.x, so that the pickle data stream is readable with Python 2.x.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003493
3494static int
3495Pickler_init(PicklerObject *self, PyObject *args, PyObject *kwds)
3496{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003497 static char *kwlist[] = {"file", "protocol", "fix_imports", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003498 PyObject *file;
3499 PyObject *proto_obj = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003500 PyObject *fix_imports = Py_True;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02003501 _Py_IDENTIFIER(persistent_id);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003502 _Py_IDENTIFIER(dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003503
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003504 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:Pickler",
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003505 kwlist, &file, &proto_obj, &fix_imports))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003506 return -1;
3507
3508 /* In case of multiple __init__() calls, clear previous content. */
3509 if (self->write != NULL)
3510 (void)Pickler_clear(self);
3511
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003512 if (_Pickler_SetProtocol(self, proto_obj, fix_imports) < 0)
3513 return -1;
3514
3515 if (_Pickler_SetOutputStream(self, file) < 0)
3516 return -1;
3517
3518 /* memo and output_buffer may have already been created in _Pickler_New */
3519 if (self->memo == NULL) {
3520 self->memo = PyMemoTable_New();
3521 if (self->memo == NULL)
3522 return -1;
3523 }
3524 self->output_len = 0;
3525 if (self->output_buffer == NULL) {
3526 self->max_output_len = WRITE_BUF_SIZE;
3527 self->output_buffer = PyBytes_FromStringAndSize(NULL,
3528 self->max_output_len);
3529 if (self->output_buffer == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003530 return -1;
3531 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003532
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003533 self->arg = NULL;
3534 self->fast = 0;
3535 self->fast_nesting = 0;
3536 self->fast_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003537 self->pers_func = NULL;
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02003538 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_id)) {
3539 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
3540 &PyId_persistent_id);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003541 if (self->pers_func == NULL)
3542 return -1;
3543 }
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003544 self->dispatch_table = NULL;
3545 if (_PyObject_HasAttrId((PyObject *)self, &PyId_dispatch_table)) {
3546 self->dispatch_table = _PyObject_GetAttrId((PyObject *)self,
3547 &PyId_dispatch_table);
3548 if (self->dispatch_table == NULL)
3549 return -1;
3550 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003551 return 0;
3552}
3553
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003554/* Define a proxy object for the Pickler's internal memo object. This is to
3555 * avoid breaking code like:
3556 * pickler.memo.clear()
3557 * and
3558 * pickler.memo = saved_memo
3559 * Is this a good idea? Not really, but we don't want to break code that uses
3560 * it. Note that we don't implement the entire mapping API here. This is
3561 * intentional, as these should be treated as black-box implementation details.
3562 */
3563
3564typedef struct {
3565 PyObject_HEAD
3566 PicklerObject *pickler; /* Pickler whose memo table we're proxying. */
3567} PicklerMemoProxyObject;
3568
3569PyDoc_STRVAR(pmp_clear_doc,
3570"memo.clear() -> None. Remove all items from memo.");
3571
3572static PyObject *
3573pmp_clear(PicklerMemoProxyObject *self)
3574{
3575 if (self->pickler->memo)
3576 PyMemoTable_Clear(self->pickler->memo);
3577 Py_RETURN_NONE;
3578}
3579
3580PyDoc_STRVAR(pmp_copy_doc,
3581"memo.copy() -> new_memo. Copy the memo to a new object.");
3582
3583static PyObject *
3584pmp_copy(PicklerMemoProxyObject *self)
3585{
3586 Py_ssize_t i;
3587 PyMemoTable *memo;
3588 PyObject *new_memo = PyDict_New();
3589 if (new_memo == NULL)
3590 return NULL;
3591
3592 memo = self->pickler->memo;
3593 for (i = 0; i < memo->mt_allocated; ++i) {
3594 PyMemoEntry entry = memo->mt_table[i];
3595 if (entry.me_key != NULL) {
3596 int status;
3597 PyObject *key, *value;
3598
3599 key = PyLong_FromVoidPtr(entry.me_key);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003600 value = Py_BuildValue("nO", entry.me_value, entry.me_key);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003601
3602 if (key == NULL || value == NULL) {
3603 Py_XDECREF(key);
3604 Py_XDECREF(value);
3605 goto error;
3606 }
3607 status = PyDict_SetItem(new_memo, key, value);
3608 Py_DECREF(key);
3609 Py_DECREF(value);
3610 if (status < 0)
3611 goto error;
3612 }
3613 }
3614 return new_memo;
3615
3616 error:
3617 Py_XDECREF(new_memo);
3618 return NULL;
3619}
3620
3621PyDoc_STRVAR(pmp_reduce_doc,
3622"memo.__reduce__(). Pickling support.");
3623
3624static PyObject *
3625pmp_reduce(PicklerMemoProxyObject *self, PyObject *args)
3626{
3627 PyObject *reduce_value, *dict_args;
3628 PyObject *contents = pmp_copy(self);
3629 if (contents == NULL)
3630 return NULL;
3631
3632 reduce_value = PyTuple_New(2);
3633 if (reduce_value == NULL) {
3634 Py_DECREF(contents);
3635 return NULL;
3636 }
3637 dict_args = PyTuple_New(1);
3638 if (dict_args == NULL) {
3639 Py_DECREF(contents);
3640 Py_DECREF(reduce_value);
3641 return NULL;
3642 }
3643 PyTuple_SET_ITEM(dict_args, 0, contents);
3644 Py_INCREF((PyObject *)&PyDict_Type);
3645 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
3646 PyTuple_SET_ITEM(reduce_value, 1, dict_args);
3647 return reduce_value;
3648}
3649
3650static PyMethodDef picklerproxy_methods[] = {
3651 {"clear", (PyCFunction)pmp_clear, METH_NOARGS, pmp_clear_doc},
3652 {"copy", (PyCFunction)pmp_copy, METH_NOARGS, pmp_copy_doc},
3653 {"__reduce__", (PyCFunction)pmp_reduce, METH_VARARGS, pmp_reduce_doc},
3654 {NULL, NULL} /* sentinel */
3655};
3656
3657static void
3658PicklerMemoProxy_dealloc(PicklerMemoProxyObject *self)
3659{
3660 PyObject_GC_UnTrack(self);
3661 Py_XDECREF(self->pickler);
3662 PyObject_GC_Del((PyObject *)self);
3663}
3664
3665static int
3666PicklerMemoProxy_traverse(PicklerMemoProxyObject *self,
3667 visitproc visit, void *arg)
3668{
3669 Py_VISIT(self->pickler);
3670 return 0;
3671}
3672
3673static int
3674PicklerMemoProxy_clear(PicklerMemoProxyObject *self)
3675{
3676 Py_CLEAR(self->pickler);
3677 return 0;
3678}
3679
3680static PyTypeObject PicklerMemoProxyType = {
3681 PyVarObject_HEAD_INIT(NULL, 0)
3682 "_pickle.PicklerMemoProxy", /*tp_name*/
3683 sizeof(PicklerMemoProxyObject), /*tp_basicsize*/
3684 0,
3685 (destructor)PicklerMemoProxy_dealloc, /* tp_dealloc */
3686 0, /* tp_print */
3687 0, /* tp_getattr */
3688 0, /* tp_setattr */
3689 0, /* tp_compare */
3690 0, /* tp_repr */
3691 0, /* tp_as_number */
3692 0, /* tp_as_sequence */
3693 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00003694 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003695 0, /* tp_call */
3696 0, /* tp_str */
3697 PyObject_GenericGetAttr, /* tp_getattro */
3698 PyObject_GenericSetAttr, /* tp_setattro */
3699 0, /* tp_as_buffer */
3700 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3701 0, /* tp_doc */
3702 (traverseproc)PicklerMemoProxy_traverse, /* tp_traverse */
3703 (inquiry)PicklerMemoProxy_clear, /* tp_clear */
3704 0, /* tp_richcompare */
3705 0, /* tp_weaklistoffset */
3706 0, /* tp_iter */
3707 0, /* tp_iternext */
3708 picklerproxy_methods, /* tp_methods */
3709};
3710
3711static PyObject *
3712PicklerMemoProxy_New(PicklerObject *pickler)
3713{
3714 PicklerMemoProxyObject *self;
3715
3716 self = PyObject_GC_New(PicklerMemoProxyObject, &PicklerMemoProxyType);
3717 if (self == NULL)
3718 return NULL;
3719 Py_INCREF(pickler);
3720 self->pickler = pickler;
3721 PyObject_GC_Track(self);
3722 return (PyObject *)self;
3723}
3724
3725/*****************************************************************************/
3726
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003727static PyObject *
3728Pickler_get_memo(PicklerObject *self)
3729{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003730 return PicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003731}
3732
3733static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003734Pickler_set_memo(PicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003735{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003736 PyMemoTable *new_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003737
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003738 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003739 PyErr_SetString(PyExc_TypeError,
3740 "attribute deletion is not supported");
3741 return -1;
3742 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003743
3744 if (Py_TYPE(obj) == &PicklerMemoProxyType) {
3745 PicklerObject *pickler =
3746 ((PicklerMemoProxyObject *)obj)->pickler;
3747
3748 new_memo = PyMemoTable_Copy(pickler->memo);
3749 if (new_memo == NULL)
3750 return -1;
3751 }
3752 else if (PyDict_Check(obj)) {
3753 Py_ssize_t i = 0;
3754 PyObject *key, *value;
3755
3756 new_memo = PyMemoTable_New();
3757 if (new_memo == NULL)
3758 return -1;
3759
3760 while (PyDict_Next(obj, &i, &key, &value)) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003761 Py_ssize_t memo_id;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003762 PyObject *memo_obj;
3763
3764 if (!PyTuple_Check(value) || Py_SIZE(value) != 2) {
3765 PyErr_SetString(PyExc_TypeError,
3766 "'memo' values must be 2-item tuples");
3767 goto error;
3768 }
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003769 memo_id = PyLong_AsSsize_t(PyTuple_GET_ITEM(value, 0));
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003770 if (memo_id == -1 && PyErr_Occurred())
3771 goto error;
3772 memo_obj = PyTuple_GET_ITEM(value, 1);
3773 if (PyMemoTable_Set(new_memo, memo_obj, memo_id) < 0)
3774 goto error;
3775 }
3776 }
3777 else {
3778 PyErr_Format(PyExc_TypeError,
3779 "'memo' attribute must be an PicklerMemoProxy object"
3780 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003781 return -1;
3782 }
3783
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003784 PyMemoTable_Del(self->memo);
3785 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003786
3787 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003788
3789 error:
3790 if (new_memo)
3791 PyMemoTable_Del(new_memo);
3792 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003793}
3794
3795static PyObject *
3796Pickler_get_persid(PicklerObject *self)
3797{
3798 if (self->pers_func == NULL)
3799 PyErr_SetString(PyExc_AttributeError, "persistent_id");
3800 else
3801 Py_INCREF(self->pers_func);
3802 return self->pers_func;
3803}
3804
3805static int
3806Pickler_set_persid(PicklerObject *self, PyObject *value)
3807{
3808 PyObject *tmp;
3809
3810 if (value == NULL) {
3811 PyErr_SetString(PyExc_TypeError,
3812 "attribute deletion is not supported");
3813 return -1;
3814 }
3815 if (!PyCallable_Check(value)) {
3816 PyErr_SetString(PyExc_TypeError,
3817 "persistent_id must be a callable taking one argument");
3818 return -1;
3819 }
3820
3821 tmp = self->pers_func;
3822 Py_INCREF(value);
3823 self->pers_func = value;
3824 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
3825
3826 return 0;
3827}
3828
3829static PyMemberDef Pickler_members[] = {
3830 {"bin", T_INT, offsetof(PicklerObject, bin)},
3831 {"fast", T_INT, offsetof(PicklerObject, fast)},
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003832 {"dispatch_table", T_OBJECT_EX, offsetof(PicklerObject, dispatch_table)},
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003833 {NULL}
3834};
3835
3836static PyGetSetDef Pickler_getsets[] = {
3837 {"memo", (getter)Pickler_get_memo,
3838 (setter)Pickler_set_memo},
3839 {"persistent_id", (getter)Pickler_get_persid,
3840 (setter)Pickler_set_persid},
3841 {NULL}
3842};
3843
3844static PyTypeObject Pickler_Type = {
3845 PyVarObject_HEAD_INIT(NULL, 0)
3846 "_pickle.Pickler" , /*tp_name*/
3847 sizeof(PicklerObject), /*tp_basicsize*/
3848 0, /*tp_itemsize*/
3849 (destructor)Pickler_dealloc, /*tp_dealloc*/
3850 0, /*tp_print*/
3851 0, /*tp_getattr*/
3852 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00003853 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003854 0, /*tp_repr*/
3855 0, /*tp_as_number*/
3856 0, /*tp_as_sequence*/
3857 0, /*tp_as_mapping*/
3858 0, /*tp_hash*/
3859 0, /*tp_call*/
3860 0, /*tp_str*/
3861 0, /*tp_getattro*/
3862 0, /*tp_setattro*/
3863 0, /*tp_as_buffer*/
3864 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3865 Pickler_doc, /*tp_doc*/
3866 (traverseproc)Pickler_traverse, /*tp_traverse*/
3867 (inquiry)Pickler_clear, /*tp_clear*/
3868 0, /*tp_richcompare*/
3869 0, /*tp_weaklistoffset*/
3870 0, /*tp_iter*/
3871 0, /*tp_iternext*/
3872 Pickler_methods, /*tp_methods*/
3873 Pickler_members, /*tp_members*/
3874 Pickler_getsets, /*tp_getset*/
3875 0, /*tp_base*/
3876 0, /*tp_dict*/
3877 0, /*tp_descr_get*/
3878 0, /*tp_descr_set*/
3879 0, /*tp_dictoffset*/
3880 (initproc)Pickler_init, /*tp_init*/
3881 PyType_GenericAlloc, /*tp_alloc*/
3882 PyType_GenericNew, /*tp_new*/
3883 PyObject_GC_Del, /*tp_free*/
3884 0, /*tp_is_gc*/
3885};
3886
Victor Stinner121aab42011-09-29 23:40:53 +02003887/* Temporary helper for calling self.find_class().
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003888
3889 XXX: It would be nice to able to avoid Python function call overhead, by
3890 using directly the C version of find_class(), when find_class() is not
3891 overridden by a subclass. Although, this could become rather hackish. A
3892 simpler optimization would be to call the C function when self is not a
3893 subclass instance. */
3894static PyObject *
3895find_class(UnpicklerObject *self, PyObject *module_name, PyObject *global_name)
3896{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02003897 _Py_IDENTIFIER(find_class);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02003898
3899 return _PyObject_CallMethodId((PyObject *)self, &PyId_find_class, "OO",
3900 module_name, global_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003901}
3902
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003903static Py_ssize_t
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003904marker(UnpicklerObject *self)
3905{
3906 if (self->num_marks < 1) {
3907 PyErr_SetString(UnpicklingError, "could not find MARK");
3908 return -1;
3909 }
3910
3911 return self->marks[--self->num_marks];
3912}
3913
3914static int
3915load_none(UnpicklerObject *self)
3916{
3917 PDATA_APPEND(self->stack, Py_None, -1);
3918 return 0;
3919}
3920
3921static int
3922bad_readline(void)
3923{
3924 PyErr_SetString(UnpicklingError, "pickle data was truncated");
3925 return -1;
3926}
3927
3928static int
3929load_int(UnpicklerObject *self)
3930{
3931 PyObject *value;
3932 char *endptr, *s;
3933 Py_ssize_t len;
3934 long x;
3935
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003936 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003937 return -1;
3938 if (len < 2)
3939 return bad_readline();
3940
3941 errno = 0;
Victor Stinner121aab42011-09-29 23:40:53 +02003942 /* XXX: Should the base argument of strtol() be explicitly set to 10?
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003943 XXX(avassalotti): Should this uses PyOS_strtol()? */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003944 x = strtol(s, &endptr, 0);
3945
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003946 if (errno || (*endptr != '\n' && *endptr != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003947 /* Hm, maybe we've got something long. Let's try reading
3948 * it as a Python long object. */
3949 errno = 0;
3950 /* XXX: Same thing about the base here. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003951 value = PyLong_FromString(s, NULL, 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003952 if (value == NULL) {
3953 PyErr_SetString(PyExc_ValueError,
3954 "could not convert string to int");
3955 return -1;
3956 }
3957 }
3958 else {
3959 if (len == 3 && (x == 0 || x == 1)) {
3960 if ((value = PyBool_FromLong(x)) == NULL)
3961 return -1;
3962 }
3963 else {
3964 if ((value = PyLong_FromLong(x)) == NULL)
3965 return -1;
3966 }
3967 }
3968
3969 PDATA_PUSH(self->stack, value, -1);
3970 return 0;
3971}
3972
3973static int
3974load_bool(UnpicklerObject *self, PyObject *boolean)
3975{
3976 assert(boolean == Py_True || boolean == Py_False);
3977 PDATA_APPEND(self->stack, boolean, -1);
3978 return 0;
3979}
3980
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003981/* s contains x bytes of an unsigned little-endian integer. Return its value
3982 * as a C Py_ssize_t, or -1 if it's higher than PY_SSIZE_T_MAX.
3983 */
3984static Py_ssize_t
3985calc_binsize(char *bytes, int size)
3986{
3987 unsigned char *s = (unsigned char *)bytes;
3988 size_t x = 0;
3989
3990 assert(size == 4);
3991
3992 x = (size_t) s[0];
3993 x |= (size_t) s[1] << 8;
3994 x |= (size_t) s[2] << 16;
3995 x |= (size_t) s[3] << 24;
3996
3997 if (x > PY_SSIZE_T_MAX)
3998 return -1;
3999 else
4000 return (Py_ssize_t) x;
4001}
4002
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004003/* s contains x bytes of a little-endian integer. Return its value as a
4004 * C int. Obscure: when x is 1 or 2, this is an unsigned little-endian
4005 * int, but when x is 4 it's a signed one. This is an historical source
4006 * of x-platform bugs.
4007 */
4008static long
4009calc_binint(char *bytes, int size)
4010{
4011 unsigned char *s = (unsigned char *)bytes;
4012 int i = size;
4013 long x = 0;
4014
4015 for (i = 0; i < size; i++) {
4016 x |= (long)s[i] << (i * 8);
4017 }
4018
4019 /* Unlike BININT1 and BININT2, BININT (more accurately BININT4)
4020 * is signed, so on a box with longs bigger than 4 bytes we need
4021 * to extend a BININT's sign bit to the full width.
4022 */
4023 if (SIZEOF_LONG > 4 && size == 4) {
4024 x |= -(x & (1L << 31));
4025 }
4026
4027 return x;
4028}
4029
4030static int
4031load_binintx(UnpicklerObject *self, char *s, int size)
4032{
4033 PyObject *value;
4034 long x;
4035
4036 x = calc_binint(s, size);
4037
4038 if ((value = PyLong_FromLong(x)) == NULL)
4039 return -1;
4040
4041 PDATA_PUSH(self->stack, value, -1);
4042 return 0;
4043}
4044
4045static int
4046load_binint(UnpicklerObject *self)
4047{
4048 char *s;
4049
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004050 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004051 return -1;
4052
4053 return load_binintx(self, s, 4);
4054}
4055
4056static int
4057load_binint1(UnpicklerObject *self)
4058{
4059 char *s;
4060
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004061 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004062 return -1;
4063
4064 return load_binintx(self, s, 1);
4065}
4066
4067static int
4068load_binint2(UnpicklerObject *self)
4069{
4070 char *s;
4071
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004072 if (_Unpickler_Read(self, &s, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004073 return -1;
4074
4075 return load_binintx(self, s, 2);
4076}
4077
4078static int
4079load_long(UnpicklerObject *self)
4080{
4081 PyObject *value;
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004082 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004083 Py_ssize_t len;
4084
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004085 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004086 return -1;
4087 if (len < 2)
4088 return bad_readline();
4089
Mark Dickinson8dd05142009-01-20 20:43:58 +00004090 /* s[len-2] will usually be 'L' (and s[len-1] is '\n'); we need to remove
4091 the 'L' before calling PyLong_FromString. In order to maintain
4092 compatibility with Python 3.0.0, we don't actually *require*
4093 the 'L' to be present. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004094 if (s[len-2] == 'L')
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004095 s[len-2] = '\0';
Alexandre Vassalottie4bccb72009-01-24 01:47:57 +00004096 /* XXX: Should the base argument explicitly set to 10? */
4097 value = PyLong_FromString(s, NULL, 0);
Mark Dickinson8dd05142009-01-20 20:43:58 +00004098 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004099 return -1;
4100
4101 PDATA_PUSH(self->stack, value, -1);
4102 return 0;
4103}
4104
4105/* 'size' bytes contain the # of bytes of little-endian 256's-complement
4106 * data following.
4107 */
4108static int
4109load_counted_long(UnpicklerObject *self, int size)
4110{
4111 PyObject *value;
4112 char *nbytes;
4113 char *pdata;
4114
4115 assert(size == 1 || size == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004116 if (_Unpickler_Read(self, &nbytes, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004117 return -1;
4118
4119 size = calc_binint(nbytes, size);
4120 if (size < 0) {
4121 /* Corrupt or hostile pickle -- we never write one like this */
4122 PyErr_SetString(UnpicklingError,
4123 "LONG pickle has negative byte count");
4124 return -1;
4125 }
4126
4127 if (size == 0)
4128 value = PyLong_FromLong(0L);
4129 else {
4130 /* Read the raw little-endian bytes and convert. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004131 if (_Unpickler_Read(self, &pdata, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004132 return -1;
4133 value = _PyLong_FromByteArray((unsigned char *)pdata, (size_t)size,
4134 1 /* little endian */ , 1 /* signed */ );
4135 }
4136 if (value == NULL)
4137 return -1;
4138 PDATA_PUSH(self->stack, value, -1);
4139 return 0;
4140}
4141
4142static int
4143load_float(UnpicklerObject *self)
4144{
4145 PyObject *value;
4146 char *endptr, *s;
4147 Py_ssize_t len;
4148 double d;
4149
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004150 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004151 return -1;
4152 if (len < 2)
4153 return bad_readline();
4154
4155 errno = 0;
Mark Dickinson725bfd82009-05-03 20:33:40 +00004156 d = PyOS_string_to_double(s, &endptr, PyExc_OverflowError);
4157 if (d == -1.0 && PyErr_Occurred())
4158 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004159 if ((endptr[0] != '\n') && (endptr[0] != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004160 PyErr_SetString(PyExc_ValueError, "could not convert string to float");
4161 return -1;
4162 }
Mark Dickinson725bfd82009-05-03 20:33:40 +00004163 value = PyFloat_FromDouble(d);
4164 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004165 return -1;
4166
4167 PDATA_PUSH(self->stack, value, -1);
4168 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004169}
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004170
4171static int
4172load_binfloat(UnpicklerObject *self)
4173{
4174 PyObject *value;
4175 double x;
4176 char *s;
4177
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004178 if (_Unpickler_Read(self, &s, 8) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004179 return -1;
4180
4181 x = _PyFloat_Unpack8((unsigned char *)s, 0);
4182 if (x == -1.0 && PyErr_Occurred())
4183 return -1;
4184
4185 if ((value = PyFloat_FromDouble(x)) == NULL)
4186 return -1;
4187
4188 PDATA_PUSH(self->stack, value, -1);
4189 return 0;
4190}
4191
4192static int
4193load_string(UnpicklerObject *self)
4194{
4195 PyObject *bytes;
4196 PyObject *str = NULL;
4197 Py_ssize_t len;
4198 char *s, *p;
4199
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004200 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004201 return -1;
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004202 /* Strip the newline */
4203 len--;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004204 /* Strip outermost quotes */
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004205 if (len >= 2 && s[0] == s[len - 1] && (s[0] == '\'' || s[0] == '"')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004206 p = s + 1;
4207 len -= 2;
4208 }
4209 else {
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004210 PyErr_SetString(UnpicklingError,
4211 "the STRING opcode argument must be quoted");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004212 return -1;
4213 }
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004214 assert(len >= 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004215
4216 /* Use the PyBytes API to decode the string, since that is what is used
4217 to encode, and then coerce the result to Unicode. */
4218 bytes = PyBytes_DecodeEscape(p, len, NULL, 0, NULL);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004219 if (bytes == NULL)
4220 return -1;
4221 str = PyUnicode_FromEncodedObject(bytes, self->encoding, self->errors);
4222 Py_DECREF(bytes);
4223 if (str == NULL)
4224 return -1;
4225
4226 PDATA_PUSH(self->stack, str, -1);
4227 return 0;
4228}
4229
4230static int
4231load_binbytes(UnpicklerObject *self)
4232{
4233 PyObject *bytes;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004234 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004235 char *s;
4236
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004237 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004238 return -1;
4239
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004240 x = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004241 if (x < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004242 PyErr_Format(PyExc_OverflowError,
4243 "BINBYTES exceeds system's maximum size of %zd bytes",
Alexandre Vassalotticc757172013-04-14 02:25:10 -07004244 PY_SSIZE_T_MAX);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004245 return -1;
4246 }
4247
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004248 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004249 return -1;
4250 bytes = PyBytes_FromStringAndSize(s, x);
4251 if (bytes == NULL)
4252 return -1;
4253
4254 PDATA_PUSH(self->stack, bytes, -1);
4255 return 0;
4256}
4257
4258static int
4259load_short_binbytes(UnpicklerObject *self)
4260{
4261 PyObject *bytes;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004262 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004263 char *s;
4264
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004265 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004266 return -1;
4267
4268 x = (unsigned char)s[0];
4269
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004270 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004271 return -1;
4272
4273 bytes = PyBytes_FromStringAndSize(s, x);
4274 if (bytes == NULL)
4275 return -1;
4276
4277 PDATA_PUSH(self->stack, bytes, -1);
4278 return 0;
4279}
4280
4281static int
4282load_binstring(UnpicklerObject *self)
4283{
4284 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004285 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004286 char *s;
4287
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004288 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004289 return -1;
4290
4291 x = calc_binint(s, 4);
4292 if (x < 0) {
Victor Stinner121aab42011-09-29 23:40:53 +02004293 PyErr_SetString(UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004294 "BINSTRING pickle has negative byte count");
4295 return -1;
4296 }
4297
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004298 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004299 return -1;
4300
4301 /* Convert Python 2.x strings to unicode. */
4302 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4303 if (str == NULL)
4304 return -1;
4305
4306 PDATA_PUSH(self->stack, str, -1);
4307 return 0;
4308}
4309
4310static int
4311load_short_binstring(UnpicklerObject *self)
4312{
4313 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004314 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004315 char *s;
4316
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004317 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004318 return -1;
4319
4320 x = (unsigned char)s[0];
4321
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004322 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004323 return -1;
4324
4325 /* Convert Python 2.x strings to unicode. */
4326 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4327 if (str == NULL)
4328 return -1;
4329
4330 PDATA_PUSH(self->stack, str, -1);
4331 return 0;
4332}
4333
4334static int
4335load_unicode(UnpicklerObject *self)
4336{
4337 PyObject *str;
4338 Py_ssize_t len;
4339 char *s;
4340
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004341 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004342 return -1;
4343 if (len < 1)
4344 return bad_readline();
4345
4346 str = PyUnicode_DecodeRawUnicodeEscape(s, len - 1, NULL);
4347 if (str == NULL)
4348 return -1;
4349
4350 PDATA_PUSH(self->stack, str, -1);
4351 return 0;
4352}
4353
4354static int
4355load_binunicode(UnpicklerObject *self)
4356{
4357 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004358 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004359 char *s;
4360
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004361 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004362 return -1;
4363
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004364 size = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004365 if (size < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004366 PyErr_Format(PyExc_OverflowError,
4367 "BINUNICODE exceeds system's maximum size of %zd bytes",
Alexandre Vassalotticc757172013-04-14 02:25:10 -07004368 PY_SSIZE_T_MAX);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004369 return -1;
4370 }
4371
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004372
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004373 if (_Unpickler_Read(self, &s, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004374 return -1;
4375
Victor Stinner485fb562010-04-13 11:07:24 +00004376 str = PyUnicode_DecodeUTF8(s, size, "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004377 if (str == NULL)
4378 return -1;
4379
4380 PDATA_PUSH(self->stack, str, -1);
4381 return 0;
4382}
4383
4384static int
4385load_tuple(UnpicklerObject *self)
4386{
4387 PyObject *tuple;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004388 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004389
4390 if ((i = marker(self)) < 0)
4391 return -1;
4392
4393 tuple = Pdata_poptuple(self->stack, i);
4394 if (tuple == NULL)
4395 return -1;
4396 PDATA_PUSH(self->stack, tuple, -1);
4397 return 0;
4398}
4399
4400static int
4401load_counted_tuple(UnpicklerObject *self, int len)
4402{
4403 PyObject *tuple;
4404
4405 tuple = PyTuple_New(len);
4406 if (tuple == NULL)
4407 return -1;
4408
4409 while (--len >= 0) {
4410 PyObject *item;
4411
4412 PDATA_POP(self->stack, item);
4413 if (item == NULL)
4414 return -1;
4415 PyTuple_SET_ITEM(tuple, len, item);
4416 }
4417 PDATA_PUSH(self->stack, tuple, -1);
4418 return 0;
4419}
4420
4421static int
4422load_empty_list(UnpicklerObject *self)
4423{
4424 PyObject *list;
4425
4426 if ((list = PyList_New(0)) == NULL)
4427 return -1;
4428 PDATA_PUSH(self->stack, list, -1);
4429 return 0;
4430}
4431
4432static int
4433load_empty_dict(UnpicklerObject *self)
4434{
4435 PyObject *dict;
4436
4437 if ((dict = PyDict_New()) == NULL)
4438 return -1;
4439 PDATA_PUSH(self->stack, dict, -1);
4440 return 0;
4441}
4442
4443static int
4444load_list(UnpicklerObject *self)
4445{
4446 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004447 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004448
4449 if ((i = marker(self)) < 0)
4450 return -1;
4451
4452 list = Pdata_poplist(self->stack, i);
4453 if (list == NULL)
4454 return -1;
4455 PDATA_PUSH(self->stack, list, -1);
4456 return 0;
4457}
4458
4459static int
4460load_dict(UnpicklerObject *self)
4461{
4462 PyObject *dict, *key, *value;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004463 Py_ssize_t i, j, k;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004464
4465 if ((i = marker(self)) < 0)
4466 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004467 j = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004468
4469 if ((dict = PyDict_New()) == NULL)
4470 return -1;
4471
4472 for (k = i + 1; k < j; k += 2) {
4473 key = self->stack->data[k - 1];
4474 value = self->stack->data[k];
4475 if (PyDict_SetItem(dict, key, value) < 0) {
4476 Py_DECREF(dict);
4477 return -1;
4478 }
4479 }
4480 Pdata_clear(self->stack, i);
4481 PDATA_PUSH(self->stack, dict, -1);
4482 return 0;
4483}
4484
4485static PyObject *
4486instantiate(PyObject *cls, PyObject *args)
4487{
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004488 PyObject *result = NULL;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02004489 _Py_IDENTIFIER(__getinitargs__);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004490 /* Caller must assure args are a tuple. Normally, args come from
4491 Pdata_poptuple which packs objects from the top of the stack
4492 into a newly created tuple. */
4493 assert(PyTuple_Check(args));
4494 if (Py_SIZE(args) > 0 || !PyType_Check(cls) ||
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02004495 _PyObject_HasAttrId(cls, &PyId___getinitargs__)) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004496 result = PyObject_CallObject(cls, args);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004497 }
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004498 else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02004499 _Py_IDENTIFIER(__new__);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02004500
4501 result = _PyObject_CallMethodId(cls, &PyId___new__, "O", cls);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004502 }
4503 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004504}
4505
4506static int
4507load_obj(UnpicklerObject *self)
4508{
4509 PyObject *cls, *args, *obj = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004510 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004511
4512 if ((i = marker(self)) < 0)
4513 return -1;
4514
4515 args = Pdata_poptuple(self->stack, i + 1);
4516 if (args == NULL)
4517 return -1;
4518
4519 PDATA_POP(self->stack, cls);
4520 if (cls) {
4521 obj = instantiate(cls, args);
4522 Py_DECREF(cls);
4523 }
4524 Py_DECREF(args);
4525 if (obj == NULL)
4526 return -1;
4527
4528 PDATA_PUSH(self->stack, obj, -1);
4529 return 0;
4530}
4531
4532static int
4533load_inst(UnpicklerObject *self)
4534{
4535 PyObject *cls = NULL;
4536 PyObject *args = NULL;
4537 PyObject *obj = NULL;
4538 PyObject *module_name;
4539 PyObject *class_name;
4540 Py_ssize_t len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004541 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004542 char *s;
4543
4544 if ((i = marker(self)) < 0)
4545 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004546 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004547 return -1;
4548 if (len < 2)
4549 return bad_readline();
4550
4551 /* Here it is safe to use PyUnicode_DecodeASCII(), even though non-ASCII
4552 identifiers are permitted in Python 3.0, since the INST opcode is only
4553 supported by older protocols on Python 2.x. */
4554 module_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
4555 if (module_name == NULL)
4556 return -1;
4557
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004558 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004559 if (len < 2)
4560 return bad_readline();
4561 class_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004562 if (class_name != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004563 cls = find_class(self, module_name, class_name);
4564 Py_DECREF(class_name);
4565 }
4566 }
4567 Py_DECREF(module_name);
4568
4569 if (cls == NULL)
4570 return -1;
4571
4572 if ((args = Pdata_poptuple(self->stack, i)) != NULL) {
4573 obj = instantiate(cls, args);
4574 Py_DECREF(args);
4575 }
4576 Py_DECREF(cls);
4577
4578 if (obj == NULL)
4579 return -1;
4580
4581 PDATA_PUSH(self->stack, obj, -1);
4582 return 0;
4583}
4584
4585static int
4586load_newobj(UnpicklerObject *self)
4587{
4588 PyObject *args = NULL;
4589 PyObject *clsraw = NULL;
4590 PyTypeObject *cls; /* clsraw cast to its true type */
4591 PyObject *obj;
4592
4593 /* Stack is ... cls argtuple, and we want to call
4594 * cls.__new__(cls, *argtuple).
4595 */
4596 PDATA_POP(self->stack, args);
4597 if (args == NULL)
4598 goto error;
4599 if (!PyTuple_Check(args)) {
4600 PyErr_SetString(UnpicklingError, "NEWOBJ expected an arg " "tuple.");
4601 goto error;
4602 }
4603
4604 PDATA_POP(self->stack, clsraw);
4605 cls = (PyTypeObject *)clsraw;
4606 if (cls == NULL)
4607 goto error;
4608 if (!PyType_Check(cls)) {
4609 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4610 "isn't a type object");
4611 goto error;
4612 }
4613 if (cls->tp_new == NULL) {
4614 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4615 "has NULL tp_new");
4616 goto error;
4617 }
4618
4619 /* Call __new__. */
4620 obj = cls->tp_new(cls, args, NULL);
4621 if (obj == NULL)
4622 goto error;
4623
4624 Py_DECREF(args);
4625 Py_DECREF(clsraw);
4626 PDATA_PUSH(self->stack, obj, -1);
4627 return 0;
4628
4629 error:
4630 Py_XDECREF(args);
4631 Py_XDECREF(clsraw);
4632 return -1;
4633}
4634
4635static int
4636load_global(UnpicklerObject *self)
4637{
4638 PyObject *global = NULL;
4639 PyObject *module_name;
4640 PyObject *global_name;
4641 Py_ssize_t len;
4642 char *s;
4643
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004644 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004645 return -1;
4646 if (len < 2)
4647 return bad_readline();
4648 module_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4649 if (!module_name)
4650 return -1;
4651
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004652 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004653 if (len < 2) {
4654 Py_DECREF(module_name);
4655 return bad_readline();
4656 }
4657 global_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4658 if (global_name) {
4659 global = find_class(self, module_name, global_name);
4660 Py_DECREF(global_name);
4661 }
4662 }
4663 Py_DECREF(module_name);
4664
4665 if (global == NULL)
4666 return -1;
4667 PDATA_PUSH(self->stack, global, -1);
4668 return 0;
4669}
4670
4671static int
4672load_persid(UnpicklerObject *self)
4673{
4674 PyObject *pid;
4675 Py_ssize_t len;
4676 char *s;
4677
4678 if (self->pers_func) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004679 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004680 return -1;
4681 if (len < 2)
4682 return bad_readline();
4683
4684 pid = PyBytes_FromStringAndSize(s, len - 1);
4685 if (pid == NULL)
4686 return -1;
4687
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004688 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004689 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004690 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004691 if (pid == NULL)
4692 return -1;
4693
4694 PDATA_PUSH(self->stack, pid, -1);
4695 return 0;
4696 }
4697 else {
4698 PyErr_SetString(UnpicklingError,
4699 "A load persistent id instruction was encountered,\n"
4700 "but no persistent_load function was specified.");
4701 return -1;
4702 }
4703}
4704
4705static int
4706load_binpersid(UnpicklerObject *self)
4707{
4708 PyObject *pid;
4709
4710 if (self->pers_func) {
4711 PDATA_POP(self->stack, pid);
4712 if (pid == NULL)
4713 return -1;
4714
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004715 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004716 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004717 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004718 if (pid == NULL)
4719 return -1;
4720
4721 PDATA_PUSH(self->stack, pid, -1);
4722 return 0;
4723 }
4724 else {
4725 PyErr_SetString(UnpicklingError,
4726 "A load persistent id instruction was encountered,\n"
4727 "but no persistent_load function was specified.");
4728 return -1;
4729 }
4730}
4731
4732static int
4733load_pop(UnpicklerObject *self)
4734{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004735 Py_ssize_t len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004736
4737 /* Note that we split the (pickle.py) stack into two stacks,
4738 * an object stack and a mark stack. We have to be clever and
4739 * pop the right one. We do this by looking at the top of the
Collin Winter8ca69de2009-05-26 16:53:41 +00004740 * mark stack first, and only signalling a stack underflow if
4741 * the object stack is empty and the mark stack doesn't match
4742 * our expectations.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004743 */
Collin Winter8ca69de2009-05-26 16:53:41 +00004744 if (self->num_marks > 0 && self->marks[self->num_marks - 1] == len) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004745 self->num_marks--;
Antoine Pitrou01a15ea2010-01-07 17:57:31 +00004746 } else if (len > 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004747 len--;
4748 Py_DECREF(self->stack->data[len]);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004749 Py_SIZE(self->stack) = len;
Collin Winter8ca69de2009-05-26 16:53:41 +00004750 } else {
4751 return stack_underflow();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004752 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004753 return 0;
4754}
4755
4756static int
4757load_pop_mark(UnpicklerObject *self)
4758{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004759 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004760
4761 if ((i = marker(self)) < 0)
4762 return -1;
4763
4764 Pdata_clear(self->stack, i);
4765
4766 return 0;
4767}
4768
4769static int
4770load_dup(UnpicklerObject *self)
4771{
4772 PyObject *last;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004773 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004774
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004775 if ((len = Py_SIZE(self->stack)) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004776 return stack_underflow();
4777 last = self->stack->data[len - 1];
4778 PDATA_APPEND(self->stack, last, -1);
4779 return 0;
4780}
4781
4782static int
4783load_get(UnpicklerObject *self)
4784{
4785 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004786 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004787 Py_ssize_t len;
4788 char *s;
4789
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004790 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004791 return -1;
4792 if (len < 2)
4793 return bad_readline();
4794
4795 key = PyLong_FromString(s, NULL, 10);
4796 if (key == NULL)
4797 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004798 idx = PyLong_AsSsize_t(key);
4799 if (idx == -1 && PyErr_Occurred()) {
4800 Py_DECREF(key);
4801 return -1;
4802 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004803
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004804 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004805 if (value == NULL) {
4806 if (!PyErr_Occurred())
4807 PyErr_SetObject(PyExc_KeyError, key);
4808 Py_DECREF(key);
4809 return -1;
4810 }
4811 Py_DECREF(key);
4812
4813 PDATA_APPEND(self->stack, value, -1);
4814 return 0;
4815}
4816
4817static int
4818load_binget(UnpicklerObject *self)
4819{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004820 PyObject *value;
4821 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004822 char *s;
4823
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004824 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004825 return -1;
4826
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004827 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004828
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004829 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004830 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004831 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004832 if (!PyErr_Occurred())
4833 PyErr_SetObject(PyExc_KeyError, key);
4834 Py_DECREF(key);
4835 return -1;
4836 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004837
4838 PDATA_APPEND(self->stack, value, -1);
4839 return 0;
4840}
4841
4842static int
4843load_long_binget(UnpicklerObject *self)
4844{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004845 PyObject *value;
4846 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004847 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004848
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004849 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004850 return -1;
4851
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004852 idx = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004853
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004854 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004855 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004856 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004857 if (!PyErr_Occurred())
4858 PyErr_SetObject(PyExc_KeyError, key);
4859 Py_DECREF(key);
4860 return -1;
4861 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004862
4863 PDATA_APPEND(self->stack, value, -1);
4864 return 0;
4865}
4866
4867/* Push an object from the extension registry (EXT[124]). nbytes is
4868 * the number of bytes following the opcode, holding the index (code) value.
4869 */
4870static int
4871load_extension(UnpicklerObject *self, int nbytes)
4872{
4873 char *codebytes; /* the nbytes bytes after the opcode */
4874 long code; /* calc_binint returns long */
4875 PyObject *py_code; /* code as a Python int */
4876 PyObject *obj; /* the object to push */
4877 PyObject *pair; /* (module_name, class_name) */
4878 PyObject *module_name, *class_name;
4879
4880 assert(nbytes == 1 || nbytes == 2 || nbytes == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004881 if (_Unpickler_Read(self, &codebytes, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004882 return -1;
4883 code = calc_binint(codebytes, nbytes);
4884 if (code <= 0) { /* note that 0 is forbidden */
4885 /* Corrupt or hostile pickle. */
4886 PyErr_SetString(UnpicklingError, "EXT specifies code <= 0");
4887 return -1;
4888 }
4889
4890 /* Look for the code in the cache. */
4891 py_code = PyLong_FromLong(code);
4892 if (py_code == NULL)
4893 return -1;
4894 obj = PyDict_GetItem(extension_cache, py_code);
4895 if (obj != NULL) {
4896 /* Bingo. */
4897 Py_DECREF(py_code);
4898 PDATA_APPEND(self->stack, obj, -1);
4899 return 0;
4900 }
4901
4902 /* Look up the (module_name, class_name) pair. */
4903 pair = PyDict_GetItem(inverted_registry, py_code);
4904 if (pair == NULL) {
4905 Py_DECREF(py_code);
4906 PyErr_Format(PyExc_ValueError, "unregistered extension "
4907 "code %ld", code);
4908 return -1;
4909 }
4910 /* Since the extension registry is manipulable via Python code,
4911 * confirm that pair is really a 2-tuple of strings.
4912 */
4913 if (!PyTuple_Check(pair) || PyTuple_Size(pair) != 2 ||
4914 !PyUnicode_Check(module_name = PyTuple_GET_ITEM(pair, 0)) ||
4915 !PyUnicode_Check(class_name = PyTuple_GET_ITEM(pair, 1))) {
4916 Py_DECREF(py_code);
4917 PyErr_Format(PyExc_ValueError, "_inverted_registry[%ld] "
4918 "isn't a 2-tuple of strings", code);
4919 return -1;
4920 }
4921 /* Load the object. */
4922 obj = find_class(self, module_name, class_name);
4923 if (obj == NULL) {
4924 Py_DECREF(py_code);
4925 return -1;
4926 }
4927 /* Cache code -> obj. */
4928 code = PyDict_SetItem(extension_cache, py_code, obj);
4929 Py_DECREF(py_code);
4930 if (code < 0) {
4931 Py_DECREF(obj);
4932 return -1;
4933 }
4934 PDATA_PUSH(self->stack, obj, -1);
4935 return 0;
4936}
4937
4938static int
4939load_put(UnpicklerObject *self)
4940{
4941 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004942 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004943 Py_ssize_t len;
4944 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004945
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004946 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004947 return -1;
4948 if (len < 2)
4949 return bad_readline();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004950 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004951 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004952 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004953
4954 key = PyLong_FromString(s, NULL, 10);
4955 if (key == NULL)
4956 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004957 idx = PyLong_AsSsize_t(key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004958 Py_DECREF(key);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02004959 if (idx < 0) {
4960 if (!PyErr_Occurred())
4961 PyErr_SetString(PyExc_ValueError,
4962 "negative PUT argument");
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004963 return -1;
Antoine Pitrou55549ec2011-08-30 00:27:10 +02004964 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004965
4966 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004967}
4968
4969static int
4970load_binput(UnpicklerObject *self)
4971{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004972 PyObject *value;
4973 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004974 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004975
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004976 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004977 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004978
4979 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004980 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004981 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004982
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004983 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004984
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004985 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004986}
4987
4988static int
4989load_long_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, 4) < 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 Pitrou82be19f2011-08-29 23:09:33 +02005002 idx = calc_binsize(s, 4);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02005003 if (idx < 0) {
5004 PyErr_SetString(PyExc_ValueError,
5005 "negative LONG_BINPUT argument");
5006 return -1;
5007 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005008
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005009 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005010}
5011
5012static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005013do_append(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005014{
5015 PyObject *value;
5016 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005017 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005018
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005019 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005020 if (x > len || x <= 0)
5021 return stack_underflow();
5022 if (len == x) /* nothing to do */
5023 return 0;
5024
5025 list = self->stack->data[x - 1];
5026
5027 if (PyList_Check(list)) {
5028 PyObject *slice;
5029 Py_ssize_t list_len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005030 int ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005031
5032 slice = Pdata_poplist(self->stack, x);
5033 if (!slice)
5034 return -1;
5035 list_len = PyList_GET_SIZE(list);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005036 ret = PyList_SetSlice(list, list_len, list_len, slice);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005037 Py_DECREF(slice);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005038 return ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005039 }
5040 else {
5041 PyObject *append_func;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005042 _Py_IDENTIFIER(append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005043
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005044 append_func = _PyObject_GetAttrId(list, &PyId_append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005045 if (append_func == NULL)
5046 return -1;
5047 for (i = x; i < len; i++) {
5048 PyObject *result;
5049
5050 value = self->stack->data[i];
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005051 result = _Unpickler_FastCall(self, append_func, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005052 if (result == NULL) {
5053 Pdata_clear(self->stack, i + 1);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005054 Py_SIZE(self->stack) = x;
Alexandre Vassalotti637c7c42013-04-20 21:28:21 -07005055 Py_DECREF(append_func);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005056 return -1;
5057 }
5058 Py_DECREF(result);
5059 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005060 Py_SIZE(self->stack) = x;
Alexandre Vassalotti637c7c42013-04-20 21:28:21 -07005061 Py_DECREF(append_func);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005062 }
5063
5064 return 0;
5065}
5066
5067static int
5068load_append(UnpicklerObject *self)
5069{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005070 return do_append(self, Py_SIZE(self->stack) - 1);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005071}
5072
5073static int
5074load_appends(UnpicklerObject *self)
5075{
5076 return do_append(self, marker(self));
5077}
5078
5079static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005080do_setitems(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005081{
5082 PyObject *value, *key;
5083 PyObject *dict;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005084 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005085 int status = 0;
5086
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005087 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005088 if (x > len || x <= 0)
5089 return stack_underflow();
5090 if (len == x) /* nothing to do */
5091 return 0;
Victor Stinner121aab42011-09-29 23:40:53 +02005092 if ((len - x) % 2 != 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005093 /* Currupt or hostile pickle -- we never write one like this. */
5094 PyErr_SetString(UnpicklingError, "odd number of items for SETITEMS");
5095 return -1;
5096 }
5097
5098 /* Here, dict does not actually need to be a PyDict; it could be anything
5099 that supports the __setitem__ attribute. */
5100 dict = self->stack->data[x - 1];
5101
5102 for (i = x + 1; i < len; i += 2) {
5103 key = self->stack->data[i - 1];
5104 value = self->stack->data[i];
5105 if (PyObject_SetItem(dict, key, value) < 0) {
5106 status = -1;
5107 break;
5108 }
5109 }
5110
5111 Pdata_clear(self->stack, x);
5112 return status;
5113}
5114
5115static int
5116load_setitem(UnpicklerObject *self)
5117{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005118 return do_setitems(self, Py_SIZE(self->stack) - 2);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005119}
5120
5121static int
5122load_setitems(UnpicklerObject *self)
5123{
5124 return do_setitems(self, marker(self));
5125}
5126
5127static int
5128load_build(UnpicklerObject *self)
5129{
5130 PyObject *state, *inst, *slotstate;
5131 PyObject *setstate;
5132 int status = 0;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005133 _Py_IDENTIFIER(__setstate__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005134
5135 /* Stack is ... instance, state. We want to leave instance at
5136 * the stack top, possibly mutated via instance.__setstate__(state).
5137 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005138 if (Py_SIZE(self->stack) < 2)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005139 return stack_underflow();
5140
5141 PDATA_POP(self->stack, state);
5142 if (state == NULL)
5143 return -1;
5144
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005145 inst = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005146
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005147 setstate = _PyObject_GetAttrId(inst, &PyId___setstate__);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005148 if (setstate == NULL) {
5149 if (PyErr_ExceptionMatches(PyExc_AttributeError))
5150 PyErr_Clear();
Antoine Pitroud79dc622008-09-05 00:03:33 +00005151 else {
5152 Py_DECREF(state);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005153 return -1;
Antoine Pitroud79dc622008-09-05 00:03:33 +00005154 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005155 }
5156 else {
5157 PyObject *result;
5158
5159 /* The explicit __setstate__ is responsible for everything. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005160 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Antoine Pitroud79dc622008-09-05 00:03:33 +00005161 reference to state first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005162 result = _Unpickler_FastCall(self, setstate, state);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005163 Py_DECREF(setstate);
5164 if (result == NULL)
5165 return -1;
5166 Py_DECREF(result);
5167 return 0;
5168 }
5169
5170 /* A default __setstate__. First see whether state embeds a
5171 * slot state dict too (a proto 2 addition).
5172 */
5173 if (PyTuple_Check(state) && Py_SIZE(state) == 2) {
5174 PyObject *tmp = state;
5175
5176 state = PyTuple_GET_ITEM(tmp, 0);
5177 slotstate = PyTuple_GET_ITEM(tmp, 1);
5178 Py_INCREF(state);
5179 Py_INCREF(slotstate);
5180 Py_DECREF(tmp);
5181 }
5182 else
5183 slotstate = NULL;
5184
5185 /* Set inst.__dict__ from the state dict (if any). */
5186 if (state != Py_None) {
5187 PyObject *dict;
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005188 PyObject *d_key, *d_value;
5189 Py_ssize_t i;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005190 _Py_IDENTIFIER(__dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005191
5192 if (!PyDict_Check(state)) {
5193 PyErr_SetString(UnpicklingError, "state is not a dictionary");
5194 goto error;
5195 }
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005196 dict = _PyObject_GetAttrId(inst, &PyId___dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005197 if (dict == NULL)
5198 goto error;
5199
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005200 i = 0;
5201 while (PyDict_Next(state, &i, &d_key, &d_value)) {
5202 /* normally the keys for instance attributes are
5203 interned. we should try to do that here. */
5204 Py_INCREF(d_key);
5205 if (PyUnicode_CheckExact(d_key))
5206 PyUnicode_InternInPlace(&d_key);
5207 if (PyObject_SetItem(dict, d_key, d_value) < 0) {
5208 Py_DECREF(d_key);
5209 goto error;
5210 }
5211 Py_DECREF(d_key);
5212 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005213 Py_DECREF(dict);
5214 }
5215
5216 /* Also set instance attributes from the slotstate dict (if any). */
5217 if (slotstate != NULL) {
5218 PyObject *d_key, *d_value;
5219 Py_ssize_t i;
5220
5221 if (!PyDict_Check(slotstate)) {
5222 PyErr_SetString(UnpicklingError,
5223 "slot state is not a dictionary");
5224 goto error;
5225 }
5226 i = 0;
5227 while (PyDict_Next(slotstate, &i, &d_key, &d_value)) {
5228 if (PyObject_SetAttr(inst, d_key, d_value) < 0)
5229 goto error;
5230 }
5231 }
5232
5233 if (0) {
5234 error:
5235 status = -1;
5236 }
5237
5238 Py_DECREF(state);
5239 Py_XDECREF(slotstate);
5240 return status;
5241}
5242
5243static int
5244load_mark(UnpicklerObject *self)
5245{
5246
5247 /* Note that we split the (pickle.py) stack into two stacks, an
5248 * object stack and a mark stack. Here we push a mark onto the
5249 * mark stack.
5250 */
5251
5252 if ((self->num_marks + 1) >= self->marks_size) {
5253 size_t alloc;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005254 Py_ssize_t *marks;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005255
5256 /* Use the size_t type to check for overflow. */
5257 alloc = ((size_t)self->num_marks << 1) + 20;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005258 if (alloc > (PY_SSIZE_T_MAX / sizeof(Py_ssize_t)) ||
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00005259 alloc <= ((size_t)self->num_marks + 1)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005260 PyErr_NoMemory();
5261 return -1;
5262 }
5263
5264 if (self->marks == NULL)
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005265 marks = (Py_ssize_t *) PyMem_Malloc(alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005266 else
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005267 marks = (Py_ssize_t *) PyMem_Realloc(self->marks,
5268 alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005269 if (marks == NULL) {
5270 PyErr_NoMemory();
5271 return -1;
5272 }
5273 self->marks = marks;
5274 self->marks_size = (Py_ssize_t)alloc;
5275 }
5276
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005277 self->marks[self->num_marks++] = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005278
5279 return 0;
5280}
5281
5282static int
5283load_reduce(UnpicklerObject *self)
5284{
5285 PyObject *callable = NULL;
5286 PyObject *argtup = NULL;
5287 PyObject *obj = NULL;
5288
5289 PDATA_POP(self->stack, argtup);
5290 if (argtup == NULL)
5291 return -1;
5292 PDATA_POP(self->stack, callable);
5293 if (callable) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005294 obj = PyObject_CallObject(callable, argtup);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005295 Py_DECREF(callable);
5296 }
5297 Py_DECREF(argtup);
5298
5299 if (obj == NULL)
5300 return -1;
5301
5302 PDATA_PUSH(self->stack, obj, -1);
5303 return 0;
5304}
5305
5306/* Just raises an error if we don't know the protocol specified. PROTO
5307 * is the first opcode for protocols >= 2.
5308 */
5309static int
5310load_proto(UnpicklerObject *self)
5311{
5312 char *s;
5313 int i;
5314
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005315 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005316 return -1;
5317
5318 i = (unsigned char)s[0];
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005319 if (i <= HIGHEST_PROTOCOL) {
5320 self->proto = i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005321 return 0;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005322 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005323
5324 PyErr_Format(PyExc_ValueError, "unsupported pickle protocol: %d", i);
5325 return -1;
5326}
5327
5328static PyObject *
5329load(UnpicklerObject *self)
5330{
5331 PyObject *err;
5332 PyObject *value = NULL;
5333 char *s;
5334
5335 self->num_marks = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005336 if (Py_SIZE(self->stack))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005337 Pdata_clear(self->stack, 0);
5338
5339 /* Convenient macros for the dispatch while-switch loop just below. */
5340#define OP(opcode, load_func) \
5341 case opcode: if (load_func(self) < 0) break; continue;
5342
5343#define OP_ARG(opcode, load_func, arg) \
5344 case opcode: if (load_func(self, (arg)) < 0) break; continue;
5345
5346 while (1) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005347 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005348 break;
5349
5350 switch ((enum opcode)s[0]) {
5351 OP(NONE, load_none)
5352 OP(BININT, load_binint)
5353 OP(BININT1, load_binint1)
5354 OP(BININT2, load_binint2)
5355 OP(INT, load_int)
5356 OP(LONG, load_long)
5357 OP_ARG(LONG1, load_counted_long, 1)
5358 OP_ARG(LONG4, load_counted_long, 4)
5359 OP(FLOAT, load_float)
5360 OP(BINFLOAT, load_binfloat)
5361 OP(BINBYTES, load_binbytes)
5362 OP(SHORT_BINBYTES, load_short_binbytes)
5363 OP(BINSTRING, load_binstring)
5364 OP(SHORT_BINSTRING, load_short_binstring)
5365 OP(STRING, load_string)
5366 OP(UNICODE, load_unicode)
5367 OP(BINUNICODE, load_binunicode)
5368 OP_ARG(EMPTY_TUPLE, load_counted_tuple, 0)
5369 OP_ARG(TUPLE1, load_counted_tuple, 1)
5370 OP_ARG(TUPLE2, load_counted_tuple, 2)
5371 OP_ARG(TUPLE3, load_counted_tuple, 3)
5372 OP(TUPLE, load_tuple)
5373 OP(EMPTY_LIST, load_empty_list)
5374 OP(LIST, load_list)
5375 OP(EMPTY_DICT, load_empty_dict)
5376 OP(DICT, load_dict)
5377 OP(OBJ, load_obj)
5378 OP(INST, load_inst)
5379 OP(NEWOBJ, load_newobj)
5380 OP(GLOBAL, load_global)
5381 OP(APPEND, load_append)
5382 OP(APPENDS, load_appends)
5383 OP(BUILD, load_build)
5384 OP(DUP, load_dup)
5385 OP(BINGET, load_binget)
5386 OP(LONG_BINGET, load_long_binget)
5387 OP(GET, load_get)
5388 OP(MARK, load_mark)
5389 OP(BINPUT, load_binput)
5390 OP(LONG_BINPUT, load_long_binput)
5391 OP(PUT, load_put)
5392 OP(POP, load_pop)
5393 OP(POP_MARK, load_pop_mark)
5394 OP(SETITEM, load_setitem)
5395 OP(SETITEMS, load_setitems)
5396 OP(PERSID, load_persid)
5397 OP(BINPERSID, load_binpersid)
5398 OP(REDUCE, load_reduce)
5399 OP(PROTO, load_proto)
5400 OP_ARG(EXT1, load_extension, 1)
5401 OP_ARG(EXT2, load_extension, 2)
5402 OP_ARG(EXT4, load_extension, 4)
5403 OP_ARG(NEWTRUE, load_bool, Py_True)
5404 OP_ARG(NEWFALSE, load_bool, Py_False)
5405
5406 case STOP:
5407 break;
5408
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005409 default:
Benjamin Petersonadde86d2011-09-23 13:41:41 -04005410 if (s[0] == '\0')
5411 PyErr_SetNone(PyExc_EOFError);
5412 else
5413 PyErr_Format(UnpicklingError,
5414 "invalid load key, '%c'.", s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005415 return NULL;
5416 }
5417
5418 break; /* and we are done! */
5419 }
5420
Antoine Pitrou04248a82010-10-12 20:51:21 +00005421 if (_Unpickler_SkipConsumed(self) < 0)
5422 return NULL;
5423
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005424 /* XXX: It is not clear what this is actually for. */
5425 if ((err = PyErr_Occurred())) {
5426 if (err == PyExc_EOFError) {
5427 PyErr_SetNone(PyExc_EOFError);
5428 }
5429 return NULL;
5430 }
5431
5432 PDATA_POP(self->stack, value);
5433 return value;
5434}
5435
5436PyDoc_STRVAR(Unpickler_load_doc,
5437"load() -> object. Load a pickle."
5438"\n"
5439"Read a pickled object representation from the open file object given in\n"
5440"the constructor, and return the reconstituted object hierarchy specified\n"
5441"therein.\n");
5442
5443static PyObject *
5444Unpickler_load(UnpicklerObject *self)
5445{
5446 /* Check whether the Unpickler was initialized correctly. This prevents
5447 segfaulting if a subclass overridden __init__ with a function that does
5448 not call Unpickler.__init__(). Here, we simply ensure that self->read
5449 is not NULL. */
5450 if (self->read == NULL) {
Victor Stinner121aab42011-09-29 23:40:53 +02005451 PyErr_Format(UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005452 "Unpickler.__init__() was not called by %s.__init__()",
5453 Py_TYPE(self)->tp_name);
5454 return NULL;
5455 }
5456
5457 return load(self);
5458}
5459
5460/* The name of find_class() is misleading. In newer pickle protocols, this
5461 function is used for loading any global (i.e., functions), not just
5462 classes. The name is kept only for backward compatibility. */
5463
5464PyDoc_STRVAR(Unpickler_find_class_doc,
5465"find_class(module_name, global_name) -> object.\n"
5466"\n"
5467"Return an object from a specified module, importing the module if\n"
5468"necessary. Subclasses may override this method (e.g. to restrict\n"
5469"unpickling of arbitrary classes and functions).\n"
5470"\n"
5471"This method is called whenever a class or a function object is\n"
5472"needed. Both arguments passed are str objects.\n");
5473
5474static PyObject *
5475Unpickler_find_class(UnpicklerObject *self, PyObject *args)
5476{
5477 PyObject *global;
5478 PyObject *modules_dict;
5479 PyObject *module;
5480 PyObject *module_name, *global_name;
5481
5482 if (!PyArg_UnpackTuple(args, "find_class", 2, 2,
5483 &module_name, &global_name))
5484 return NULL;
5485
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005486 /* Try to map the old names used in Python 2.x to the new ones used in
5487 Python 3.x. We do this only with old pickle protocols and when the
5488 user has not disabled the feature. */
5489 if (self->proto < 3 && self->fix_imports) {
5490 PyObject *key;
5491 PyObject *item;
5492
5493 /* Check if the global (i.e., a function or a class) was renamed
5494 or moved to another module. */
5495 key = PyTuple_Pack(2, module_name, global_name);
5496 if (key == NULL)
5497 return NULL;
5498 item = PyDict_GetItemWithError(name_mapping_2to3, key);
5499 Py_DECREF(key);
5500 if (item) {
5501 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
5502 PyErr_Format(PyExc_RuntimeError,
5503 "_compat_pickle.NAME_MAPPING values should be "
5504 "2-tuples, not %.200s", Py_TYPE(item)->tp_name);
5505 return NULL;
5506 }
5507 module_name = PyTuple_GET_ITEM(item, 0);
5508 global_name = PyTuple_GET_ITEM(item, 1);
5509 if (!PyUnicode_Check(module_name) ||
5510 !PyUnicode_Check(global_name)) {
5511 PyErr_Format(PyExc_RuntimeError,
5512 "_compat_pickle.NAME_MAPPING values should be "
5513 "pairs of str, not (%.200s, %.200s)",
5514 Py_TYPE(module_name)->tp_name,
5515 Py_TYPE(global_name)->tp_name);
5516 return NULL;
5517 }
5518 }
5519 else if (PyErr_Occurred()) {
5520 return NULL;
5521 }
5522
5523 /* Check if the module was renamed. */
5524 item = PyDict_GetItemWithError(import_mapping_2to3, module_name);
5525 if (item) {
5526 if (!PyUnicode_Check(item)) {
5527 PyErr_Format(PyExc_RuntimeError,
5528 "_compat_pickle.IMPORT_MAPPING values should be "
5529 "strings, not %.200s", Py_TYPE(item)->tp_name);
5530 return NULL;
5531 }
5532 module_name = item;
5533 }
5534 else if (PyErr_Occurred()) {
5535 return NULL;
5536 }
5537 }
5538
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005539 modules_dict = PySys_GetObject("modules");
5540 if (modules_dict == NULL)
5541 return NULL;
5542
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005543 module = PyDict_GetItemWithError(modules_dict, module_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005544 if (module == NULL) {
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005545 if (PyErr_Occurred())
5546 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005547 module = PyImport_Import(module_name);
5548 if (module == NULL)
5549 return NULL;
5550 global = PyObject_GetAttr(module, global_name);
5551 Py_DECREF(module);
5552 }
Victor Stinner121aab42011-09-29 23:40:53 +02005553 else {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005554 global = PyObject_GetAttr(module, global_name);
5555 }
5556 return global;
5557}
5558
5559static struct PyMethodDef Unpickler_methods[] = {
5560 {"load", (PyCFunction)Unpickler_load, METH_NOARGS,
5561 Unpickler_load_doc},
5562 {"find_class", (PyCFunction)Unpickler_find_class, METH_VARARGS,
5563 Unpickler_find_class_doc},
5564 {NULL, NULL} /* sentinel */
5565};
5566
5567static void
5568Unpickler_dealloc(UnpicklerObject *self)
5569{
5570 PyObject_GC_UnTrack((PyObject *)self);
5571 Py_XDECREF(self->readline);
5572 Py_XDECREF(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005573 Py_XDECREF(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005574 Py_XDECREF(self->stack);
5575 Py_XDECREF(self->pers_func);
5576 Py_XDECREF(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005577 if (self->buffer.buf != NULL) {
5578 PyBuffer_Release(&self->buffer);
5579 self->buffer.buf = NULL;
5580 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005581
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005582 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005583 PyMem_Free(self->marks);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005584 PyMem_Free(self->input_line);
Victor Stinner49fc8ec2013-07-07 23:30:24 +02005585 PyMem_Free(self->encoding);
5586 PyMem_Free(self->errors);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005587
5588 Py_TYPE(self)->tp_free((PyObject *)self);
5589}
5590
5591static int
5592Unpickler_traverse(UnpicklerObject *self, visitproc visit, void *arg)
5593{
5594 Py_VISIT(self->readline);
5595 Py_VISIT(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005596 Py_VISIT(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005597 Py_VISIT(self->stack);
5598 Py_VISIT(self->pers_func);
5599 Py_VISIT(self->arg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005600 return 0;
5601}
5602
5603static int
5604Unpickler_clear(UnpicklerObject *self)
5605{
5606 Py_CLEAR(self->readline);
5607 Py_CLEAR(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005608 Py_CLEAR(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005609 Py_CLEAR(self->stack);
5610 Py_CLEAR(self->pers_func);
5611 Py_CLEAR(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005612 if (self->buffer.buf != NULL) {
5613 PyBuffer_Release(&self->buffer);
5614 self->buffer.buf = NULL;
5615 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005616
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005617 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005618 PyMem_Free(self->marks);
5619 self->marks = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005620 PyMem_Free(self->input_line);
5621 self->input_line = NULL;
Victor Stinner49fc8ec2013-07-07 23:30:24 +02005622 PyMem_Free(self->encoding);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005623 self->encoding = NULL;
Victor Stinner49fc8ec2013-07-07 23:30:24 +02005624 PyMem_Free(self->errors);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005625 self->errors = NULL;
5626
5627 return 0;
5628}
5629
5630PyDoc_STRVAR(Unpickler_doc,
5631"Unpickler(file, *, encoding='ASCII', errors='strict')"
5632"\n"
5633"This takes a binary file for reading a pickle data stream.\n"
5634"\n"
5635"The protocol version of the pickle is detected automatically, so no\n"
5636"proto argument is needed.\n"
5637"\n"
5638"The file-like object must have two methods, a read() method\n"
5639"that takes an integer argument, and a readline() method that\n"
5640"requires no arguments. Both methods should return bytes.\n"
5641"Thus file-like object can be a binary file object opened for\n"
5642"reading, a BytesIO object, or any other custom object that\n"
5643"meets this interface.\n"
5644"\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005645"Optional keyword arguments are *fix_imports*, *encoding* and *errors*,\n"
5646"which are used to control compatiblity support for pickle stream\n"
5647"generated by Python 2.x. If *fix_imports* is True, pickle will try to\n"
5648"map the old Python 2.x names to the new names used in Python 3.x. The\n"
5649"*encoding* and *errors* tell pickle how to decode 8-bit string\n"
5650"instances pickled by Python 2.x; these default to 'ASCII' and\n"
5651"'strict', respectively.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005652
5653static int
5654Unpickler_init(UnpicklerObject *self, PyObject *args, PyObject *kwds)
5655{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005656 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005657 PyObject *file;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005658 PyObject *fix_imports = Py_True;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005659 char *encoding = NULL;
5660 char *errors = NULL;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005661 _Py_IDENTIFIER(persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005662
5663 /* XXX: That is an horrible error message. But, I don't know how to do
5664 better... */
5665 if (Py_SIZE(args) != 1) {
5666 PyErr_Format(PyExc_TypeError,
5667 "%s takes exactly one positional argument (%zd given)",
5668 Py_TYPE(self)->tp_name, Py_SIZE(args));
5669 return -1;
5670 }
5671
5672 /* Arguments parsing needs to be done in the __init__() method to allow
5673 subclasses to define their own __init__() method, which may (or may
5674 not) support Unpickler arguments. However, this means we need to be
5675 extra careful in the other Unpickler methods, since a subclass could
5676 forget to call Unpickler.__init__() thus breaking our internal
5677 invariants. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005678 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:Unpickler", kwlist,
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005679 &file, &fix_imports, &encoding, &errors))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005680 return -1;
5681
5682 /* In case of multiple __init__() calls, clear previous content. */
5683 if (self->read != NULL)
5684 (void)Unpickler_clear(self);
5685
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005686 if (_Unpickler_SetInputStream(self, file) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005687 return -1;
5688
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005689 if (_Unpickler_SetInputEncoding(self, encoding, errors) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005690 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005691
5692 self->fix_imports = PyObject_IsTrue(fix_imports);
5693 if (self->fix_imports == -1)
5694 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005695
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005696 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_load)) {
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005697 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
5698 &PyId_persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005699 if (self->pers_func == NULL)
5700 return -1;
5701 }
5702 else {
5703 self->pers_func = NULL;
5704 }
5705
5706 self->stack = (Pdata *)Pdata_New();
5707 if (self->stack == NULL)
5708 return -1;
5709
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005710 self->memo_size = 32;
5711 self->memo = _Unpickler_NewMemo(self->memo_size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005712 if (self->memo == NULL)
5713 return -1;
5714
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005715 self->arg = NULL;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005716 self->proto = 0;
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005717
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005718 return 0;
5719}
5720
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005721/* Define a proxy object for the Unpickler's internal memo object. This is to
5722 * avoid breaking code like:
5723 * unpickler.memo.clear()
5724 * and
5725 * unpickler.memo = saved_memo
5726 * Is this a good idea? Not really, but we don't want to break code that uses
5727 * it. Note that we don't implement the entire mapping API here. This is
5728 * intentional, as these should be treated as black-box implementation details.
5729 *
5730 * We do, however, have to implement pickling/unpickling support because of
Victor Stinner121aab42011-09-29 23:40:53 +02005731 * real-world code like cvs2svn.
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005732 */
5733
5734typedef struct {
5735 PyObject_HEAD
5736 UnpicklerObject *unpickler;
5737} UnpicklerMemoProxyObject;
5738
5739PyDoc_STRVAR(ump_clear_doc,
5740"memo.clear() -> None. Remove all items from memo.");
5741
5742static PyObject *
5743ump_clear(UnpicklerMemoProxyObject *self)
5744{
5745 _Unpickler_MemoCleanup(self->unpickler);
5746 self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size);
5747 if (self->unpickler->memo == NULL)
5748 return NULL;
5749 Py_RETURN_NONE;
5750}
5751
5752PyDoc_STRVAR(ump_copy_doc,
5753"memo.copy() -> new_memo. Copy the memo to a new object.");
5754
5755static PyObject *
5756ump_copy(UnpicklerMemoProxyObject *self)
5757{
5758 Py_ssize_t i;
5759 PyObject *new_memo = PyDict_New();
5760 if (new_memo == NULL)
5761 return NULL;
5762
5763 for (i = 0; i < self->unpickler->memo_size; i++) {
5764 int status;
5765 PyObject *key, *value;
5766
5767 value = self->unpickler->memo[i];
5768 if (value == NULL)
5769 continue;
5770
5771 key = PyLong_FromSsize_t(i);
5772 if (key == NULL)
5773 goto error;
5774 status = PyDict_SetItem(new_memo, key, value);
5775 Py_DECREF(key);
5776 if (status < 0)
5777 goto error;
5778 }
5779 return new_memo;
5780
5781error:
5782 Py_DECREF(new_memo);
5783 return NULL;
5784}
5785
5786PyDoc_STRVAR(ump_reduce_doc,
5787"memo.__reduce__(). Pickling support.");
5788
5789static PyObject *
5790ump_reduce(UnpicklerMemoProxyObject *self, PyObject *args)
5791{
5792 PyObject *reduce_value;
5793 PyObject *constructor_args;
5794 PyObject *contents = ump_copy(self);
5795 if (contents == NULL)
5796 return NULL;
5797
5798 reduce_value = PyTuple_New(2);
5799 if (reduce_value == NULL) {
5800 Py_DECREF(contents);
5801 return NULL;
5802 }
5803 constructor_args = PyTuple_New(1);
5804 if (constructor_args == NULL) {
5805 Py_DECREF(contents);
5806 Py_DECREF(reduce_value);
5807 return NULL;
5808 }
5809 PyTuple_SET_ITEM(constructor_args, 0, contents);
5810 Py_INCREF((PyObject *)&PyDict_Type);
5811 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
5812 PyTuple_SET_ITEM(reduce_value, 1, constructor_args);
5813 return reduce_value;
5814}
5815
5816static PyMethodDef unpicklerproxy_methods[] = {
5817 {"clear", (PyCFunction)ump_clear, METH_NOARGS, ump_clear_doc},
5818 {"copy", (PyCFunction)ump_copy, METH_NOARGS, ump_copy_doc},
5819 {"__reduce__", (PyCFunction)ump_reduce, METH_VARARGS, ump_reduce_doc},
5820 {NULL, NULL} /* sentinel */
5821};
5822
5823static void
5824UnpicklerMemoProxy_dealloc(UnpicklerMemoProxyObject *self)
5825{
5826 PyObject_GC_UnTrack(self);
5827 Py_XDECREF(self->unpickler);
5828 PyObject_GC_Del((PyObject *)self);
5829}
5830
5831static int
5832UnpicklerMemoProxy_traverse(UnpicklerMemoProxyObject *self,
5833 visitproc visit, void *arg)
5834{
5835 Py_VISIT(self->unpickler);
5836 return 0;
5837}
5838
5839static int
5840UnpicklerMemoProxy_clear(UnpicklerMemoProxyObject *self)
5841{
5842 Py_CLEAR(self->unpickler);
5843 return 0;
5844}
5845
5846static PyTypeObject UnpicklerMemoProxyType = {
5847 PyVarObject_HEAD_INIT(NULL, 0)
5848 "_pickle.UnpicklerMemoProxy", /*tp_name*/
5849 sizeof(UnpicklerMemoProxyObject), /*tp_basicsize*/
5850 0,
5851 (destructor)UnpicklerMemoProxy_dealloc, /* tp_dealloc */
5852 0, /* tp_print */
5853 0, /* tp_getattr */
5854 0, /* tp_setattr */
5855 0, /* tp_compare */
5856 0, /* tp_repr */
5857 0, /* tp_as_number */
5858 0, /* tp_as_sequence */
5859 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00005860 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005861 0, /* tp_call */
5862 0, /* tp_str */
5863 PyObject_GenericGetAttr, /* tp_getattro */
5864 PyObject_GenericSetAttr, /* tp_setattro */
5865 0, /* tp_as_buffer */
5866 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5867 0, /* tp_doc */
5868 (traverseproc)UnpicklerMemoProxy_traverse, /* tp_traverse */
5869 (inquiry)UnpicklerMemoProxy_clear, /* tp_clear */
5870 0, /* tp_richcompare */
5871 0, /* tp_weaklistoffset */
5872 0, /* tp_iter */
5873 0, /* tp_iternext */
5874 unpicklerproxy_methods, /* tp_methods */
5875};
5876
5877static PyObject *
5878UnpicklerMemoProxy_New(UnpicklerObject *unpickler)
5879{
5880 UnpicklerMemoProxyObject *self;
5881
5882 self = PyObject_GC_New(UnpicklerMemoProxyObject,
5883 &UnpicklerMemoProxyType);
5884 if (self == NULL)
5885 return NULL;
5886 Py_INCREF(unpickler);
5887 self->unpickler = unpickler;
5888 PyObject_GC_Track(self);
5889 return (PyObject *)self;
5890}
5891
5892/*****************************************************************************/
5893
5894
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005895static PyObject *
5896Unpickler_get_memo(UnpicklerObject *self)
5897{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005898 return UnpicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005899}
5900
5901static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005902Unpickler_set_memo(UnpicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005903{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005904 PyObject **new_memo;
5905 Py_ssize_t new_memo_size = 0;
5906 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005907
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005908 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005909 PyErr_SetString(PyExc_TypeError,
5910 "attribute deletion is not supported");
5911 return -1;
5912 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005913
5914 if (Py_TYPE(obj) == &UnpicklerMemoProxyType) {
5915 UnpicklerObject *unpickler =
5916 ((UnpicklerMemoProxyObject *)obj)->unpickler;
5917
5918 new_memo_size = unpickler->memo_size;
5919 new_memo = _Unpickler_NewMemo(new_memo_size);
5920 if (new_memo == NULL)
5921 return -1;
5922
5923 for (i = 0; i < new_memo_size; i++) {
5924 Py_XINCREF(unpickler->memo[i]);
5925 new_memo[i] = unpickler->memo[i];
5926 }
5927 }
5928 else if (PyDict_Check(obj)) {
5929 Py_ssize_t i = 0;
5930 PyObject *key, *value;
5931
5932 new_memo_size = PyDict_Size(obj);
5933 new_memo = _Unpickler_NewMemo(new_memo_size);
5934 if (new_memo == NULL)
5935 return -1;
5936
5937 while (PyDict_Next(obj, &i, &key, &value)) {
5938 Py_ssize_t idx;
5939 if (!PyLong_Check(key)) {
5940 PyErr_SetString(PyExc_TypeError,
5941 "memo key must be integers");
5942 goto error;
5943 }
5944 idx = PyLong_AsSsize_t(key);
5945 if (idx == -1 && PyErr_Occurred())
5946 goto error;
Christian Heimesa24b4d22013-07-01 15:17:45 +02005947 if (idx < 0) {
5948 PyErr_SetString(PyExc_ValueError,
Christian Heimes80878792013-07-01 15:23:39 +02005949 "memo key must be positive integers.");
Christian Heimesa24b4d22013-07-01 15:17:45 +02005950 goto error;
5951 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005952 if (_Unpickler_MemoPut(self, idx, value) < 0)
5953 goto error;
5954 }
5955 }
5956 else {
5957 PyErr_Format(PyExc_TypeError,
5958 "'memo' attribute must be an UnpicklerMemoProxy object"
5959 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005960 return -1;
5961 }
5962
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005963 _Unpickler_MemoCleanup(self);
5964 self->memo_size = new_memo_size;
5965 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005966
5967 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005968
5969 error:
5970 if (new_memo_size) {
5971 i = new_memo_size;
5972 while (--i >= 0) {
5973 Py_XDECREF(new_memo[i]);
5974 }
5975 PyMem_FREE(new_memo);
5976 }
5977 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005978}
5979
5980static PyObject *
5981Unpickler_get_persload(UnpicklerObject *self)
5982{
5983 if (self->pers_func == NULL)
5984 PyErr_SetString(PyExc_AttributeError, "persistent_load");
5985 else
5986 Py_INCREF(self->pers_func);
5987 return self->pers_func;
5988}
5989
5990static int
5991Unpickler_set_persload(UnpicklerObject *self, PyObject *value)
5992{
5993 PyObject *tmp;
5994
5995 if (value == NULL) {
5996 PyErr_SetString(PyExc_TypeError,
5997 "attribute deletion is not supported");
5998 return -1;
5999 }
6000 if (!PyCallable_Check(value)) {
6001 PyErr_SetString(PyExc_TypeError,
6002 "persistent_load must be a callable taking "
6003 "one argument");
6004 return -1;
6005 }
6006
6007 tmp = self->pers_func;
6008 Py_INCREF(value);
6009 self->pers_func = value;
6010 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
6011
6012 return 0;
6013}
6014
6015static PyGetSetDef Unpickler_getsets[] = {
6016 {"memo", (getter)Unpickler_get_memo, (setter)Unpickler_set_memo},
6017 {"persistent_load", (getter)Unpickler_get_persload,
6018 (setter)Unpickler_set_persload},
6019 {NULL}
6020};
6021
6022static PyTypeObject Unpickler_Type = {
6023 PyVarObject_HEAD_INIT(NULL, 0)
6024 "_pickle.Unpickler", /*tp_name*/
6025 sizeof(UnpicklerObject), /*tp_basicsize*/
6026 0, /*tp_itemsize*/
6027 (destructor)Unpickler_dealloc, /*tp_dealloc*/
6028 0, /*tp_print*/
6029 0, /*tp_getattr*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006030 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00006031 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006032 0, /*tp_repr*/
6033 0, /*tp_as_number*/
6034 0, /*tp_as_sequence*/
6035 0, /*tp_as_mapping*/
6036 0, /*tp_hash*/
6037 0, /*tp_call*/
6038 0, /*tp_str*/
6039 0, /*tp_getattro*/
6040 0, /*tp_setattro*/
6041 0, /*tp_as_buffer*/
6042 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
6043 Unpickler_doc, /*tp_doc*/
6044 (traverseproc)Unpickler_traverse, /*tp_traverse*/
6045 (inquiry)Unpickler_clear, /*tp_clear*/
6046 0, /*tp_richcompare*/
6047 0, /*tp_weaklistoffset*/
6048 0, /*tp_iter*/
6049 0, /*tp_iternext*/
6050 Unpickler_methods, /*tp_methods*/
6051 0, /*tp_members*/
6052 Unpickler_getsets, /*tp_getset*/
6053 0, /*tp_base*/
6054 0, /*tp_dict*/
6055 0, /*tp_descr_get*/
6056 0, /*tp_descr_set*/
6057 0, /*tp_dictoffset*/
6058 (initproc)Unpickler_init, /*tp_init*/
6059 PyType_GenericAlloc, /*tp_alloc*/
6060 PyType_GenericNew, /*tp_new*/
6061 PyObject_GC_Del, /*tp_free*/
6062 0, /*tp_is_gc*/
6063};
6064
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006065PyDoc_STRVAR(pickle_dump_doc,
6066"dump(obj, file, protocol=None, *, fix_imports=True) -> None\n"
6067"\n"
6068"Write a pickled representation of obj to the open file object file. This\n"
6069"is equivalent to ``Pickler(file, protocol).dump(obj)``, but may be more\n"
6070"efficient.\n"
6071"\n"
6072"The optional protocol argument tells the pickler to use the given protocol;\n"
6073"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
6074"backward-incompatible protocol designed for Python 3.0.\n"
6075"\n"
6076"Specifying a negative protocol version selects the highest protocol version\n"
6077"supported. The higher the protocol used, the more recent the version of\n"
6078"Python needed to read the pickle produced.\n"
6079"\n"
6080"The file argument must have a write() method that accepts a single bytes\n"
6081"argument. It can thus be a file object opened for binary writing, a\n"
6082"io.BytesIO instance, or any other custom object that meets this interface.\n"
6083"\n"
6084"If fix_imports is True and protocol is less than 3, pickle will try to\n"
6085"map the new Python 3.x names to the old module names used in Python 2.x,\n"
6086"so that the pickle data stream is readable with Python 2.x.\n");
6087
6088static PyObject *
6089pickle_dump(PyObject *self, PyObject *args, PyObject *kwds)
6090{
6091 static char *kwlist[] = {"obj", "file", "protocol", "fix_imports", 0};
6092 PyObject *obj;
6093 PyObject *file;
6094 PyObject *proto = NULL;
6095 PyObject *fix_imports = Py_True;
6096 PicklerObject *pickler;
6097
6098 /* fix_imports is a keyword-only argument. */
6099 if (Py_SIZE(args) > 3) {
6100 PyErr_Format(PyExc_TypeError,
6101 "pickle.dump() takes at most 3 positional "
6102 "argument (%zd given)", Py_SIZE(args));
6103 return NULL;
6104 }
6105
6106 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|OO:dump", kwlist,
6107 &obj, &file, &proto, &fix_imports))
6108 return NULL;
6109
6110 pickler = _Pickler_New();
6111 if (pickler == NULL)
6112 return NULL;
6113
6114 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6115 goto error;
6116
6117 if (_Pickler_SetOutputStream(pickler, file) < 0)
6118 goto error;
6119
6120 if (dump(pickler, obj) < 0)
6121 goto error;
6122
6123 if (_Pickler_FlushToFile(pickler) < 0)
6124 goto error;
6125
6126 Py_DECREF(pickler);
6127 Py_RETURN_NONE;
6128
6129 error:
6130 Py_XDECREF(pickler);
6131 return NULL;
6132}
6133
6134PyDoc_STRVAR(pickle_dumps_doc,
6135"dumps(obj, protocol=None, *, fix_imports=True) -> bytes\n"
6136"\n"
6137"Return the pickled representation of the object as a bytes\n"
6138"object, instead of writing it to a file.\n"
6139"\n"
6140"The optional protocol argument tells the pickler to use the given protocol;\n"
6141"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
6142"backward-incompatible protocol designed for Python 3.0.\n"
6143"\n"
6144"Specifying a negative protocol version selects the highest protocol version\n"
6145"supported. The higher the protocol used, the more recent the version of\n"
6146"Python needed to read the pickle produced.\n"
6147"\n"
6148"If fix_imports is True and *protocol* is less than 3, pickle will try to\n"
6149"map the new Python 3.x names to the old module names used in Python 2.x,\n"
6150"so that the pickle data stream is readable with Python 2.x.\n");
6151
6152static PyObject *
6153pickle_dumps(PyObject *self, PyObject *args, PyObject *kwds)
6154{
6155 static char *kwlist[] = {"obj", "protocol", "fix_imports", 0};
6156 PyObject *obj;
6157 PyObject *proto = NULL;
6158 PyObject *result;
6159 PyObject *fix_imports = Py_True;
6160 PicklerObject *pickler;
6161
6162 /* fix_imports is a keyword-only argument. */
6163 if (Py_SIZE(args) > 2) {
6164 PyErr_Format(PyExc_TypeError,
6165 "pickle.dumps() takes at most 2 positional "
6166 "argument (%zd given)", Py_SIZE(args));
6167 return NULL;
6168 }
6169
6170 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:dumps", kwlist,
6171 &obj, &proto, &fix_imports))
6172 return NULL;
6173
6174 pickler = _Pickler_New();
6175 if (pickler == NULL)
6176 return NULL;
6177
6178 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6179 goto error;
6180
6181 if (dump(pickler, obj) < 0)
6182 goto error;
6183
6184 result = _Pickler_GetString(pickler);
6185 Py_DECREF(pickler);
6186 return result;
6187
6188 error:
6189 Py_XDECREF(pickler);
6190 return NULL;
6191}
6192
6193PyDoc_STRVAR(pickle_load_doc,
6194"load(file, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6195"\n"
6196"Read a pickled object representation from the open file object file and\n"
6197"return the reconstituted object hierarchy specified therein. This is\n"
6198"equivalent to ``Unpickler(file).load()``, but may be more efficient.\n"
6199"\n"
6200"The protocol version of the pickle is detected automatically, so no protocol\n"
6201"argument is needed. Bytes past the pickled object's representation are\n"
6202"ignored.\n"
6203"\n"
6204"The argument file must have two methods, a read() method that takes an\n"
6205"integer argument, and a readline() method that requires no arguments. Both\n"
6206"methods should return bytes. Thus *file* can be a binary file object opened\n"
6207"for reading, a BytesIO object, or any other custom object that meets this\n"
6208"interface.\n"
6209"\n"
6210"Optional keyword arguments are fix_imports, encoding and errors,\n"
6211"which are used to control compatiblity support for pickle stream generated\n"
6212"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6213"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6214"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6215"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6216
6217static PyObject *
6218pickle_load(PyObject *self, PyObject *args, PyObject *kwds)
6219{
6220 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
6221 PyObject *file;
6222 PyObject *fix_imports = Py_True;
6223 PyObject *result;
6224 char *encoding = NULL;
6225 char *errors = NULL;
6226 UnpicklerObject *unpickler;
6227
6228 /* fix_imports, encoding and errors are a keyword-only argument. */
6229 if (Py_SIZE(args) != 1) {
6230 PyErr_Format(PyExc_TypeError,
6231 "pickle.load() takes exactly one positional "
6232 "argument (%zd given)", Py_SIZE(args));
6233 return NULL;
6234 }
6235
6236 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:load", kwlist,
6237 &file, &fix_imports, &encoding, &errors))
6238 return NULL;
6239
6240 unpickler = _Unpickler_New();
6241 if (unpickler == NULL)
6242 return NULL;
6243
6244 if (_Unpickler_SetInputStream(unpickler, file) < 0)
6245 goto error;
6246
6247 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6248 goto error;
6249
6250 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6251 if (unpickler->fix_imports == -1)
6252 goto error;
6253
6254 result = load(unpickler);
6255 Py_DECREF(unpickler);
6256 return result;
6257
6258 error:
6259 Py_XDECREF(unpickler);
6260 return NULL;
6261}
6262
6263PyDoc_STRVAR(pickle_loads_doc,
6264"loads(input, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6265"\n"
6266"Read a pickled object hierarchy from a bytes object and return the\n"
6267"reconstituted object hierarchy specified therein\n"
6268"\n"
6269"The protocol version of the pickle is detected automatically, so no protocol\n"
6270"argument is needed. Bytes past the pickled object's representation are\n"
6271"ignored.\n"
6272"\n"
6273"Optional keyword arguments are fix_imports, encoding and errors, which\n"
6274"are used to control compatiblity support for pickle stream generated\n"
6275"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6276"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6277"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6278"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6279
6280static PyObject *
6281pickle_loads(PyObject *self, PyObject *args, PyObject *kwds)
6282{
6283 static char *kwlist[] = {"input", "fix_imports", "encoding", "errors", 0};
6284 PyObject *input;
6285 PyObject *fix_imports = Py_True;
6286 PyObject *result;
6287 char *encoding = NULL;
6288 char *errors = NULL;
6289 UnpicklerObject *unpickler;
6290
6291 /* fix_imports, encoding and errors are a keyword-only argument. */
6292 if (Py_SIZE(args) != 1) {
6293 PyErr_Format(PyExc_TypeError,
6294 "pickle.loads() takes exactly one positional "
6295 "argument (%zd given)", Py_SIZE(args));
6296 return NULL;
6297 }
6298
6299 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:loads", kwlist,
6300 &input, &fix_imports, &encoding, &errors))
6301 return NULL;
6302
6303 unpickler = _Unpickler_New();
6304 if (unpickler == NULL)
6305 return NULL;
6306
6307 if (_Unpickler_SetStringInput(unpickler, input) < 0)
6308 goto error;
6309
6310 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6311 goto error;
6312
6313 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6314 if (unpickler->fix_imports == -1)
6315 goto error;
6316
6317 result = load(unpickler);
6318 Py_DECREF(unpickler);
6319 return result;
6320
6321 error:
6322 Py_XDECREF(unpickler);
6323 return NULL;
6324}
6325
6326
6327static struct PyMethodDef pickle_methods[] = {
6328 {"dump", (PyCFunction)pickle_dump, METH_VARARGS|METH_KEYWORDS,
6329 pickle_dump_doc},
6330 {"dumps", (PyCFunction)pickle_dumps, METH_VARARGS|METH_KEYWORDS,
6331 pickle_dumps_doc},
6332 {"load", (PyCFunction)pickle_load, METH_VARARGS|METH_KEYWORDS,
6333 pickle_load_doc},
6334 {"loads", (PyCFunction)pickle_loads, METH_VARARGS|METH_KEYWORDS,
6335 pickle_loads_doc},
6336 {NULL, NULL} /* sentinel */
6337};
6338
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006339static int
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006340initmodule(void)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006341{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006342 PyObject *copyreg = NULL;
6343 PyObject *compat_pickle = NULL;
6344
6345 /* XXX: We should ensure that the types of the dictionaries imported are
6346 exactly PyDict objects. Otherwise, it is possible to crash the pickle
6347 since we use the PyDict API directly to access these dictionaries. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006348
6349 copyreg = PyImport_ImportModule("copyreg");
6350 if (!copyreg)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006351 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006352 dispatch_table = PyObject_GetAttrString(copyreg, "dispatch_table");
6353 if (!dispatch_table)
6354 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006355 extension_registry = \
6356 PyObject_GetAttrString(copyreg, "_extension_registry");
6357 if (!extension_registry)
6358 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006359 inverted_registry = PyObject_GetAttrString(copyreg, "_inverted_registry");
6360 if (!inverted_registry)
6361 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006362 extension_cache = PyObject_GetAttrString(copyreg, "_extension_cache");
6363 if (!extension_cache)
6364 goto error;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006365 Py_CLEAR(copyreg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006366
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006367 /* Load the 2.x -> 3.x stdlib module mapping tables */
6368 compat_pickle = PyImport_ImportModule("_compat_pickle");
6369 if (!compat_pickle)
6370 goto error;
6371 name_mapping_2to3 = PyObject_GetAttrString(compat_pickle, "NAME_MAPPING");
6372 if (!name_mapping_2to3)
6373 goto error;
6374 if (!PyDict_CheckExact(name_mapping_2to3)) {
6375 PyErr_Format(PyExc_RuntimeError,
6376 "_compat_pickle.NAME_MAPPING should be a dict, not %.200s",
6377 Py_TYPE(name_mapping_2to3)->tp_name);
6378 goto error;
6379 }
6380 import_mapping_2to3 = PyObject_GetAttrString(compat_pickle,
6381 "IMPORT_MAPPING");
6382 if (!import_mapping_2to3)
6383 goto error;
6384 if (!PyDict_CheckExact(import_mapping_2to3)) {
6385 PyErr_Format(PyExc_RuntimeError,
6386 "_compat_pickle.IMPORT_MAPPING should be a dict, "
6387 "not %.200s", Py_TYPE(import_mapping_2to3)->tp_name);
6388 goto error;
6389 }
6390 /* ... and the 3.x -> 2.x mapping tables */
6391 name_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6392 "REVERSE_NAME_MAPPING");
6393 if (!name_mapping_3to2)
6394 goto error;
6395 if (!PyDict_CheckExact(name_mapping_3to2)) {
6396 PyErr_Format(PyExc_RuntimeError,
Ezio Melotti13925002011-03-16 11:05:33 +02006397 "_compat_pickle.REVERSE_NAME_MAPPING should be a dict, "
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006398 "not %.200s", Py_TYPE(name_mapping_3to2)->tp_name);
6399 goto error;
6400 }
6401 import_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6402 "REVERSE_IMPORT_MAPPING");
6403 if (!import_mapping_3to2)
6404 goto error;
6405 if (!PyDict_CheckExact(import_mapping_3to2)) {
6406 PyErr_Format(PyExc_RuntimeError,
6407 "_compat_pickle.REVERSE_IMPORT_MAPPING should be a dict, "
6408 "not %.200s", Py_TYPE(import_mapping_3to2)->tp_name);
6409 goto error;
6410 }
6411 Py_CLEAR(compat_pickle);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006412
6413 empty_tuple = PyTuple_New(0);
6414 if (empty_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006415 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006416 two_tuple = PyTuple_New(2);
6417 if (two_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006418 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006419 /* We use this temp container with no regard to refcounts, or to
6420 * keeping containees alive. Exempt from GC, because we don't
6421 * want anything looking at two_tuple() by magic.
6422 */
6423 PyObject_GC_UnTrack(two_tuple);
6424
6425 return 0;
6426
6427 error:
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006428 Py_CLEAR(copyreg);
6429 Py_CLEAR(dispatch_table);
6430 Py_CLEAR(extension_registry);
6431 Py_CLEAR(inverted_registry);
6432 Py_CLEAR(extension_cache);
6433 Py_CLEAR(compat_pickle);
6434 Py_CLEAR(name_mapping_2to3);
6435 Py_CLEAR(import_mapping_2to3);
6436 Py_CLEAR(name_mapping_3to2);
6437 Py_CLEAR(import_mapping_3to2);
6438 Py_CLEAR(empty_tuple);
6439 Py_CLEAR(two_tuple);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006440 return -1;
6441}
6442
6443static struct PyModuleDef _picklemodule = {
6444 PyModuleDef_HEAD_INIT,
6445 "_pickle",
6446 pickle_module_doc,
6447 -1,
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006448 pickle_methods,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006449 NULL,
6450 NULL,
6451 NULL,
6452 NULL
6453};
6454
6455PyMODINIT_FUNC
6456PyInit__pickle(void)
6457{
6458 PyObject *m;
6459
6460 if (PyType_Ready(&Unpickler_Type) < 0)
6461 return NULL;
6462 if (PyType_Ready(&Pickler_Type) < 0)
6463 return NULL;
6464 if (PyType_Ready(&Pdata_Type) < 0)
6465 return NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006466 if (PyType_Ready(&PicklerMemoProxyType) < 0)
6467 return NULL;
6468 if (PyType_Ready(&UnpicklerMemoProxyType) < 0)
6469 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006470
6471 /* Create the module and add the functions. */
6472 m = PyModule_Create(&_picklemodule);
6473 if (m == NULL)
6474 return NULL;
6475
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006476 Py_INCREF(&Pickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006477 if (PyModule_AddObject(m, "Pickler", (PyObject *)&Pickler_Type) < 0)
6478 return NULL;
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006479 Py_INCREF(&Unpickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006480 if (PyModule_AddObject(m, "Unpickler", (PyObject *)&Unpickler_Type) < 0)
6481 return NULL;
6482
6483 /* Initialize the exceptions. */
6484 PickleError = PyErr_NewException("_pickle.PickleError", NULL, NULL);
6485 if (PickleError == NULL)
6486 return NULL;
6487 PicklingError = \
6488 PyErr_NewException("_pickle.PicklingError", PickleError, NULL);
6489 if (PicklingError == NULL)
6490 return NULL;
6491 UnpicklingError = \
6492 PyErr_NewException("_pickle.UnpicklingError", PickleError, NULL);
6493 if (UnpicklingError == NULL)
6494 return NULL;
6495
6496 if (PyModule_AddObject(m, "PickleError", PickleError) < 0)
6497 return NULL;
6498 if (PyModule_AddObject(m, "PicklingError", PicklingError) < 0)
6499 return NULL;
6500 if (PyModule_AddObject(m, "UnpicklingError", UnpicklingError) < 0)
6501 return NULL;
6502
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006503 if (initmodule() < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006504 return NULL;
6505
6506 return m;
6507}