blob: 5b30931310838628e9fb1baf0c73172d7416e5d3 [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;
777
778 self->memo = PyMemoTable_New();
779 if (self->memo == NULL) {
780 Py_DECREF(self);
781 return NULL;
782 }
783 self->max_output_len = WRITE_BUF_SIZE;
784 self->output_len = 0;
785 self->output_buffer = PyBytes_FromStringAndSize(NULL,
786 self->max_output_len);
787 if (self->output_buffer == NULL) {
788 Py_DECREF(self);
789 return NULL;
790 }
791 return self;
792}
793
794static int
795_Pickler_SetProtocol(PicklerObject *self, PyObject *proto_obj,
796 PyObject *fix_imports_obj)
797{
798 long proto = 0;
799 int fix_imports;
800
801 if (proto_obj == NULL || proto_obj == Py_None)
802 proto = DEFAULT_PROTOCOL;
803 else {
804 proto = PyLong_AsLong(proto_obj);
805 if (proto == -1 && PyErr_Occurred())
806 return -1;
807 }
808 if (proto < 0)
809 proto = HIGHEST_PROTOCOL;
810 if (proto > HIGHEST_PROTOCOL) {
811 PyErr_Format(PyExc_ValueError, "pickle protocol must be <= %d",
812 HIGHEST_PROTOCOL);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000813 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000814 }
815 fix_imports = PyObject_IsTrue(fix_imports_obj);
816 if (fix_imports == -1)
817 return -1;
Victor Stinner121aab42011-09-29 23:40:53 +0200818
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000819 self->proto = proto;
820 self->bin = proto > 0;
821 self->fix_imports = fix_imports && proto < 3;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000822
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000823 return 0;
824}
825
826/* Returns -1 (with an exception set) on failure, 0 on success. This may
827 be called once on a freshly created Pickler. */
828static int
829_Pickler_SetOutputStream(PicklerObject *self, PyObject *file)
830{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +0200831 _Py_IDENTIFIER(write);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000832 assert(file != NULL);
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +0200833 self->write = _PyObject_GetAttrId(file, &PyId_write);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000834 if (self->write == NULL) {
835 if (PyErr_ExceptionMatches(PyExc_AttributeError))
836 PyErr_SetString(PyExc_TypeError,
837 "file must have a 'write' attribute");
838 return -1;
839 }
840
841 return 0;
842}
843
844/* See documentation for _Pickler_FastCall(). */
845static PyObject *
846_Unpickler_FastCall(UnpicklerObject *self, PyObject *func, PyObject *arg)
847{
848 PyObject *result = NULL;
849
850 ARG_TUP(self, arg);
851 if (self->arg) {
852 result = PyObject_Call(func, self->arg, NULL);
853 FREE_ARG_TUP(self);
854 }
855 return result;
856}
857
858/* Returns the size of the input on success, -1 on failure. This takes its
859 own reference to `input`. */
860static Py_ssize_t
861_Unpickler_SetStringInput(UnpicklerObject *self, PyObject *input)
862{
863 if (self->buffer.buf != NULL)
864 PyBuffer_Release(&self->buffer);
865 if (PyObject_GetBuffer(input, &self->buffer, PyBUF_CONTIG_RO) < 0)
866 return -1;
867 self->input_buffer = self->buffer.buf;
868 self->input_len = self->buffer.len;
869 self->next_read_idx = 0;
Antoine Pitrou04248a82010-10-12 20:51:21 +0000870 self->prefetched_idx = self->input_len;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000871 return self->input_len;
872}
873
Antoine Pitrou04248a82010-10-12 20:51:21 +0000874static int
875_Unpickler_SkipConsumed(UnpicklerObject *self)
876{
877 Py_ssize_t consumed = self->next_read_idx - self->prefetched_idx;
878
879 if (consumed > 0) {
880 PyObject *r;
881 assert(self->peek); /* otherwise we did something wrong */
882 /* This makes an useless copy... */
883 r = PyObject_CallFunction(self->read, "n", consumed);
884 if (r == NULL)
885 return -1;
886 Py_DECREF(r);
887 self->prefetched_idx = self->next_read_idx;
888 }
889 return 0;
890}
891
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000892static const Py_ssize_t READ_WHOLE_LINE = -1;
893
894/* If reading from a file, we need to only pull the bytes we need, since there
895 may be multiple pickle objects arranged contiguously in the same input
896 buffer.
897
898 If `n` is READ_WHOLE_LINE, read a whole line. Otherwise, read up to `n`
899 bytes from the input stream/buffer.
900
901 Update the unpickler's input buffer with the newly-read data. Returns -1 on
902 failure; on success, returns the number of bytes read from the file.
903
904 On success, self->input_len will be 0; this is intentional so that when
905 unpickling from a file, the "we've run out of data" code paths will trigger,
906 causing the Unpickler to go back to the file for more data. Use the returned
907 size to tell you how much data you can process. */
908static Py_ssize_t
909_Unpickler_ReadFromFile(UnpicklerObject *self, Py_ssize_t n)
910{
911 PyObject *data;
Antoine Pitrou04248a82010-10-12 20:51:21 +0000912 Py_ssize_t read_size, prefetched_size = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000913
914 assert(self->read != NULL);
Victor Stinner121aab42011-09-29 23:40:53 +0200915
Antoine Pitrou04248a82010-10-12 20:51:21 +0000916 if (_Unpickler_SkipConsumed(self) < 0)
917 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000918
919 if (n == READ_WHOLE_LINE)
920 data = PyObject_Call(self->readline, empty_tuple, NULL);
921 else {
922 PyObject *len = PyLong_FromSsize_t(n);
923 if (len == NULL)
924 return -1;
925 data = _Unpickler_FastCall(self, self->read, len);
926 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000927 if (data == NULL)
928 return -1;
929
Antoine Pitrou04248a82010-10-12 20:51:21 +0000930 /* Prefetch some data without advancing the file pointer, if possible */
931 if (self->peek) {
932 PyObject *len, *prefetched;
933 len = PyLong_FromSsize_t(PREFETCH);
934 if (len == NULL) {
935 Py_DECREF(data);
936 return -1;
937 }
938 prefetched = _Unpickler_FastCall(self, self->peek, len);
939 if (prefetched == NULL) {
940 if (PyErr_ExceptionMatches(PyExc_NotImplementedError)) {
941 /* peek() is probably not supported by the given file object */
942 PyErr_Clear();
943 Py_CLEAR(self->peek);
944 }
945 else {
946 Py_DECREF(data);
947 return -1;
948 }
949 }
950 else {
951 assert(PyBytes_Check(prefetched));
952 prefetched_size = PyBytes_GET_SIZE(prefetched);
953 PyBytes_ConcatAndDel(&data, prefetched);
954 if (data == NULL)
955 return -1;
956 }
957 }
958
959 read_size = _Unpickler_SetStringInput(self, data) - prefetched_size;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000960 Py_DECREF(data);
Antoine Pitrou04248a82010-10-12 20:51:21 +0000961 self->prefetched_idx = read_size;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000962 return read_size;
963}
964
965/* Read `n` bytes from the unpickler's data source, storing the result in `*s`.
966
967 This should be used for all data reads, rather than accessing the unpickler's
968 input buffer directly. This method deals correctly with reading from input
969 streams, which the input buffer doesn't deal with.
970
971 Note that when reading from a file-like object, self->next_read_idx won't
972 be updated (it should remain at 0 for the entire unpickling process). You
973 should use this function's return value to know how many bytes you can
974 consume.
975
976 Returns -1 (with an exception set) on failure. On success, return the
977 number of chars read. */
978static Py_ssize_t
979_Unpickler_Read(UnpicklerObject *self, char **s, Py_ssize_t n)
980{
Antoine Pitrou04248a82010-10-12 20:51:21 +0000981 Py_ssize_t num_read;
982
Antoine Pitrou04248a82010-10-12 20:51:21 +0000983 if (self->next_read_idx + n <= self->input_len) {
984 *s = self->input_buffer + self->next_read_idx;
985 self->next_read_idx += n;
986 return n;
987 }
988 if (!self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000989 PyErr_Format(PyExc_EOFError, "Ran out of input");
Amaury Forgeot d'Arc3e4e72f2008-11-11 20:05:06 +0000990 return -1;
991 }
Antoine Pitrou04248a82010-10-12 20:51:21 +0000992 num_read = _Unpickler_ReadFromFile(self, n);
993 if (num_read < 0)
994 return -1;
995 if (num_read < n) {
996 PyErr_Format(PyExc_EOFError, "Ran out of input");
997 return -1;
998 }
999 *s = self->input_buffer;
1000 self->next_read_idx = n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001001 return n;
1002}
1003
1004static Py_ssize_t
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001005_Unpickler_CopyLine(UnpicklerObject *self, char *line, Py_ssize_t len,
1006 char **result)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001007{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001008 char *input_line = PyMem_Realloc(self->input_line, len + 1);
1009 if (input_line == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001010 return -1;
1011
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001012 memcpy(input_line, line, len);
1013 input_line[len] = '\0';
1014 self->input_line = input_line;
1015 *result = self->input_line;
1016 return len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001017}
1018
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001019/* Read a line from the input stream/buffer. If we run off the end of the input
1020 before hitting \n, return the data we found.
1021
1022 Returns the number of chars read, or -1 on failure. */
1023static Py_ssize_t
1024_Unpickler_Readline(UnpicklerObject *self, char **result)
1025{
1026 Py_ssize_t i, num_read;
1027
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001028 for (i = self->next_read_idx; i < self->input_len; i++) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001029 if (self->input_buffer[i] == '\n') {
1030 char *line_start = self->input_buffer + self->next_read_idx;
1031 num_read = i - self->next_read_idx + 1;
1032 self->next_read_idx = i + 1;
1033 return _Unpickler_CopyLine(self, line_start, num_read, result);
1034 }
1035 }
1036 if (self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001037 num_read = _Unpickler_ReadFromFile(self, READ_WHOLE_LINE);
1038 if (num_read < 0)
1039 return -1;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001040 self->next_read_idx = num_read;
Antoine Pitrouf6c7a852011-08-11 21:04:02 +02001041 return _Unpickler_CopyLine(self, self->input_buffer, num_read, result);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001042 }
Victor Stinner121aab42011-09-29 23:40:53 +02001043
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001044 /* If we get here, we've run off the end of the input string. Return the
1045 remaining string and let the caller figure it out. */
1046 *result = self->input_buffer + self->next_read_idx;
1047 num_read = i - self->next_read_idx;
1048 self->next_read_idx = i;
1049 return num_read;
1050}
1051
1052/* Returns -1 (with an exception set) on failure, 0 on success. The memo array
1053 will be modified in place. */
1054static int
1055_Unpickler_ResizeMemoList(UnpicklerObject *self, Py_ssize_t new_size)
1056{
1057 Py_ssize_t i;
1058 PyObject **memo;
1059
1060 assert(new_size > self->memo_size);
1061
1062 memo = PyMem_REALLOC(self->memo, new_size * sizeof(PyObject *));
1063 if (memo == NULL) {
1064 PyErr_NoMemory();
1065 return -1;
1066 }
1067 self->memo = memo;
1068 for (i = self->memo_size; i < new_size; i++)
1069 self->memo[i] = NULL;
1070 self->memo_size = new_size;
1071 return 0;
1072}
1073
1074/* Returns NULL if idx is out of bounds. */
1075static PyObject *
1076_Unpickler_MemoGet(UnpicklerObject *self, Py_ssize_t idx)
1077{
1078 if (idx < 0 || idx >= self->memo_size)
1079 return NULL;
1080
1081 return self->memo[idx];
1082}
1083
1084/* Returns -1 (with an exception set) on failure, 0 on success.
1085 This takes its own reference to `value`. */
1086static int
1087_Unpickler_MemoPut(UnpicklerObject *self, Py_ssize_t idx, PyObject *value)
1088{
1089 PyObject *old_item;
1090
1091 if (idx >= self->memo_size) {
1092 if (_Unpickler_ResizeMemoList(self, idx * 2) < 0)
1093 return -1;
1094 assert(idx < self->memo_size);
1095 }
1096 Py_INCREF(value);
1097 old_item = self->memo[idx];
1098 self->memo[idx] = value;
1099 Py_XDECREF(old_item);
1100 return 0;
1101}
1102
1103static PyObject **
1104_Unpickler_NewMemo(Py_ssize_t new_size)
1105{
1106 PyObject **memo = PyMem_MALLOC(new_size * sizeof(PyObject *));
1107 if (memo == NULL)
1108 return NULL;
1109 memset(memo, 0, new_size * sizeof(PyObject *));
1110 return memo;
1111}
1112
1113/* Free the unpickler's memo, taking care to decref any items left in it. */
1114static void
1115_Unpickler_MemoCleanup(UnpicklerObject *self)
1116{
1117 Py_ssize_t i;
1118 PyObject **memo = self->memo;
1119
1120 if (self->memo == NULL)
1121 return;
1122 self->memo = NULL;
1123 i = self->memo_size;
1124 while (--i >= 0) {
1125 Py_XDECREF(memo[i]);
1126 }
1127 PyMem_FREE(memo);
1128}
1129
1130static UnpicklerObject *
1131_Unpickler_New(void)
1132{
1133 UnpicklerObject *self;
1134
1135 self = PyObject_GC_New(UnpicklerObject, &Unpickler_Type);
1136 if (self == NULL)
1137 return NULL;
1138
1139 self->stack = (Pdata *)Pdata_New();
1140 if (self->stack == NULL) {
1141 Py_DECREF(self);
1142 return NULL;
1143 }
1144 memset(&self->buffer, 0, sizeof(Py_buffer));
1145
1146 self->memo_size = 32;
1147 self->memo = _Unpickler_NewMemo(self->memo_size);
1148 if (self->memo == NULL) {
1149 Py_DECREF(self);
1150 return NULL;
1151 }
1152
1153 self->arg = NULL;
1154 self->pers_func = NULL;
1155 self->input_buffer = NULL;
1156 self->input_line = NULL;
1157 self->input_len = 0;
1158 self->next_read_idx = 0;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001159 self->prefetched_idx = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001160 self->read = NULL;
1161 self->readline = NULL;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001162 self->peek = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001163 self->encoding = NULL;
1164 self->errors = NULL;
1165 self->marks = NULL;
1166 self->num_marks = 0;
1167 self->marks_size = 0;
1168 self->proto = 0;
1169 self->fix_imports = 0;
1170
1171 return self;
1172}
1173
1174/* Returns -1 (with an exception set) on failure, 0 on success. This may
1175 be called once on a freshly created Pickler. */
1176static int
1177_Unpickler_SetInputStream(UnpicklerObject *self, PyObject *file)
1178{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02001179 _Py_IDENTIFIER(peek);
1180 _Py_IDENTIFIER(read);
1181 _Py_IDENTIFIER(readline);
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02001182
1183 self->peek = _PyObject_GetAttrId(file, &PyId_peek);
Antoine Pitrou04248a82010-10-12 20:51:21 +00001184 if (self->peek == NULL) {
1185 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1186 PyErr_Clear();
1187 else
1188 return -1;
1189 }
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02001190 self->read = _PyObject_GetAttrId(file, &PyId_read);
1191 self->readline = _PyObject_GetAttrId(file, &PyId_readline);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001192 if (self->readline == NULL || self->read == NULL) {
1193 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1194 PyErr_SetString(PyExc_TypeError,
1195 "file must have 'read' and 'readline' attributes");
1196 Py_CLEAR(self->read);
1197 Py_CLEAR(self->readline);
Antoine Pitrou04248a82010-10-12 20:51:21 +00001198 Py_CLEAR(self->peek);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001199 return -1;
1200 }
1201 return 0;
1202}
1203
1204/* Returns -1 (with an exception set) on failure, 0 on success. This may
1205 be called once on a freshly created Pickler. */
1206static int
1207_Unpickler_SetInputEncoding(UnpicklerObject *self,
1208 const char *encoding,
1209 const char *errors)
1210{
1211 if (encoding == NULL)
1212 encoding = "ASCII";
1213 if (errors == NULL)
1214 errors = "strict";
1215
1216 self->encoding = strdup(encoding);
1217 self->errors = strdup(errors);
1218 if (self->encoding == NULL || self->errors == NULL) {
1219 PyErr_NoMemory();
1220 return -1;
1221 }
1222 return 0;
1223}
1224
1225/* Generate a GET opcode for an object stored in the memo. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001226static int
1227memo_get(PicklerObject *self, PyObject *key)
1228{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001229 Py_ssize_t *value;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001230 char pdata[30];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001231 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001232
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001233 value = PyMemoTable_Get(self->memo, key);
1234 if (value == NULL) {
1235 PyErr_SetObject(PyExc_KeyError, key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001236 return -1;
1237 }
1238
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001239 if (!self->bin) {
1240 pdata[0] = GET;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001241 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1,
1242 "%" PY_FORMAT_SIZE_T "d\n", *value);
1243 len = strlen(pdata);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001244 }
1245 else {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001246 if (*value < 256) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001247 pdata[0] = BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001248 pdata[1] = (unsigned char)(*value & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001249 len = 2;
1250 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001251 else if (*value <= 0xffffffffL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001252 pdata[0] = LONG_BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001253 pdata[1] = (unsigned char)(*value & 0xff);
1254 pdata[2] = (unsigned char)((*value >> 8) & 0xff);
1255 pdata[3] = (unsigned char)((*value >> 16) & 0xff);
1256 pdata[4] = (unsigned char)((*value >> 24) & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001257 len = 5;
1258 }
1259 else { /* unlikely */
1260 PyErr_SetString(PicklingError,
1261 "memo id too large for LONG_BINGET");
1262 return -1;
1263 }
1264 }
1265
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001266 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001267 return -1;
1268
1269 return 0;
1270}
1271
1272/* Store an object in the memo, assign it a new unique ID based on the number
1273 of objects currently stored in the memo and generate a PUT opcode. */
1274static int
1275memo_put(PicklerObject *self, PyObject *obj)
1276{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001277 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001278 char pdata[30];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001279 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001280 int status = 0;
1281
1282 if (self->fast)
1283 return 0;
1284
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001285 x = PyMemoTable_Size(self->memo);
1286 if (PyMemoTable_Set(self->memo, obj, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001287 goto error;
1288
1289 if (!self->bin) {
1290 pdata[0] = PUT;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001291 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1,
1292 "%" PY_FORMAT_SIZE_T "d\n", x);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001293 len = strlen(pdata);
1294 }
1295 else {
1296 if (x < 256) {
1297 pdata[0] = BINPUT;
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00001298 pdata[1] = (unsigned char)x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001299 len = 2;
1300 }
1301 else if (x <= 0xffffffffL) {
1302 pdata[0] = LONG_BINPUT;
1303 pdata[1] = (unsigned char)(x & 0xff);
1304 pdata[2] = (unsigned char)((x >> 8) & 0xff);
1305 pdata[3] = (unsigned char)((x >> 16) & 0xff);
1306 pdata[4] = (unsigned char)((x >> 24) & 0xff);
1307 len = 5;
1308 }
1309 else { /* unlikely */
1310 PyErr_SetString(PicklingError,
1311 "memo id too large for LONG_BINPUT");
1312 return -1;
1313 }
1314 }
1315
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001316 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001317 goto error;
1318
1319 if (0) {
1320 error:
1321 status = -1;
1322 }
1323
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001324 return status;
1325}
1326
1327static PyObject *
1328whichmodule(PyObject *global, PyObject *global_name)
1329{
1330 Py_ssize_t i, j;
1331 static PyObject *module_str = NULL;
1332 static PyObject *main_str = NULL;
1333 PyObject *module_name;
1334 PyObject *modules_dict;
1335 PyObject *module;
1336 PyObject *obj;
1337
1338 if (module_str == NULL) {
1339 module_str = PyUnicode_InternFromString("__module__");
1340 if (module_str == NULL)
1341 return NULL;
1342 main_str = PyUnicode_InternFromString("__main__");
1343 if (main_str == NULL)
1344 return NULL;
1345 }
1346
1347 module_name = PyObject_GetAttr(global, module_str);
1348
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00001349 /* In some rare cases (e.g., bound methods of extension types),
1350 __module__ can be None. If it is so, then search sys.modules
1351 for the module of global. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001352 if (module_name == Py_None) {
1353 Py_DECREF(module_name);
1354 goto search;
1355 }
1356
1357 if (module_name) {
1358 return module_name;
1359 }
1360 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1361 PyErr_Clear();
1362 else
1363 return NULL;
1364
1365 search:
1366 modules_dict = PySys_GetObject("modules");
1367 if (modules_dict == NULL)
1368 return NULL;
1369
1370 i = 0;
1371 module_name = NULL;
1372 while ((j = PyDict_Next(modules_dict, &i, &module_name, &module))) {
Mark Dickinson211c6252009-02-01 10:28:51 +00001373 if (PyObject_RichCompareBool(module_name, main_str, Py_EQ) == 1)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001374 continue;
1375
1376 obj = PyObject_GetAttr(module, global_name);
1377 if (obj == NULL) {
1378 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1379 PyErr_Clear();
1380 else
1381 return NULL;
1382 continue;
1383 }
1384
1385 if (obj != global) {
1386 Py_DECREF(obj);
1387 continue;
1388 }
1389
1390 Py_DECREF(obj);
1391 break;
1392 }
1393
1394 /* If no module is found, use __main__. */
1395 if (!j) {
1396 module_name = main_str;
1397 }
1398
1399 Py_INCREF(module_name);
1400 return module_name;
1401}
1402
1403/* fast_save_enter() and fast_save_leave() are guards against recursive
1404 objects when Pickler is used with the "fast mode" (i.e., with object
1405 memoization disabled). If the nesting of a list or dict object exceed
1406 FAST_NESTING_LIMIT, these guards will start keeping an internal
1407 reference to the seen list or dict objects and check whether these objects
1408 are recursive. These are not strictly necessary, since save() has a
1409 hard-coded recursion limit, but they give a nicer error message than the
1410 typical RuntimeError. */
1411static int
1412fast_save_enter(PicklerObject *self, PyObject *obj)
1413{
1414 /* if fast_nesting < 0, we're doing an error exit. */
1415 if (++self->fast_nesting >= FAST_NESTING_LIMIT) {
1416 PyObject *key = NULL;
1417 if (self->fast_memo == NULL) {
1418 self->fast_memo = PyDict_New();
1419 if (self->fast_memo == NULL) {
1420 self->fast_nesting = -1;
1421 return 0;
1422 }
1423 }
1424 key = PyLong_FromVoidPtr(obj);
1425 if (key == NULL)
1426 return 0;
1427 if (PyDict_GetItem(self->fast_memo, key)) {
1428 Py_DECREF(key);
1429 PyErr_Format(PyExc_ValueError,
1430 "fast mode: can't pickle cyclic objects "
1431 "including object type %.200s at %p",
1432 obj->ob_type->tp_name, obj);
1433 self->fast_nesting = -1;
1434 return 0;
1435 }
1436 if (PyDict_SetItem(self->fast_memo, key, Py_None) < 0) {
1437 Py_DECREF(key);
1438 self->fast_nesting = -1;
1439 return 0;
1440 }
1441 Py_DECREF(key);
1442 }
1443 return 1;
1444}
1445
1446static int
1447fast_save_leave(PicklerObject *self, PyObject *obj)
1448{
1449 if (self->fast_nesting-- >= FAST_NESTING_LIMIT) {
1450 PyObject *key = PyLong_FromVoidPtr(obj);
1451 if (key == NULL)
1452 return 0;
1453 if (PyDict_DelItem(self->fast_memo, key) < 0) {
1454 Py_DECREF(key);
1455 return 0;
1456 }
1457 Py_DECREF(key);
1458 }
1459 return 1;
1460}
1461
1462static int
1463save_none(PicklerObject *self, PyObject *obj)
1464{
1465 const char none_op = NONE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001466 if (_Pickler_Write(self, &none_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001467 return -1;
1468
1469 return 0;
1470}
1471
1472static int
1473save_bool(PicklerObject *self, PyObject *obj)
1474{
1475 static const char *buf[2] = { FALSE, TRUE };
1476 const char len[2] = {sizeof(FALSE) - 1, sizeof(TRUE) - 1};
1477 int p = (obj == Py_True);
1478
1479 if (self->proto >= 2) {
1480 const char bool_op = p ? NEWTRUE : NEWFALSE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001481 if (_Pickler_Write(self, &bool_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001482 return -1;
1483 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001484 else if (_Pickler_Write(self, buf[p], len[p]) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001485 return -1;
1486
1487 return 0;
1488}
1489
1490static int
1491save_int(PicklerObject *self, long x)
1492{
1493 char pdata[32];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001494 Py_ssize_t len = 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001495
1496 if (!self->bin
1497#if SIZEOF_LONG > 4
1498 || x > 0x7fffffffL || x < -0x80000000L
1499#endif
1500 ) {
1501 /* Text-mode pickle, or long too big to fit in the 4-byte
1502 * signed BININT format: store as a string.
1503 */
Mark Dickinson8dd05142009-01-20 20:43:58 +00001504 pdata[0] = LONG; /* use LONG for consistency with pickle.py */
1505 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1, "%ldL\n", x);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001506 if (_Pickler_Write(self, pdata, strlen(pdata)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001507 return -1;
1508 }
1509 else {
1510 /* Binary pickle and x fits in a signed 4-byte int. */
1511 pdata[1] = (unsigned char)(x & 0xff);
1512 pdata[2] = (unsigned char)((x >> 8) & 0xff);
1513 pdata[3] = (unsigned char)((x >> 16) & 0xff);
1514 pdata[4] = (unsigned char)((x >> 24) & 0xff);
1515
1516 if ((pdata[4] == 0) && (pdata[3] == 0)) {
1517 if (pdata[2] == 0) {
1518 pdata[0] = BININT1;
1519 len = 2;
1520 }
1521 else {
1522 pdata[0] = BININT2;
1523 len = 3;
1524 }
1525 }
1526 else {
1527 pdata[0] = BININT;
1528 len = 5;
1529 }
1530
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001531 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001532 return -1;
1533 }
1534
1535 return 0;
1536}
1537
1538static int
1539save_long(PicklerObject *self, PyObject *obj)
1540{
1541 PyObject *repr = NULL;
1542 Py_ssize_t size;
1543 long val = PyLong_AsLong(obj);
1544 int status = 0;
1545
1546 const char long_op = LONG;
1547
1548 if (val == -1 && PyErr_Occurred()) {
1549 /* out of range for int pickling */
1550 PyErr_Clear();
1551 }
Antoine Pitroue58bffb2011-08-13 20:40:32 +02001552 else
1553#if SIZEOF_LONG > 4
1554 if (val <= 0x7fffffffL && val >= -0x80000000L)
1555#endif
1556 return save_int(self, val);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001557
1558 if (self->proto >= 2) {
1559 /* Linear-time pickling. */
1560 size_t nbits;
1561 size_t nbytes;
1562 unsigned char *pdata;
1563 char header[5];
1564 int i;
1565 int sign = _PyLong_Sign(obj);
1566
1567 if (sign == 0) {
1568 header[0] = LONG1;
1569 header[1] = 0; /* It's 0 -- an empty bytestring. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001570 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001571 goto error;
1572 return 0;
1573 }
1574 nbits = _PyLong_NumBits(obj);
1575 if (nbits == (size_t)-1 && PyErr_Occurred())
1576 goto error;
1577 /* How many bytes do we need? There are nbits >> 3 full
1578 * bytes of data, and nbits & 7 leftover bits. If there
1579 * are any leftover bits, then we clearly need another
1580 * byte. Wnat's not so obvious is that we *probably*
1581 * need another byte even if there aren't any leftovers:
1582 * the most-significant bit of the most-significant byte
1583 * acts like a sign bit, and it's usually got a sense
1584 * opposite of the one we need. The exception is longs
1585 * of the form -(2**(8*j-1)) for j > 0. Such a long is
1586 * its own 256's-complement, so has the right sign bit
1587 * even without the extra byte. That's a pain to check
1588 * for in advance, though, so we always grab an extra
1589 * byte at the start, and cut it back later if possible.
1590 */
1591 nbytes = (nbits >> 3) + 1;
Antoine Pitroubf6ecf92012-11-24 20:40:21 +01001592 if (nbytes > 0x7fffffffL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001593 PyErr_SetString(PyExc_OverflowError,
1594 "long too large to pickle");
1595 goto error;
1596 }
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001597 repr = PyBytes_FromStringAndSize(NULL, (Py_ssize_t)nbytes);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001598 if (repr == NULL)
1599 goto error;
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001600 pdata = (unsigned char *)PyBytes_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001601 i = _PyLong_AsByteArray((PyLongObject *)obj,
1602 pdata, nbytes,
1603 1 /* little endian */ , 1 /* signed */ );
1604 if (i < 0)
1605 goto error;
1606 /* If the long is negative, this may be a byte more than
1607 * needed. This is so iff the MSB is all redundant sign
1608 * bits.
1609 */
1610 if (sign < 0 &&
Victor Stinner121aab42011-09-29 23:40:53 +02001611 nbytes > 1 &&
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001612 pdata[nbytes - 1] == 0xff &&
1613 (pdata[nbytes - 2] & 0x80) != 0) {
1614 nbytes--;
1615 }
1616
1617 if (nbytes < 256) {
1618 header[0] = LONG1;
1619 header[1] = (unsigned char)nbytes;
1620 size = 2;
1621 }
1622 else {
1623 header[0] = LONG4;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001624 size = (Py_ssize_t) nbytes;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001625 for (i = 1; i < 5; i++) {
1626 header[i] = (unsigned char)(size & 0xff);
1627 size >>= 8;
1628 }
1629 size = 5;
1630 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001631 if (_Pickler_Write(self, header, size) < 0 ||
1632 _Pickler_Write(self, (char *)pdata, (int)nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001633 goto error;
1634 }
1635 else {
1636 char *string;
1637
Mark Dickinson8dd05142009-01-20 20:43:58 +00001638 /* proto < 2: write the repr and newline. This is quadratic-time (in
1639 the number of digits), in both directions. We add a trailing 'L'
1640 to the repr, for compatibility with Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001641
1642 repr = PyObject_Repr(obj);
1643 if (repr == NULL)
1644 goto error;
1645
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00001646 string = _PyUnicode_AsStringAndSize(repr, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001647 if (string == NULL)
1648 goto error;
1649
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001650 if (_Pickler_Write(self, &long_op, 1) < 0 ||
1651 _Pickler_Write(self, string, size) < 0 ||
1652 _Pickler_Write(self, "L\n", 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001653 goto error;
1654 }
1655
1656 if (0) {
1657 error:
1658 status = -1;
1659 }
1660 Py_XDECREF(repr);
1661
1662 return status;
1663}
1664
1665static int
1666save_float(PicklerObject *self, PyObject *obj)
1667{
1668 double x = PyFloat_AS_DOUBLE((PyFloatObject *)obj);
1669
1670 if (self->bin) {
1671 char pdata[9];
1672 pdata[0] = BINFLOAT;
1673 if (_PyFloat_Pack8(x, (unsigned char *)&pdata[1], 0) < 0)
1674 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001675 if (_Pickler_Write(self, pdata, 9) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001676 return -1;
Victor Stinner121aab42011-09-29 23:40:53 +02001677 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001678 else {
Eric Smith0923d1d2009-04-16 20:16:10 +00001679 int result = -1;
1680 char *buf = NULL;
1681 char op = FLOAT;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001682
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001683 if (_Pickler_Write(self, &op, 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001684 goto done;
1685
Mark Dickinson3e09f432009-04-17 08:41:23 +00001686 buf = PyOS_double_to_string(x, 'g', 17, 0, NULL);
Eric Smith0923d1d2009-04-16 20:16:10 +00001687 if (!buf) {
1688 PyErr_NoMemory();
1689 goto done;
1690 }
1691
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001692 if (_Pickler_Write(self, buf, strlen(buf)) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001693 goto done;
1694
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001695 if (_Pickler_Write(self, "\n", 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001696 goto done;
1697
1698 result = 0;
1699done:
1700 PyMem_Free(buf);
1701 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001702 }
1703
1704 return 0;
1705}
1706
1707static int
1708save_bytes(PicklerObject *self, PyObject *obj)
1709{
1710 if (self->proto < 3) {
1711 /* Older pickle protocols do not have an opcode for pickling bytes
1712 objects. Therefore, we need to fake the copy protocol (i.e.,
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001713 the __reduce__ method) to permit bytes object unpickling.
1714
1715 Here we use a hack to be compatible with Python 2. Since in Python
1716 2 'bytes' is just an alias for 'str' (which has different
1717 parameters than the actual bytes object), we use codecs.encode
1718 to create the appropriate 'str' object when unpickled using
1719 Python 2 *and* the appropriate 'bytes' object when unpickled
1720 using Python 3. Again this is a hack and we don't need to do this
1721 with newer protocols. */
1722 static PyObject *codecs_encode = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001723 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001724 int status;
1725
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001726 if (codecs_encode == NULL) {
1727 PyObject *codecs_module = PyImport_ImportModule("codecs");
1728 if (codecs_module == NULL) {
1729 return -1;
1730 }
1731 codecs_encode = PyObject_GetAttrString(codecs_module, "encode");
1732 Py_DECREF(codecs_module);
1733 if (codecs_encode == NULL) {
1734 return -1;
1735 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001736 }
1737
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001738 if (PyBytes_GET_SIZE(obj) == 0) {
1739 reduce_value = Py_BuildValue("(O())", (PyObject*)&PyBytes_Type);
1740 }
1741 else {
1742 static PyObject *latin1 = NULL;
1743 PyObject *unicode_str =
1744 PyUnicode_DecodeLatin1(PyBytes_AS_STRING(obj),
1745 PyBytes_GET_SIZE(obj),
1746 "strict");
1747 if (unicode_str == NULL)
1748 return -1;
1749 if (latin1 == NULL) {
1750 latin1 = PyUnicode_InternFromString("latin1");
Christian Heimes82e6b942013-06-29 21:37:34 +02001751 if (latin1 == NULL) {
1752 Py_DECREF(unicode_str);
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001753 return -1;
Christian Heimes82e6b942013-06-29 21:37:34 +02001754 }
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001755 }
1756 reduce_value = Py_BuildValue("(O(OO))",
1757 codecs_encode, unicode_str, latin1);
1758 Py_DECREF(unicode_str);
1759 }
1760
1761 if (reduce_value == NULL)
1762 return -1;
1763
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001764 /* save_reduce() will memoize the object automatically. */
1765 status = save_reduce(self, reduce_value, obj);
1766 Py_DECREF(reduce_value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001767 return status;
1768 }
1769 else {
1770 Py_ssize_t size;
1771 char header[5];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001772 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001773
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001774 size = PyBytes_GET_SIZE(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001775 if (size < 0)
1776 return -1;
1777
1778 if (size < 256) {
1779 header[0] = SHORT_BINBYTES;
1780 header[1] = (unsigned char)size;
1781 len = 2;
1782 }
1783 else if (size <= 0xffffffffL) {
1784 header[0] = BINBYTES;
1785 header[1] = (unsigned char)(size & 0xff);
1786 header[2] = (unsigned char)((size >> 8) & 0xff);
1787 header[3] = (unsigned char)((size >> 16) & 0xff);
1788 header[4] = (unsigned char)((size >> 24) & 0xff);
1789 len = 5;
1790 }
1791 else {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001792 PyErr_SetString(PyExc_OverflowError,
Serhiy Storchakaf8def282013-02-16 17:29:56 +02001793 "cannot serialize a bytes object larger than 4 GiB");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001794 return -1; /* string too large */
1795 }
1796
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001797 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001798 return -1;
1799
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001800 if (_Pickler_Write(self, PyBytes_AS_STRING(obj), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001801 return -1;
1802
1803 if (memo_put(self, obj) < 0)
1804 return -1;
1805
1806 return 0;
1807 }
1808}
1809
1810/* A copy of PyUnicode_EncodeRawUnicodeEscape() that also translates
1811 backslash and newline characters to \uXXXX escapes. */
1812static PyObject *
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001813raw_unicode_escape(PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001814{
1815 PyObject *repr, *result;
1816 char *p;
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001817 Py_ssize_t i, size, expandsize;
1818 void *data;
1819 unsigned int kind;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001820
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001821 if (PyUnicode_READY(obj))
1822 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001823
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001824 size = PyUnicode_GET_LENGTH(obj);
1825 data = PyUnicode_DATA(obj);
1826 kind = PyUnicode_KIND(obj);
1827 if (kind == PyUnicode_4BYTE_KIND)
1828 expandsize = 10;
1829 else
1830 expandsize = 6;
Victor Stinner121aab42011-09-29 23:40:53 +02001831
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001832 if (size > PY_SSIZE_T_MAX / expandsize)
1833 return PyErr_NoMemory();
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001834 repr = PyByteArray_FromStringAndSize(NULL, expandsize * size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001835 if (repr == NULL)
1836 return NULL;
1837 if (size == 0)
1838 goto done;
1839
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001840 p = PyByteArray_AS_STRING(repr);
1841 for (i=0; i < size; i++) {
1842 Py_UCS4 ch = PyUnicode_READ(kind, data, i);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001843 /* Map 32-bit characters to '\Uxxxxxxxx' */
1844 if (ch >= 0x10000) {
1845 *p++ = '\\';
1846 *p++ = 'U';
Victor Stinnerf5cff562011-10-14 02:13:11 +02001847 *p++ = Py_hexdigits[(ch >> 28) & 0xf];
1848 *p++ = Py_hexdigits[(ch >> 24) & 0xf];
1849 *p++ = Py_hexdigits[(ch >> 20) & 0xf];
1850 *p++ = Py_hexdigits[(ch >> 16) & 0xf];
1851 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
1852 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
1853 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
1854 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001855 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001856 /* Map 16-bit characters to '\uxxxx' */
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001857 else if (ch >= 256 || ch == '\\' || ch == '\n') {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001858 *p++ = '\\';
1859 *p++ = 'u';
Victor Stinnerf5cff562011-10-14 02:13:11 +02001860 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
1861 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
1862 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
1863 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001864 }
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001865 /* Copy everything else as-is */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001866 else
1867 *p++ = (char) ch;
1868 }
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001869 size = p - PyByteArray_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001870
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001871done:
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001872 result = PyBytes_FromStringAndSize(PyByteArray_AS_STRING(repr), size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001873 Py_DECREF(repr);
1874 return result;
1875}
1876
1877static int
Antoine Pitrou299978d2013-04-07 17:38:11 +02001878write_utf8(PicklerObject *self, char *data, Py_ssize_t size)
1879{
1880 char pdata[5];
1881
1882#if SIZEOF_SIZE_T > 4
1883 if (size > 0xffffffffUL) {
1884 /* string too large */
1885 PyErr_SetString(PyExc_OverflowError,
Antoine Pitrou4b7b0f02013-04-07 23:46:52 +02001886 "cannot serialize a string larger than 4GiB");
Antoine Pitrou299978d2013-04-07 17:38:11 +02001887 return -1;
1888 }
1889#endif
1890
1891 pdata[0] = BINUNICODE;
1892 pdata[1] = (unsigned char)(size & 0xff);
1893 pdata[2] = (unsigned char)((size >> 8) & 0xff);
1894 pdata[3] = (unsigned char)((size >> 16) & 0xff);
1895 pdata[4] = (unsigned char)((size >> 24) & 0xff);
1896
1897 if (_Pickler_Write(self, pdata, sizeof(pdata)) < 0)
1898 return -1;
1899
1900 if (_Pickler_Write(self, data, size) < 0)
1901 return -1;
1902
1903 return 0;
1904}
1905
1906static int
1907write_unicode_binary(PicklerObject *self, PyObject *obj)
1908{
1909 PyObject *encoded = NULL;
1910 Py_ssize_t size;
1911 char *data;
1912 int r;
1913
1914 if (PyUnicode_READY(obj))
1915 return -1;
1916
1917 data = PyUnicode_AsUTF8AndSize(obj, &size);
1918 if (data != NULL)
1919 return write_utf8(self, data, size);
1920
1921 /* Issue #8383: for strings with lone surrogates, fallback on the
1922 "surrogatepass" error handler. */
1923 PyErr_Clear();
1924 encoded = PyUnicode_AsEncodedString(obj, "utf-8", "surrogatepass");
1925 if (encoded == NULL)
1926 return -1;
1927
1928 r = write_utf8(self, PyBytes_AS_STRING(encoded),
1929 PyBytes_GET_SIZE(encoded));
1930 Py_DECREF(encoded);
1931 return r;
1932}
1933
1934static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001935save_unicode(PicklerObject *self, PyObject *obj)
1936{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001937 if (self->bin) {
Antoine Pitrou299978d2013-04-07 17:38:11 +02001938 if (write_unicode_binary(self, obj) < 0)
1939 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001940 }
1941 else {
Antoine Pitrou299978d2013-04-07 17:38:11 +02001942 PyObject *encoded;
1943 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001944 const char unicode_op = UNICODE;
1945
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001946 encoded = raw_unicode_escape(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001947 if (encoded == NULL)
Antoine Pitrou299978d2013-04-07 17:38:11 +02001948 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001949
Antoine Pitrou299978d2013-04-07 17:38:11 +02001950 if (_Pickler_Write(self, &unicode_op, 1) < 0) {
1951 Py_DECREF(encoded);
1952 return -1;
1953 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001954
1955 size = PyBytes_GET_SIZE(encoded);
Antoine Pitrou299978d2013-04-07 17:38:11 +02001956 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0) {
1957 Py_DECREF(encoded);
1958 return -1;
1959 }
1960 Py_DECREF(encoded);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001961
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001962 if (_Pickler_Write(self, "\n", 1) < 0)
Antoine Pitrou299978d2013-04-07 17:38:11 +02001963 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001964 }
1965 if (memo_put(self, obj) < 0)
Antoine Pitrou299978d2013-04-07 17:38:11 +02001966 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001967
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001968 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001969}
1970
1971/* A helper for save_tuple. Push the len elements in tuple t on the stack. */
1972static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001973store_tuple_elements(PicklerObject *self, PyObject *t, Py_ssize_t len)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001974{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001975 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001976
1977 assert(PyTuple_Size(t) == len);
1978
1979 for (i = 0; i < len; i++) {
1980 PyObject *element = PyTuple_GET_ITEM(t, i);
1981
1982 if (element == NULL)
1983 return -1;
1984 if (save(self, element, 0) < 0)
1985 return -1;
1986 }
1987
1988 return 0;
1989}
1990
1991/* Tuples are ubiquitous in the pickle protocols, so many techniques are
1992 * used across protocols to minimize the space needed to pickle them.
1993 * Tuples are also the only builtin immutable type that can be recursive
1994 * (a tuple can be reached from itself), and that requires some subtle
1995 * magic so that it works in all cases. IOW, this is a long routine.
1996 */
1997static int
1998save_tuple(PicklerObject *self, PyObject *obj)
1999{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002000 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002001
2002 const char mark_op = MARK;
2003 const char tuple_op = TUPLE;
2004 const char pop_op = POP;
2005 const char pop_mark_op = POP_MARK;
2006 const char len2opcode[] = {EMPTY_TUPLE, TUPLE1, TUPLE2, TUPLE3};
2007
2008 if ((len = PyTuple_Size(obj)) < 0)
2009 return -1;
2010
2011 if (len == 0) {
2012 char pdata[2];
2013
2014 if (self->proto) {
2015 pdata[0] = EMPTY_TUPLE;
2016 len = 1;
2017 }
2018 else {
2019 pdata[0] = MARK;
2020 pdata[1] = TUPLE;
2021 len = 2;
2022 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002023 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002024 return -1;
2025 return 0;
2026 }
2027
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002028 /* The tuple isn't in the memo now. If it shows up there after
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002029 * saving the tuple elements, the tuple must be recursive, in
2030 * which case we'll pop everything we put on the stack, and fetch
2031 * its value from the memo.
2032 */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002033 if (len <= 3 && self->proto >= 2) {
2034 /* Use TUPLE{1,2,3} opcodes. */
2035 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002036 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002037
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002038 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002039 /* pop the len elements */
2040 for (i = 0; i < len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002041 if (_Pickler_Write(self, &pop_op, 1) < 0)
2042 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002043 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002044 if (memo_get(self, obj) < 0)
2045 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002046
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002047 return 0;
2048 }
2049 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002050 if (_Pickler_Write(self, len2opcode + len, 1) < 0)
2051 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002052 }
2053 goto memoize;
2054 }
2055
2056 /* proto < 2 and len > 0, or proto >= 2 and len > 3.
2057 * Generate MARK e1 e2 ... TUPLE
2058 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002059 if (_Pickler_Write(self, &mark_op, 1) < 0)
2060 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002061
2062 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002063 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002064
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002065 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002066 /* pop the stack stuff we pushed */
2067 if (self->bin) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002068 if (_Pickler_Write(self, &pop_mark_op, 1) < 0)
2069 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002070 }
2071 else {
2072 /* Note that we pop one more than len, to remove
2073 * the MARK too.
2074 */
2075 for (i = 0; i <= len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002076 if (_Pickler_Write(self, &pop_op, 1) < 0)
2077 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002078 }
2079 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002080 if (memo_get(self, obj) < 0)
2081 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002082
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002083 return 0;
2084 }
2085 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002086 if (_Pickler_Write(self, &tuple_op, 1) < 0)
2087 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002088 }
2089
2090 memoize:
2091 if (memo_put(self, obj) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002092 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002093
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002094 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002095}
2096
2097/* iter is an iterator giving items, and we batch up chunks of
2098 * MARK item item ... item APPENDS
2099 * opcode sequences. Calling code should have arranged to first create an
2100 * empty list, or list-like object, for the APPENDS to operate on.
2101 * Returns 0 on success, <0 on error.
2102 */
2103static int
2104batch_list(PicklerObject *self, PyObject *iter)
2105{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002106 PyObject *obj = NULL;
2107 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002108 int i, n;
2109
2110 const char mark_op = MARK;
2111 const char append_op = APPEND;
2112 const char appends_op = APPENDS;
2113
2114 assert(iter != NULL);
2115
2116 /* XXX: I think this function could be made faster by avoiding the
2117 iterator interface and fetching objects directly from list using
2118 PyList_GET_ITEM.
2119 */
2120
2121 if (self->proto == 0) {
2122 /* APPENDS isn't available; do one at a time. */
2123 for (;;) {
2124 obj = PyIter_Next(iter);
2125 if (obj == NULL) {
2126 if (PyErr_Occurred())
2127 return -1;
2128 break;
2129 }
2130 i = save(self, obj, 0);
2131 Py_DECREF(obj);
2132 if (i < 0)
2133 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002134 if (_Pickler_Write(self, &append_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002135 return -1;
2136 }
2137 return 0;
2138 }
2139
2140 /* proto > 0: write in batches of BATCHSIZE. */
2141 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002142 /* Get first item */
2143 firstitem = PyIter_Next(iter);
2144 if (firstitem == NULL) {
2145 if (PyErr_Occurred())
2146 goto error;
2147
2148 /* nothing more to add */
2149 break;
2150 }
2151
2152 /* Try to get a second item */
2153 obj = PyIter_Next(iter);
2154 if (obj == NULL) {
2155 if (PyErr_Occurred())
2156 goto error;
2157
2158 /* Only one item to write */
2159 if (save(self, firstitem, 0) < 0)
2160 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002161 if (_Pickler_Write(self, &append_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002162 goto error;
2163 Py_CLEAR(firstitem);
2164 break;
2165 }
2166
2167 /* More than one item to write */
2168
2169 /* Pump out MARK, items, APPENDS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002170 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002171 goto error;
2172
2173 if (save(self, firstitem, 0) < 0)
2174 goto error;
2175 Py_CLEAR(firstitem);
2176 n = 1;
2177
2178 /* Fetch and save up to BATCHSIZE items */
2179 while (obj) {
2180 if (save(self, obj, 0) < 0)
2181 goto error;
2182 Py_CLEAR(obj);
2183 n += 1;
2184
2185 if (n == BATCHSIZE)
2186 break;
2187
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002188 obj = PyIter_Next(iter);
2189 if (obj == NULL) {
2190 if (PyErr_Occurred())
2191 goto error;
2192 break;
2193 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002194 }
2195
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002196 if (_Pickler_Write(self, &appends_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002197 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002198
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002199 } while (n == BATCHSIZE);
2200 return 0;
2201
2202 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002203 Py_XDECREF(firstitem);
2204 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002205 return -1;
2206}
2207
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002208/* This is a variant of batch_list() above, specialized for lists (with no
2209 * support for list subclasses). Like batch_list(), we batch up chunks of
2210 * MARK item item ... item APPENDS
2211 * opcode sequences. Calling code should have arranged to first create an
2212 * empty list, or list-like object, for the APPENDS to operate on.
2213 * Returns 0 on success, -1 on error.
2214 *
2215 * This version is considerably faster than batch_list(), if less general.
2216 *
2217 * Note that this only works for protocols > 0.
2218 */
2219static int
2220batch_list_exact(PicklerObject *self, PyObject *obj)
2221{
2222 PyObject *item = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002223 Py_ssize_t this_batch, total;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002224
2225 const char append_op = APPEND;
2226 const char appends_op = APPENDS;
2227 const char mark_op = MARK;
2228
2229 assert(obj != NULL);
2230 assert(self->proto > 0);
2231 assert(PyList_CheckExact(obj));
2232
2233 if (PyList_GET_SIZE(obj) == 1) {
2234 item = PyList_GET_ITEM(obj, 0);
2235 if (save(self, item, 0) < 0)
2236 return -1;
2237 if (_Pickler_Write(self, &append_op, 1) < 0)
2238 return -1;
2239 return 0;
2240 }
2241
2242 /* Write in batches of BATCHSIZE. */
2243 total = 0;
2244 do {
2245 this_batch = 0;
2246 if (_Pickler_Write(self, &mark_op, 1) < 0)
2247 return -1;
2248 while (total < PyList_GET_SIZE(obj)) {
2249 item = PyList_GET_ITEM(obj, total);
2250 if (save(self, item, 0) < 0)
2251 return -1;
2252 total++;
2253 if (++this_batch == BATCHSIZE)
2254 break;
2255 }
2256 if (_Pickler_Write(self, &appends_op, 1) < 0)
2257 return -1;
2258
2259 } while (total < PyList_GET_SIZE(obj));
2260
2261 return 0;
2262}
2263
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002264static int
2265save_list(PicklerObject *self, PyObject *obj)
2266{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002267 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002268 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002269 int status = 0;
2270
2271 if (self->fast && !fast_save_enter(self, obj))
2272 goto error;
2273
2274 /* Create an empty list. */
2275 if (self->bin) {
2276 header[0] = EMPTY_LIST;
2277 len = 1;
2278 }
2279 else {
2280 header[0] = MARK;
2281 header[1] = LIST;
2282 len = 2;
2283 }
2284
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002285 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002286 goto error;
2287
2288 /* Get list length, and bow out early if empty. */
2289 if ((len = PyList_Size(obj)) < 0)
2290 goto error;
2291
2292 if (memo_put(self, obj) < 0)
2293 goto error;
2294
2295 if (len != 0) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002296 /* Materialize the list elements. */
2297 if (PyList_CheckExact(obj) && self->proto > 0) {
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002298 if (Py_EnterRecursiveCall(" while pickling an object"))
2299 goto error;
2300 status = batch_list_exact(self, obj);
2301 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002302 } else {
2303 PyObject *iter = PyObject_GetIter(obj);
2304 if (iter == NULL)
2305 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002306
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002307 if (Py_EnterRecursiveCall(" while pickling an object")) {
2308 Py_DECREF(iter);
2309 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002310 }
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002311 status = batch_list(self, iter);
2312 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002313 Py_DECREF(iter);
2314 }
2315 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002316 if (0) {
2317 error:
2318 status = -1;
2319 }
2320
2321 if (self->fast && !fast_save_leave(self, obj))
2322 status = -1;
2323
2324 return status;
2325}
2326
2327/* iter is an iterator giving (key, value) pairs, and we batch up chunks of
2328 * MARK key value ... key value SETITEMS
2329 * opcode sequences. Calling code should have arranged to first create an
2330 * empty dict, or dict-like object, for the SETITEMS to operate on.
2331 * Returns 0 on success, <0 on error.
2332 *
2333 * This is very much like batch_list(). The difference between saving
2334 * elements directly, and picking apart two-tuples, is so long-winded at
2335 * the C level, though, that attempts to combine these routines were too
2336 * ugly to bear.
2337 */
2338static int
2339batch_dict(PicklerObject *self, PyObject *iter)
2340{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002341 PyObject *obj = NULL;
2342 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002343 int i, n;
2344
2345 const char mark_op = MARK;
2346 const char setitem_op = SETITEM;
2347 const char setitems_op = SETITEMS;
2348
2349 assert(iter != NULL);
2350
2351 if (self->proto == 0) {
2352 /* SETITEMS isn't available; do one at a time. */
2353 for (;;) {
2354 obj = PyIter_Next(iter);
2355 if (obj == NULL) {
2356 if (PyErr_Occurred())
2357 return -1;
2358 break;
2359 }
2360 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2361 PyErr_SetString(PyExc_TypeError, "dict items "
2362 "iterator must return 2-tuples");
2363 return -1;
2364 }
2365 i = save(self, PyTuple_GET_ITEM(obj, 0), 0);
2366 if (i >= 0)
2367 i = save(self, PyTuple_GET_ITEM(obj, 1), 0);
2368 Py_DECREF(obj);
2369 if (i < 0)
2370 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002371 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002372 return -1;
2373 }
2374 return 0;
2375 }
2376
2377 /* proto > 0: write in batches of BATCHSIZE. */
2378 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002379 /* Get first item */
2380 firstitem = PyIter_Next(iter);
2381 if (firstitem == NULL) {
2382 if (PyErr_Occurred())
2383 goto error;
2384
2385 /* nothing more to add */
2386 break;
2387 }
2388 if (!PyTuple_Check(firstitem) || PyTuple_Size(firstitem) != 2) {
2389 PyErr_SetString(PyExc_TypeError, "dict items "
2390 "iterator must return 2-tuples");
2391 goto error;
2392 }
2393
2394 /* Try to get a second item */
2395 obj = PyIter_Next(iter);
2396 if (obj == NULL) {
2397 if (PyErr_Occurred())
2398 goto error;
2399
2400 /* Only one item to write */
2401 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2402 goto error;
2403 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2404 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002405 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002406 goto error;
2407 Py_CLEAR(firstitem);
2408 break;
2409 }
2410
2411 /* More than one item to write */
2412
2413 /* Pump out MARK, items, SETITEMS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002414 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002415 goto error;
2416
2417 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2418 goto error;
2419 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2420 goto error;
2421 Py_CLEAR(firstitem);
2422 n = 1;
2423
2424 /* Fetch and save up to BATCHSIZE items */
2425 while (obj) {
2426 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2427 PyErr_SetString(PyExc_TypeError, "dict items "
2428 "iterator must return 2-tuples");
2429 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002430 }
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002431 if (save(self, PyTuple_GET_ITEM(obj, 0), 0) < 0 ||
2432 save(self, PyTuple_GET_ITEM(obj, 1), 0) < 0)
2433 goto error;
2434 Py_CLEAR(obj);
2435 n += 1;
2436
2437 if (n == BATCHSIZE)
2438 break;
2439
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002440 obj = PyIter_Next(iter);
2441 if (obj == NULL) {
2442 if (PyErr_Occurred())
2443 goto error;
2444 break;
2445 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002446 }
2447
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002448 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002449 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002450
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002451 } while (n == BATCHSIZE);
2452 return 0;
2453
2454 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002455 Py_XDECREF(firstitem);
2456 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002457 return -1;
2458}
2459
Collin Winter5c9b02d2009-05-25 05:43:30 +00002460/* This is a variant of batch_dict() above that specializes for dicts, with no
2461 * support for dict subclasses. Like batch_dict(), we batch up chunks of
2462 * MARK key value ... key value SETITEMS
2463 * opcode sequences. Calling code should have arranged to first create an
2464 * empty dict, or dict-like object, for the SETITEMS to operate on.
2465 * Returns 0 on success, -1 on error.
2466 *
2467 * Note that this currently doesn't work for protocol 0.
2468 */
2469static int
2470batch_dict_exact(PicklerObject *self, PyObject *obj)
2471{
2472 PyObject *key = NULL, *value = NULL;
2473 int i;
2474 Py_ssize_t dict_size, ppos = 0;
2475
Alexandre Vassalottif70b1292009-05-25 18:00:52 +00002476 const char mark_op = MARK;
2477 const char setitem_op = SETITEM;
2478 const char setitems_op = SETITEMS;
Collin Winter5c9b02d2009-05-25 05:43:30 +00002479
2480 assert(obj != NULL);
2481 assert(self->proto > 0);
2482
2483 dict_size = PyDict_Size(obj);
2484
2485 /* Special-case len(d) == 1 to save space. */
2486 if (dict_size == 1) {
2487 PyDict_Next(obj, &ppos, &key, &value);
2488 if (save(self, key, 0) < 0)
2489 return -1;
2490 if (save(self, value, 0) < 0)
2491 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002492 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002493 return -1;
2494 return 0;
2495 }
2496
2497 /* Write in batches of BATCHSIZE. */
2498 do {
2499 i = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002500 if (_Pickler_Write(self, &mark_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002501 return -1;
2502 while (PyDict_Next(obj, &ppos, &key, &value)) {
2503 if (save(self, key, 0) < 0)
2504 return -1;
2505 if (save(self, value, 0) < 0)
2506 return -1;
2507 if (++i == BATCHSIZE)
2508 break;
2509 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002510 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002511 return -1;
2512 if (PyDict_Size(obj) != dict_size) {
2513 PyErr_Format(
2514 PyExc_RuntimeError,
2515 "dictionary changed size during iteration");
2516 return -1;
2517 }
2518
2519 } while (i == BATCHSIZE);
2520 return 0;
2521}
2522
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002523static int
2524save_dict(PicklerObject *self, PyObject *obj)
2525{
2526 PyObject *items, *iter;
2527 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002528 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002529 int status = 0;
2530
2531 if (self->fast && !fast_save_enter(self, obj))
2532 goto error;
2533
2534 /* Create an empty dict. */
2535 if (self->bin) {
2536 header[0] = EMPTY_DICT;
2537 len = 1;
2538 }
2539 else {
2540 header[0] = MARK;
2541 header[1] = DICT;
2542 len = 2;
2543 }
2544
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002545 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002546 goto error;
2547
2548 /* Get dict size, and bow out early if empty. */
2549 if ((len = PyDict_Size(obj)) < 0)
2550 goto error;
2551
2552 if (memo_put(self, obj) < 0)
2553 goto error;
2554
2555 if (len != 0) {
2556 /* Save the dict items. */
Collin Winter5c9b02d2009-05-25 05:43:30 +00002557 if (PyDict_CheckExact(obj) && self->proto > 0) {
2558 /* We can take certain shortcuts if we know this is a dict and
2559 not a dict subclass. */
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002560 if (Py_EnterRecursiveCall(" while pickling an object"))
2561 goto error;
2562 status = batch_dict_exact(self, obj);
2563 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002564 } else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02002565 _Py_IDENTIFIER(items);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02002566
2567 items = _PyObject_CallMethodId(obj, &PyId_items, "()");
Collin Winter5c9b02d2009-05-25 05:43:30 +00002568 if (items == NULL)
2569 goto error;
2570 iter = PyObject_GetIter(items);
2571 Py_DECREF(items);
2572 if (iter == NULL)
2573 goto error;
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002574 if (Py_EnterRecursiveCall(" while pickling an object")) {
2575 Py_DECREF(iter);
2576 goto error;
2577 }
Collin Winter5c9b02d2009-05-25 05:43:30 +00002578 status = batch_dict(self, iter);
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002579 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002580 Py_DECREF(iter);
2581 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002582 }
2583
2584 if (0) {
2585 error:
2586 status = -1;
2587 }
2588
2589 if (self->fast && !fast_save_leave(self, obj))
2590 status = -1;
2591
2592 return status;
2593}
2594
2595static int
2596save_global(PicklerObject *self, PyObject *obj, PyObject *name)
2597{
2598 static PyObject *name_str = NULL;
2599 PyObject *global_name = NULL;
2600 PyObject *module_name = NULL;
2601 PyObject *module = NULL;
2602 PyObject *cls;
2603 int status = 0;
2604
2605 const char global_op = GLOBAL;
2606
2607 if (name_str == NULL) {
2608 name_str = PyUnicode_InternFromString("__name__");
2609 if (name_str == NULL)
2610 goto error;
2611 }
2612
2613 if (name) {
2614 global_name = name;
2615 Py_INCREF(global_name);
2616 }
2617 else {
2618 global_name = PyObject_GetAttr(obj, name_str);
2619 if (global_name == NULL)
2620 goto error;
2621 }
2622
2623 module_name = whichmodule(obj, global_name);
2624 if (module_name == NULL)
2625 goto error;
2626
2627 /* XXX: Change to use the import C API directly with level=0 to disallow
2628 relative imports.
2629
2630 XXX: PyImport_ImportModuleLevel could be used. However, this bypasses
2631 builtins.__import__. Therefore, _pickle, unlike pickle.py, will ignore
2632 custom import functions (IMHO, this would be a nice security
2633 feature). The import C API would need to be extended to support the
2634 extra parameters of __import__ to fix that. */
2635 module = PyImport_Import(module_name);
2636 if (module == NULL) {
2637 PyErr_Format(PicklingError,
2638 "Can't pickle %R: import of module %R failed",
2639 obj, module_name);
2640 goto error;
2641 }
2642 cls = PyObject_GetAttr(module, global_name);
2643 if (cls == NULL) {
2644 PyErr_Format(PicklingError,
2645 "Can't pickle %R: attribute lookup %S.%S failed",
2646 obj, module_name, global_name);
2647 goto error;
2648 }
2649 if (cls != obj) {
2650 Py_DECREF(cls);
2651 PyErr_Format(PicklingError,
2652 "Can't pickle %R: it's not the same object as %S.%S",
2653 obj, module_name, global_name);
2654 goto error;
2655 }
2656 Py_DECREF(cls);
2657
2658 if (self->proto >= 2) {
2659 /* See whether this is in the extension registry, and if
2660 * so generate an EXT opcode.
2661 */
2662 PyObject *code_obj; /* extension code as Python object */
2663 long code; /* extension code as C value */
2664 char pdata[5];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002665 Py_ssize_t n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002666
2667 PyTuple_SET_ITEM(two_tuple, 0, module_name);
2668 PyTuple_SET_ITEM(two_tuple, 1, global_name);
2669 code_obj = PyDict_GetItem(extension_registry, two_tuple);
2670 /* The object is not registered in the extension registry.
2671 This is the most likely code path. */
2672 if (code_obj == NULL)
2673 goto gen_global;
2674
2675 /* XXX: pickle.py doesn't check neither the type, nor the range
2676 of the value returned by the extension_registry. It should for
2677 consistency. */
2678
2679 /* Verify code_obj has the right type and value. */
2680 if (!PyLong_Check(code_obj)) {
2681 PyErr_Format(PicklingError,
2682 "Can't pickle %R: extension code %R isn't an integer",
2683 obj, code_obj);
2684 goto error;
2685 }
2686 code = PyLong_AS_LONG(code_obj);
2687 if (code <= 0 || code > 0x7fffffffL) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002688 if (!PyErr_Occurred())
2689 PyErr_Format(PicklingError,
2690 "Can't pickle %R: extension code %ld is out of range",
2691 obj, code);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002692 goto error;
2693 }
2694
2695 /* Generate an EXT opcode. */
2696 if (code <= 0xff) {
2697 pdata[0] = EXT1;
2698 pdata[1] = (unsigned char)code;
2699 n = 2;
2700 }
2701 else if (code <= 0xffff) {
2702 pdata[0] = EXT2;
2703 pdata[1] = (unsigned char)(code & 0xff);
2704 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2705 n = 3;
2706 }
2707 else {
2708 pdata[0] = EXT4;
2709 pdata[1] = (unsigned char)(code & 0xff);
2710 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2711 pdata[3] = (unsigned char)((code >> 16) & 0xff);
2712 pdata[4] = (unsigned char)((code >> 24) & 0xff);
2713 n = 5;
2714 }
2715
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002716 if (_Pickler_Write(self, pdata, n) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002717 goto error;
2718 }
2719 else {
2720 /* Generate a normal global opcode if we are using a pickle
2721 protocol <= 2, or if the object is not registered in the
2722 extension registry. */
2723 PyObject *encoded;
2724 PyObject *(*unicode_encoder)(PyObject *);
2725
2726 gen_global:
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002727 if (_Pickler_Write(self, &global_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002728 goto error;
2729
2730 /* Since Python 3.0 now supports non-ASCII identifiers, we encode both
2731 the module name and the global name using UTF-8. We do so only when
2732 we are using the pickle protocol newer than version 3. This is to
2733 ensure compatibility with older Unpickler running on Python 2.x. */
2734 if (self->proto >= 3) {
2735 unicode_encoder = PyUnicode_AsUTF8String;
2736 }
2737 else {
2738 unicode_encoder = PyUnicode_AsASCIIString;
2739 }
2740
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00002741 /* For protocol < 3 and if the user didn't request against doing so,
2742 we convert module names to the old 2.x module names. */
2743 if (self->fix_imports) {
2744 PyObject *key;
2745 PyObject *item;
2746
2747 key = PyTuple_Pack(2, module_name, global_name);
2748 if (key == NULL)
2749 goto error;
2750 item = PyDict_GetItemWithError(name_mapping_3to2, key);
2751 Py_DECREF(key);
2752 if (item) {
2753 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
2754 PyErr_Format(PyExc_RuntimeError,
2755 "_compat_pickle.REVERSE_NAME_MAPPING values "
2756 "should be 2-tuples, not %.200s",
2757 Py_TYPE(item)->tp_name);
2758 goto error;
2759 }
2760 Py_CLEAR(module_name);
2761 Py_CLEAR(global_name);
2762 module_name = PyTuple_GET_ITEM(item, 0);
2763 global_name = PyTuple_GET_ITEM(item, 1);
2764 if (!PyUnicode_Check(module_name) ||
2765 !PyUnicode_Check(global_name)) {
2766 PyErr_Format(PyExc_RuntimeError,
2767 "_compat_pickle.REVERSE_NAME_MAPPING values "
2768 "should be pairs of str, not (%.200s, %.200s)",
2769 Py_TYPE(module_name)->tp_name,
2770 Py_TYPE(global_name)->tp_name);
2771 goto error;
2772 }
2773 Py_INCREF(module_name);
2774 Py_INCREF(global_name);
2775 }
2776 else if (PyErr_Occurred()) {
2777 goto error;
2778 }
2779
2780 item = PyDict_GetItemWithError(import_mapping_3to2, module_name);
2781 if (item) {
2782 if (!PyUnicode_Check(item)) {
2783 PyErr_Format(PyExc_RuntimeError,
2784 "_compat_pickle.REVERSE_IMPORT_MAPPING values "
2785 "should be strings, not %.200s",
2786 Py_TYPE(item)->tp_name);
2787 goto error;
2788 }
2789 Py_CLEAR(module_name);
2790 module_name = item;
2791 Py_INCREF(module_name);
2792 }
2793 else if (PyErr_Occurred()) {
2794 goto error;
2795 }
2796 }
2797
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002798 /* Save the name of the module. */
2799 encoded = unicode_encoder(module_name);
2800 if (encoded == NULL) {
2801 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2802 PyErr_Format(PicklingError,
2803 "can't pickle module identifier '%S' using "
2804 "pickle protocol %i", module_name, self->proto);
2805 goto error;
2806 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002807 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002808 PyBytes_GET_SIZE(encoded)) < 0) {
2809 Py_DECREF(encoded);
2810 goto error;
2811 }
2812 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002813 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002814 goto error;
2815
2816 /* Save the name of the module. */
2817 encoded = unicode_encoder(global_name);
2818 if (encoded == NULL) {
2819 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2820 PyErr_Format(PicklingError,
2821 "can't pickle global identifier '%S' using "
2822 "pickle protocol %i", global_name, self->proto);
2823 goto error;
2824 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002825 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002826 PyBytes_GET_SIZE(encoded)) < 0) {
2827 Py_DECREF(encoded);
2828 goto error;
2829 }
2830 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002831 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002832 goto error;
2833
2834 /* Memoize the object. */
2835 if (memo_put(self, obj) < 0)
2836 goto error;
2837 }
2838
2839 if (0) {
2840 error:
2841 status = -1;
2842 }
2843 Py_XDECREF(module_name);
2844 Py_XDECREF(global_name);
2845 Py_XDECREF(module);
2846
2847 return status;
2848}
2849
2850static int
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002851save_ellipsis(PicklerObject *self, PyObject *obj)
2852{
Łukasz Langadbd78252012-03-12 22:59:11 +01002853 PyObject *str = PyUnicode_FromString("Ellipsis");
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002854 int res;
Łukasz Langadbd78252012-03-12 22:59:11 +01002855 if (str == NULL)
Łukasz Langacad1a072012-03-12 23:41:07 +01002856 return -1;
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002857 res = save_global(self, Py_Ellipsis, str);
2858 Py_DECREF(str);
2859 return res;
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002860}
2861
2862static int
2863save_notimplemented(PicklerObject *self, PyObject *obj)
2864{
Łukasz Langadbd78252012-03-12 22:59:11 +01002865 PyObject *str = PyUnicode_FromString("NotImplemented");
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002866 int res;
Łukasz Langadbd78252012-03-12 22:59:11 +01002867 if (str == NULL)
Łukasz Langacad1a072012-03-12 23:41:07 +01002868 return -1;
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002869 res = save_global(self, Py_NotImplemented, str);
2870 Py_DECREF(str);
2871 return res;
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002872}
2873
2874static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002875save_pers(PicklerObject *self, PyObject *obj, PyObject *func)
2876{
2877 PyObject *pid = NULL;
2878 int status = 0;
2879
2880 const char persid_op = PERSID;
2881 const char binpersid_op = BINPERSID;
2882
2883 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002884 pid = _Pickler_FastCall(self, func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002885 if (pid == NULL)
2886 return -1;
2887
2888 if (pid != Py_None) {
2889 if (self->bin) {
2890 if (save(self, pid, 1) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002891 _Pickler_Write(self, &binpersid_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002892 goto error;
2893 }
2894 else {
2895 PyObject *pid_str = NULL;
2896 char *pid_ascii_bytes;
2897 Py_ssize_t size;
2898
2899 pid_str = PyObject_Str(pid);
2900 if (pid_str == NULL)
2901 goto error;
2902
2903 /* XXX: Should it check whether the persistent id only contains
2904 ASCII characters? And what if the pid contains embedded
2905 newlines? */
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00002906 pid_ascii_bytes = _PyUnicode_AsStringAndSize(pid_str, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002907 Py_DECREF(pid_str);
2908 if (pid_ascii_bytes == NULL)
2909 goto error;
2910
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002911 if (_Pickler_Write(self, &persid_op, 1) < 0 ||
2912 _Pickler_Write(self, pid_ascii_bytes, size) < 0 ||
2913 _Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002914 goto error;
2915 }
2916 status = 1;
2917 }
2918
2919 if (0) {
2920 error:
2921 status = -1;
2922 }
2923 Py_XDECREF(pid);
2924
2925 return status;
2926}
2927
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002928static PyObject *
2929get_class(PyObject *obj)
2930{
2931 PyObject *cls;
2932 static PyObject *str_class;
2933
2934 if (str_class == NULL) {
2935 str_class = PyUnicode_InternFromString("__class__");
2936 if (str_class == NULL)
2937 return NULL;
2938 }
2939 cls = PyObject_GetAttr(obj, str_class);
2940 if (cls == NULL) {
2941 if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
2942 PyErr_Clear();
2943 cls = (PyObject *) Py_TYPE(obj);
2944 Py_INCREF(cls);
2945 }
2946 }
2947 return cls;
2948}
2949
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002950/* We're saving obj, and args is the 2-thru-5 tuple returned by the
2951 * appropriate __reduce__ method for obj.
2952 */
2953static int
2954save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
2955{
2956 PyObject *callable;
2957 PyObject *argtup;
2958 PyObject *state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002959 PyObject *listitems = Py_None;
2960 PyObject *dictitems = Py_None;
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002961 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002962
2963 int use_newobj = self->proto >= 2;
2964
2965 const char reduce_op = REDUCE;
2966 const char build_op = BUILD;
2967 const char newobj_op = NEWOBJ;
2968
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002969 size = PyTuple_Size(args);
2970 if (size < 2 || size > 5) {
2971 PyErr_SetString(PicklingError, "tuple returned by "
2972 "__reduce__ must contain 2 through 5 elements");
2973 return -1;
2974 }
2975
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002976 if (!PyArg_UnpackTuple(args, "save_reduce", 2, 5,
2977 &callable, &argtup, &state, &listitems, &dictitems))
2978 return -1;
2979
2980 if (!PyCallable_Check(callable)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002981 PyErr_SetString(PicklingError, "first item of the tuple "
2982 "returned by __reduce__ must be callable");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002983 return -1;
2984 }
2985 if (!PyTuple_Check(argtup)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002986 PyErr_SetString(PicklingError, "second item of the tuple "
2987 "returned by __reduce__ must be a tuple");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002988 return -1;
2989 }
2990
2991 if (state == Py_None)
2992 state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002993
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002994 if (listitems == Py_None)
2995 listitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002996 else if (!PyIter_Check(listitems)) {
Alexandre Vassalotti00d83f22013-04-14 01:28:01 -07002997 PyErr_Format(PicklingError, "fourth element of the tuple "
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002998 "returned by __reduce__ must be an iterator, not %s",
2999 Py_TYPE(listitems)->tp_name);
3000 return -1;
3001 }
3002
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003003 if (dictitems == Py_None)
3004 dictitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003005 else if (!PyIter_Check(dictitems)) {
Alexandre Vassalotti00d83f22013-04-14 01:28:01 -07003006 PyErr_Format(PicklingError, "fifth element of the tuple "
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003007 "returned by __reduce__ must be an iterator, not %s",
3008 Py_TYPE(dictitems)->tp_name);
3009 return -1;
3010 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003011
3012 /* Protocol 2 special case: if callable's name is __newobj__, use
3013 NEWOBJ. */
3014 if (use_newobj) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003015 static PyObject *newobj_str = NULL, *name_str = NULL;
3016 PyObject *name;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003017
3018 if (newobj_str == NULL) {
3019 newobj_str = PyUnicode_InternFromString("__newobj__");
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003020 name_str = PyUnicode_InternFromString("__name__");
3021 if (newobj_str == NULL || name_str == NULL)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003022 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003023 }
3024
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003025 name = PyObject_GetAttr(callable, name_str);
3026 if (name == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003027 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3028 PyErr_Clear();
3029 else
3030 return -1;
3031 use_newobj = 0;
3032 }
3033 else {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003034 use_newobj = PyUnicode_Check(name) &&
3035 PyUnicode_Compare(name, newobj_str) == 0;
3036 Py_DECREF(name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003037 }
3038 }
3039 if (use_newobj) {
3040 PyObject *cls;
3041 PyObject *newargtup;
3042 PyObject *obj_class;
3043 int p;
3044
3045 /* Sanity checks. */
3046 if (Py_SIZE(argtup) < 1) {
3047 PyErr_SetString(PicklingError, "__newobj__ arglist is empty");
3048 return -1;
3049 }
3050
3051 cls = PyTuple_GET_ITEM(argtup, 0);
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003052 if (!PyType_Check(cls)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003053 PyErr_SetString(PicklingError, "args[0] from "
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003054 "__newobj__ args is not a type");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003055 return -1;
3056 }
3057
3058 if (obj != NULL) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003059 obj_class = get_class(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003060 p = obj_class != cls; /* true iff a problem */
3061 Py_DECREF(obj_class);
3062 if (p) {
3063 PyErr_SetString(PicklingError, "args[0] from "
3064 "__newobj__ args has the wrong class");
3065 return -1;
3066 }
3067 }
3068 /* XXX: These calls save() are prone to infinite recursion. Imagine
3069 what happen if the value returned by the __reduce__() method of
3070 some extension type contains another object of the same type. Ouch!
3071
3072 Here is a quick example, that I ran into, to illustrate what I
3073 mean:
3074
3075 >>> import pickle, copyreg
3076 >>> copyreg.dispatch_table.pop(complex)
3077 >>> pickle.dumps(1+2j)
3078 Traceback (most recent call last):
3079 ...
3080 RuntimeError: maximum recursion depth exceeded
3081
3082 Removing the complex class from copyreg.dispatch_table made the
3083 __reduce_ex__() method emit another complex object:
3084
3085 >>> (1+1j).__reduce_ex__(2)
3086 (<function __newobj__ at 0xb7b71c3c>,
3087 (<class 'complex'>, (1+1j)), None, None, None)
3088
3089 Thus when save() was called on newargstup (the 2nd item) recursion
3090 ensued. Of course, the bug was in the complex class which had a
3091 broken __getnewargs__() that emitted another complex object. But,
3092 the point, here, is it is quite easy to end up with a broken reduce
3093 function. */
3094
3095 /* Save the class and its __new__ arguments. */
3096 if (save(self, cls, 0) < 0)
3097 return -1;
3098
3099 newargtup = PyTuple_GetSlice(argtup, 1, Py_SIZE(argtup));
3100 if (newargtup == NULL)
3101 return -1;
3102
3103 p = save(self, newargtup, 0);
3104 Py_DECREF(newargtup);
3105 if (p < 0)
3106 return -1;
3107
3108 /* Add NEWOBJ opcode. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003109 if (_Pickler_Write(self, &newobj_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003110 return -1;
3111 }
3112 else { /* Not using NEWOBJ. */
3113 if (save(self, callable, 0) < 0 ||
3114 save(self, argtup, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003115 _Pickler_Write(self, &reduce_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003116 return -1;
3117 }
3118
3119 /* obj can be NULL when save_reduce() is used directly. A NULL obj means
3120 the caller do not want to memoize the object. Not particularly useful,
3121 but that is to mimic the behavior save_reduce() in pickle.py when
3122 obj is None. */
3123 if (obj && memo_put(self, obj) < 0)
3124 return -1;
3125
3126 if (listitems && batch_list(self, listitems) < 0)
3127 return -1;
3128
3129 if (dictitems && batch_dict(self, dictitems) < 0)
3130 return -1;
3131
3132 if (state) {
Victor Stinner121aab42011-09-29 23:40:53 +02003133 if (save(self, state, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003134 _Pickler_Write(self, &build_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003135 return -1;
3136 }
3137
3138 return 0;
3139}
3140
3141static int
3142save(PicklerObject *self, PyObject *obj, int pers_save)
3143{
3144 PyTypeObject *type;
3145 PyObject *reduce_func = NULL;
3146 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003147 int status = 0;
3148
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00003149 if (Py_EnterRecursiveCall(" while pickling an object"))
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003150 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003151
3152 /* The extra pers_save argument is necessary to avoid calling save_pers()
3153 on its returned object. */
3154 if (!pers_save && self->pers_func) {
3155 /* save_pers() returns:
3156 -1 to signal an error;
3157 0 if it did nothing successfully;
3158 1 if a persistent id was saved.
3159 */
3160 if ((status = save_pers(self, obj, self->pers_func)) != 0)
3161 goto done;
3162 }
3163
3164 type = Py_TYPE(obj);
3165
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003166 /* The old cPickle had an optimization that used switch-case statement
3167 dispatching on the first letter of the type name. This has was removed
3168 since benchmarks shown that this optimization was actually slowing
3169 things down. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003170
3171 /* Atom types; these aren't memoized, so don't check the memo. */
3172
3173 if (obj == Py_None) {
3174 status = save_none(self, obj);
3175 goto done;
3176 }
Łukasz Langaf3078fb2012-03-12 19:46:12 +01003177 else if (obj == Py_Ellipsis) {
3178 status = save_ellipsis(self, obj);
3179 goto done;
3180 }
3181 else if (obj == Py_NotImplemented) {
3182 status = save_notimplemented(self, obj);
3183 goto done;
3184 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003185 else if (obj == Py_False || obj == Py_True) {
3186 status = save_bool(self, obj);
3187 goto done;
3188 }
3189 else if (type == &PyLong_Type) {
3190 status = save_long(self, obj);
3191 goto done;
3192 }
3193 else if (type == &PyFloat_Type) {
3194 status = save_float(self, obj);
3195 goto done;
3196 }
3197
3198 /* Check the memo to see if it has the object. If so, generate
3199 a GET (or BINGET) opcode, instead of pickling the object
3200 once again. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003201 if (PyMemoTable_Get(self->memo, obj)) {
3202 if (memo_get(self, obj) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003203 goto error;
3204 goto done;
3205 }
3206
3207 if (type == &PyBytes_Type) {
3208 status = save_bytes(self, obj);
3209 goto done;
3210 }
3211 else if (type == &PyUnicode_Type) {
3212 status = save_unicode(self, obj);
3213 goto done;
3214 }
3215 else if (type == &PyDict_Type) {
3216 status = save_dict(self, obj);
3217 goto done;
3218 }
3219 else if (type == &PyList_Type) {
3220 status = save_list(self, obj);
3221 goto done;
3222 }
3223 else if (type == &PyTuple_Type) {
3224 status = save_tuple(self, obj);
3225 goto done;
3226 }
3227 else if (type == &PyType_Type) {
3228 status = save_global(self, obj, NULL);
3229 goto done;
3230 }
3231 else if (type == &PyFunction_Type) {
3232 status = save_global(self, obj, NULL);
3233 if (status < 0 && PyErr_ExceptionMatches(PickleError)) {
3234 /* fall back to reduce */
3235 PyErr_Clear();
3236 }
3237 else {
3238 goto done;
3239 }
3240 }
3241 else if (type == &PyCFunction_Type) {
3242 status = save_global(self, obj, NULL);
3243 goto done;
3244 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003245
3246 /* XXX: This part needs some unit tests. */
3247
3248 /* Get a reduction callable, and call it. This may come from
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003249 * self.dispatch_table, copyreg.dispatch_table, the object's
3250 * __reduce_ex__ method, or the object's __reduce__ method.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003251 */
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003252 if (self->dispatch_table == NULL) {
3253 reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
3254 /* PyDict_GetItem() unlike PyObject_GetItem() and
3255 PyObject_GetAttr() returns a borrowed ref */
3256 Py_XINCREF(reduce_func);
3257 } else {
3258 reduce_func = PyObject_GetItem(self->dispatch_table, (PyObject *)type);
3259 if (reduce_func == NULL) {
3260 if (PyErr_ExceptionMatches(PyExc_KeyError))
3261 PyErr_Clear();
3262 else
3263 goto error;
3264 }
3265 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003266 if (reduce_func != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003267 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003268 reduce_value = _Pickler_FastCall(self, reduce_func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003269 }
Antoine Pitrouffd41d92011-10-04 09:23:04 +02003270 else if (PyType_IsSubtype(type, &PyType_Type)) {
3271 status = save_global(self, obj, NULL);
3272 goto done;
3273 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003274 else {
3275 static PyObject *reduce_str = NULL;
3276 static PyObject *reduce_ex_str = NULL;
3277
3278 /* Cache the name of the reduce methods. */
3279 if (reduce_str == NULL) {
3280 reduce_str = PyUnicode_InternFromString("__reduce__");
3281 if (reduce_str == NULL)
3282 goto error;
3283 reduce_ex_str = PyUnicode_InternFromString("__reduce_ex__");
3284 if (reduce_ex_str == NULL)
3285 goto error;
3286 }
3287
3288 /* XXX: If the __reduce__ method is defined, __reduce_ex__ is
3289 automatically defined as __reduce__. While this is convenient, this
3290 make it impossible to know which method was actually called. Of
3291 course, this is not a big deal. But still, it would be nice to let
3292 the user know which method was called when something go
3293 wrong. Incidentally, this means if __reduce_ex__ is not defined, we
3294 don't actually have to check for a __reduce__ method. */
3295
3296 /* Check for a __reduce_ex__ method. */
3297 reduce_func = PyObject_GetAttr(obj, reduce_ex_str);
3298 if (reduce_func != NULL) {
3299 PyObject *proto;
3300 proto = PyLong_FromLong(self->proto);
3301 if (proto != NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003302 reduce_value = _Pickler_FastCall(self, reduce_func, proto);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003303 }
3304 }
3305 else {
3306 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3307 PyErr_Clear();
3308 else
3309 goto error;
3310 /* Check for a __reduce__ method. */
3311 reduce_func = PyObject_GetAttr(obj, reduce_str);
3312 if (reduce_func != NULL) {
3313 reduce_value = PyObject_Call(reduce_func, empty_tuple, NULL);
3314 }
3315 else {
3316 PyErr_Format(PicklingError, "can't pickle '%.200s' object: %R",
3317 type->tp_name, obj);
3318 goto error;
3319 }
3320 }
3321 }
3322
3323 if (reduce_value == NULL)
3324 goto error;
3325
3326 if (PyUnicode_Check(reduce_value)) {
3327 status = save_global(self, obj, reduce_value);
3328 goto done;
3329 }
3330
3331 if (!PyTuple_Check(reduce_value)) {
3332 PyErr_SetString(PicklingError,
3333 "__reduce__ must return a string or tuple");
3334 goto error;
3335 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003336
3337 status = save_reduce(self, reduce_value, obj);
3338
3339 if (0) {
3340 error:
3341 status = -1;
3342 }
3343 done:
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003344 Py_LeaveRecursiveCall();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003345 Py_XDECREF(reduce_func);
3346 Py_XDECREF(reduce_value);
3347
3348 return status;
3349}
3350
3351static int
3352dump(PicklerObject *self, PyObject *obj)
3353{
3354 const char stop_op = STOP;
3355
3356 if (self->proto >= 2) {
3357 char header[2];
3358
3359 header[0] = PROTO;
3360 assert(self->proto >= 0 && self->proto < 256);
3361 header[1] = (unsigned char)self->proto;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003362 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003363 return -1;
3364 }
3365
3366 if (save(self, obj, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003367 _Pickler_Write(self, &stop_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003368 return -1;
3369
3370 return 0;
3371}
3372
3373PyDoc_STRVAR(Pickler_clear_memo_doc,
3374"clear_memo() -> None. Clears the pickler's \"memo\"."
3375"\n"
3376"The memo is the data structure that remembers which objects the\n"
3377"pickler has already seen, so that shared or recursive objects are\n"
3378"pickled by reference and not by value. This method is useful when\n"
3379"re-using picklers.");
3380
3381static PyObject *
3382Pickler_clear_memo(PicklerObject *self)
3383{
3384 if (self->memo)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003385 PyMemoTable_Clear(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003386
3387 Py_RETURN_NONE;
3388}
3389
3390PyDoc_STRVAR(Pickler_dump_doc,
3391"dump(obj) -> None. Write a pickled representation of obj to the open file.");
3392
3393static PyObject *
3394Pickler_dump(PicklerObject *self, PyObject *args)
3395{
3396 PyObject *obj;
3397
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003398 /* Check whether the Pickler was initialized correctly (issue3664).
3399 Developers often forget to call __init__() in their subclasses, which
3400 would trigger a segfault without this check. */
3401 if (self->write == NULL) {
Victor Stinner121aab42011-09-29 23:40:53 +02003402 PyErr_Format(PicklingError,
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003403 "Pickler.__init__() was not called by %s.__init__()",
3404 Py_TYPE(self)->tp_name);
3405 return NULL;
3406 }
3407
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003408 if (!PyArg_ParseTuple(args, "O:dump", &obj))
3409 return NULL;
3410
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003411 if (_Pickler_ClearBuffer(self) < 0)
3412 return NULL;
3413
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003414 if (dump(self, obj) < 0)
3415 return NULL;
3416
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003417 if (_Pickler_FlushToFile(self) < 0)
3418 return NULL;
3419
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003420 Py_RETURN_NONE;
3421}
3422
3423static struct PyMethodDef Pickler_methods[] = {
3424 {"dump", (PyCFunction)Pickler_dump, METH_VARARGS,
3425 Pickler_dump_doc},
3426 {"clear_memo", (PyCFunction)Pickler_clear_memo, METH_NOARGS,
3427 Pickler_clear_memo_doc},
3428 {NULL, NULL} /* sentinel */
3429};
3430
3431static void
3432Pickler_dealloc(PicklerObject *self)
3433{
3434 PyObject_GC_UnTrack(self);
3435
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003436 Py_XDECREF(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003437 Py_XDECREF(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003438 Py_XDECREF(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003439 Py_XDECREF(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003440 Py_XDECREF(self->arg);
3441 Py_XDECREF(self->fast_memo);
3442
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003443 PyMemoTable_Del(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003444
3445 Py_TYPE(self)->tp_free((PyObject *)self);
3446}
3447
3448static int
3449Pickler_traverse(PicklerObject *self, visitproc visit, void *arg)
3450{
3451 Py_VISIT(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003452 Py_VISIT(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003453 Py_VISIT(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003454 Py_VISIT(self->arg);
3455 Py_VISIT(self->fast_memo);
3456 return 0;
3457}
3458
3459static int
3460Pickler_clear(PicklerObject *self)
3461{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003462 Py_CLEAR(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003463 Py_CLEAR(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003464 Py_CLEAR(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003465 Py_CLEAR(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003466 Py_CLEAR(self->arg);
3467 Py_CLEAR(self->fast_memo);
3468
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003469 if (self->memo != NULL) {
3470 PyMemoTable *memo = self->memo;
3471 self->memo = NULL;
3472 PyMemoTable_Del(memo);
3473 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003474 return 0;
3475}
3476
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003477
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003478PyDoc_STRVAR(Pickler_doc,
3479"Pickler(file, protocol=None)"
3480"\n"
3481"This takes a binary file for writing a pickle data stream.\n"
3482"\n"
3483"The optional protocol argument tells the pickler to use the\n"
3484"given protocol; supported protocols are 0, 1, 2, 3. The default\n"
3485"protocol is 3; a backward-incompatible protocol designed for\n"
3486"Python 3.0.\n"
3487"\n"
3488"Specifying a negative protocol version selects the highest\n"
3489"protocol version supported. The higher the protocol used, the\n"
3490"more recent the version of Python needed to read the pickle\n"
3491"produced.\n"
3492"\n"
3493"The file argument must have a write() method that accepts a single\n"
3494"bytes argument. It can thus be a file object opened for binary\n"
3495"writing, a io.BytesIO instance, or any other custom object that\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003496"meets this interface.\n"
3497"\n"
3498"If fix_imports is True and protocol is less than 3, pickle will try to\n"
3499"map the new Python 3.x names to the old module names used in Python\n"
3500"2.x, so that the pickle data stream is readable with Python 2.x.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003501
3502static int
3503Pickler_init(PicklerObject *self, PyObject *args, PyObject *kwds)
3504{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003505 static char *kwlist[] = {"file", "protocol", "fix_imports", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003506 PyObject *file;
3507 PyObject *proto_obj = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003508 PyObject *fix_imports = Py_True;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02003509 _Py_IDENTIFIER(persistent_id);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003510 _Py_IDENTIFIER(dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003511
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003512 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:Pickler",
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003513 kwlist, &file, &proto_obj, &fix_imports))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003514 return -1;
3515
3516 /* In case of multiple __init__() calls, clear previous content. */
3517 if (self->write != NULL)
3518 (void)Pickler_clear(self);
3519
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003520 if (_Pickler_SetProtocol(self, proto_obj, fix_imports) < 0)
3521 return -1;
3522
3523 if (_Pickler_SetOutputStream(self, file) < 0)
3524 return -1;
3525
3526 /* memo and output_buffer may have already been created in _Pickler_New */
3527 if (self->memo == NULL) {
3528 self->memo = PyMemoTable_New();
3529 if (self->memo == NULL)
3530 return -1;
3531 }
3532 self->output_len = 0;
3533 if (self->output_buffer == NULL) {
3534 self->max_output_len = WRITE_BUF_SIZE;
3535 self->output_buffer = PyBytes_FromStringAndSize(NULL,
3536 self->max_output_len);
3537 if (self->output_buffer == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003538 return -1;
3539 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003540
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003541 self->arg = NULL;
3542 self->fast = 0;
3543 self->fast_nesting = 0;
3544 self->fast_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003545 self->pers_func = NULL;
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02003546 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_id)) {
3547 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
3548 &PyId_persistent_id);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003549 if (self->pers_func == NULL)
3550 return -1;
3551 }
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003552 self->dispatch_table = NULL;
3553 if (_PyObject_HasAttrId((PyObject *)self, &PyId_dispatch_table)) {
3554 self->dispatch_table = _PyObject_GetAttrId((PyObject *)self,
3555 &PyId_dispatch_table);
3556 if (self->dispatch_table == NULL)
3557 return -1;
3558 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003559 return 0;
3560}
3561
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003562/* Define a proxy object for the Pickler's internal memo object. This is to
3563 * avoid breaking code like:
3564 * pickler.memo.clear()
3565 * and
3566 * pickler.memo = saved_memo
3567 * Is this a good idea? Not really, but we don't want to break code that uses
3568 * it. Note that we don't implement the entire mapping API here. This is
3569 * intentional, as these should be treated as black-box implementation details.
3570 */
3571
3572typedef struct {
3573 PyObject_HEAD
3574 PicklerObject *pickler; /* Pickler whose memo table we're proxying. */
3575} PicklerMemoProxyObject;
3576
3577PyDoc_STRVAR(pmp_clear_doc,
3578"memo.clear() -> None. Remove all items from memo.");
3579
3580static PyObject *
3581pmp_clear(PicklerMemoProxyObject *self)
3582{
3583 if (self->pickler->memo)
3584 PyMemoTable_Clear(self->pickler->memo);
3585 Py_RETURN_NONE;
3586}
3587
3588PyDoc_STRVAR(pmp_copy_doc,
3589"memo.copy() -> new_memo. Copy the memo to a new object.");
3590
3591static PyObject *
3592pmp_copy(PicklerMemoProxyObject *self)
3593{
3594 Py_ssize_t i;
3595 PyMemoTable *memo;
3596 PyObject *new_memo = PyDict_New();
3597 if (new_memo == NULL)
3598 return NULL;
3599
3600 memo = self->pickler->memo;
3601 for (i = 0; i < memo->mt_allocated; ++i) {
3602 PyMemoEntry entry = memo->mt_table[i];
3603 if (entry.me_key != NULL) {
3604 int status;
3605 PyObject *key, *value;
3606
3607 key = PyLong_FromVoidPtr(entry.me_key);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003608 value = Py_BuildValue("nO", entry.me_value, entry.me_key);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003609
3610 if (key == NULL || value == NULL) {
3611 Py_XDECREF(key);
3612 Py_XDECREF(value);
3613 goto error;
3614 }
3615 status = PyDict_SetItem(new_memo, key, value);
3616 Py_DECREF(key);
3617 Py_DECREF(value);
3618 if (status < 0)
3619 goto error;
3620 }
3621 }
3622 return new_memo;
3623
3624 error:
3625 Py_XDECREF(new_memo);
3626 return NULL;
3627}
3628
3629PyDoc_STRVAR(pmp_reduce_doc,
3630"memo.__reduce__(). Pickling support.");
3631
3632static PyObject *
3633pmp_reduce(PicklerMemoProxyObject *self, PyObject *args)
3634{
3635 PyObject *reduce_value, *dict_args;
3636 PyObject *contents = pmp_copy(self);
3637 if (contents == NULL)
3638 return NULL;
3639
3640 reduce_value = PyTuple_New(2);
3641 if (reduce_value == NULL) {
3642 Py_DECREF(contents);
3643 return NULL;
3644 }
3645 dict_args = PyTuple_New(1);
3646 if (dict_args == NULL) {
3647 Py_DECREF(contents);
3648 Py_DECREF(reduce_value);
3649 return NULL;
3650 }
3651 PyTuple_SET_ITEM(dict_args, 0, contents);
3652 Py_INCREF((PyObject *)&PyDict_Type);
3653 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
3654 PyTuple_SET_ITEM(reduce_value, 1, dict_args);
3655 return reduce_value;
3656}
3657
3658static PyMethodDef picklerproxy_methods[] = {
3659 {"clear", (PyCFunction)pmp_clear, METH_NOARGS, pmp_clear_doc},
3660 {"copy", (PyCFunction)pmp_copy, METH_NOARGS, pmp_copy_doc},
3661 {"__reduce__", (PyCFunction)pmp_reduce, METH_VARARGS, pmp_reduce_doc},
3662 {NULL, NULL} /* sentinel */
3663};
3664
3665static void
3666PicklerMemoProxy_dealloc(PicklerMemoProxyObject *self)
3667{
3668 PyObject_GC_UnTrack(self);
3669 Py_XDECREF(self->pickler);
3670 PyObject_GC_Del((PyObject *)self);
3671}
3672
3673static int
3674PicklerMemoProxy_traverse(PicklerMemoProxyObject *self,
3675 visitproc visit, void *arg)
3676{
3677 Py_VISIT(self->pickler);
3678 return 0;
3679}
3680
3681static int
3682PicklerMemoProxy_clear(PicklerMemoProxyObject *self)
3683{
3684 Py_CLEAR(self->pickler);
3685 return 0;
3686}
3687
3688static PyTypeObject PicklerMemoProxyType = {
3689 PyVarObject_HEAD_INIT(NULL, 0)
3690 "_pickle.PicklerMemoProxy", /*tp_name*/
3691 sizeof(PicklerMemoProxyObject), /*tp_basicsize*/
3692 0,
3693 (destructor)PicklerMemoProxy_dealloc, /* tp_dealloc */
3694 0, /* tp_print */
3695 0, /* tp_getattr */
3696 0, /* tp_setattr */
3697 0, /* tp_compare */
3698 0, /* tp_repr */
3699 0, /* tp_as_number */
3700 0, /* tp_as_sequence */
3701 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00003702 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003703 0, /* tp_call */
3704 0, /* tp_str */
3705 PyObject_GenericGetAttr, /* tp_getattro */
3706 PyObject_GenericSetAttr, /* tp_setattro */
3707 0, /* tp_as_buffer */
3708 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3709 0, /* tp_doc */
3710 (traverseproc)PicklerMemoProxy_traverse, /* tp_traverse */
3711 (inquiry)PicklerMemoProxy_clear, /* tp_clear */
3712 0, /* tp_richcompare */
3713 0, /* tp_weaklistoffset */
3714 0, /* tp_iter */
3715 0, /* tp_iternext */
3716 picklerproxy_methods, /* tp_methods */
3717};
3718
3719static PyObject *
3720PicklerMemoProxy_New(PicklerObject *pickler)
3721{
3722 PicklerMemoProxyObject *self;
3723
3724 self = PyObject_GC_New(PicklerMemoProxyObject, &PicklerMemoProxyType);
3725 if (self == NULL)
3726 return NULL;
3727 Py_INCREF(pickler);
3728 self->pickler = pickler;
3729 PyObject_GC_Track(self);
3730 return (PyObject *)self;
3731}
3732
3733/*****************************************************************************/
3734
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003735static PyObject *
3736Pickler_get_memo(PicklerObject *self)
3737{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003738 return PicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003739}
3740
3741static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003742Pickler_set_memo(PicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003743{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003744 PyMemoTable *new_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003745
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003746 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003747 PyErr_SetString(PyExc_TypeError,
3748 "attribute deletion is not supported");
3749 return -1;
3750 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003751
3752 if (Py_TYPE(obj) == &PicklerMemoProxyType) {
3753 PicklerObject *pickler =
3754 ((PicklerMemoProxyObject *)obj)->pickler;
3755
3756 new_memo = PyMemoTable_Copy(pickler->memo);
3757 if (new_memo == NULL)
3758 return -1;
3759 }
3760 else if (PyDict_Check(obj)) {
3761 Py_ssize_t i = 0;
3762 PyObject *key, *value;
3763
3764 new_memo = PyMemoTable_New();
3765 if (new_memo == NULL)
3766 return -1;
3767
3768 while (PyDict_Next(obj, &i, &key, &value)) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003769 Py_ssize_t memo_id;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003770 PyObject *memo_obj;
3771
3772 if (!PyTuple_Check(value) || Py_SIZE(value) != 2) {
3773 PyErr_SetString(PyExc_TypeError,
3774 "'memo' values must be 2-item tuples");
3775 goto error;
3776 }
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003777 memo_id = PyLong_AsSsize_t(PyTuple_GET_ITEM(value, 0));
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003778 if (memo_id == -1 && PyErr_Occurred())
3779 goto error;
3780 memo_obj = PyTuple_GET_ITEM(value, 1);
3781 if (PyMemoTable_Set(new_memo, memo_obj, memo_id) < 0)
3782 goto error;
3783 }
3784 }
3785 else {
3786 PyErr_Format(PyExc_TypeError,
3787 "'memo' attribute must be an PicklerMemoProxy object"
3788 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003789 return -1;
3790 }
3791
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003792 PyMemoTable_Del(self->memo);
3793 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003794
3795 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003796
3797 error:
3798 if (new_memo)
3799 PyMemoTable_Del(new_memo);
3800 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003801}
3802
3803static PyObject *
3804Pickler_get_persid(PicklerObject *self)
3805{
3806 if (self->pers_func == NULL)
3807 PyErr_SetString(PyExc_AttributeError, "persistent_id");
3808 else
3809 Py_INCREF(self->pers_func);
3810 return self->pers_func;
3811}
3812
3813static int
3814Pickler_set_persid(PicklerObject *self, PyObject *value)
3815{
3816 PyObject *tmp;
3817
3818 if (value == NULL) {
3819 PyErr_SetString(PyExc_TypeError,
3820 "attribute deletion is not supported");
3821 return -1;
3822 }
3823 if (!PyCallable_Check(value)) {
3824 PyErr_SetString(PyExc_TypeError,
3825 "persistent_id must be a callable taking one argument");
3826 return -1;
3827 }
3828
3829 tmp = self->pers_func;
3830 Py_INCREF(value);
3831 self->pers_func = value;
3832 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
3833
3834 return 0;
3835}
3836
3837static PyMemberDef Pickler_members[] = {
3838 {"bin", T_INT, offsetof(PicklerObject, bin)},
3839 {"fast", T_INT, offsetof(PicklerObject, fast)},
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003840 {"dispatch_table", T_OBJECT_EX, offsetof(PicklerObject, dispatch_table)},
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003841 {NULL}
3842};
3843
3844static PyGetSetDef Pickler_getsets[] = {
3845 {"memo", (getter)Pickler_get_memo,
3846 (setter)Pickler_set_memo},
3847 {"persistent_id", (getter)Pickler_get_persid,
3848 (setter)Pickler_set_persid},
3849 {NULL}
3850};
3851
3852static PyTypeObject Pickler_Type = {
3853 PyVarObject_HEAD_INIT(NULL, 0)
3854 "_pickle.Pickler" , /*tp_name*/
3855 sizeof(PicklerObject), /*tp_basicsize*/
3856 0, /*tp_itemsize*/
3857 (destructor)Pickler_dealloc, /*tp_dealloc*/
3858 0, /*tp_print*/
3859 0, /*tp_getattr*/
3860 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00003861 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003862 0, /*tp_repr*/
3863 0, /*tp_as_number*/
3864 0, /*tp_as_sequence*/
3865 0, /*tp_as_mapping*/
3866 0, /*tp_hash*/
3867 0, /*tp_call*/
3868 0, /*tp_str*/
3869 0, /*tp_getattro*/
3870 0, /*tp_setattro*/
3871 0, /*tp_as_buffer*/
3872 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3873 Pickler_doc, /*tp_doc*/
3874 (traverseproc)Pickler_traverse, /*tp_traverse*/
3875 (inquiry)Pickler_clear, /*tp_clear*/
3876 0, /*tp_richcompare*/
3877 0, /*tp_weaklistoffset*/
3878 0, /*tp_iter*/
3879 0, /*tp_iternext*/
3880 Pickler_methods, /*tp_methods*/
3881 Pickler_members, /*tp_members*/
3882 Pickler_getsets, /*tp_getset*/
3883 0, /*tp_base*/
3884 0, /*tp_dict*/
3885 0, /*tp_descr_get*/
3886 0, /*tp_descr_set*/
3887 0, /*tp_dictoffset*/
3888 (initproc)Pickler_init, /*tp_init*/
3889 PyType_GenericAlloc, /*tp_alloc*/
3890 PyType_GenericNew, /*tp_new*/
3891 PyObject_GC_Del, /*tp_free*/
3892 0, /*tp_is_gc*/
3893};
3894
Victor Stinner121aab42011-09-29 23:40:53 +02003895/* Temporary helper for calling self.find_class().
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003896
3897 XXX: It would be nice to able to avoid Python function call overhead, by
3898 using directly the C version of find_class(), when find_class() is not
3899 overridden by a subclass. Although, this could become rather hackish. A
3900 simpler optimization would be to call the C function when self is not a
3901 subclass instance. */
3902static PyObject *
3903find_class(UnpicklerObject *self, PyObject *module_name, PyObject *global_name)
3904{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02003905 _Py_IDENTIFIER(find_class);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02003906
3907 return _PyObject_CallMethodId((PyObject *)self, &PyId_find_class, "OO",
3908 module_name, global_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003909}
3910
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003911static Py_ssize_t
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003912marker(UnpicklerObject *self)
3913{
3914 if (self->num_marks < 1) {
3915 PyErr_SetString(UnpicklingError, "could not find MARK");
3916 return -1;
3917 }
3918
3919 return self->marks[--self->num_marks];
3920}
3921
3922static int
3923load_none(UnpicklerObject *self)
3924{
3925 PDATA_APPEND(self->stack, Py_None, -1);
3926 return 0;
3927}
3928
3929static int
3930bad_readline(void)
3931{
3932 PyErr_SetString(UnpicklingError, "pickle data was truncated");
3933 return -1;
3934}
3935
3936static int
3937load_int(UnpicklerObject *self)
3938{
3939 PyObject *value;
3940 char *endptr, *s;
3941 Py_ssize_t len;
3942 long x;
3943
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003944 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003945 return -1;
3946 if (len < 2)
3947 return bad_readline();
3948
3949 errno = 0;
Victor Stinner121aab42011-09-29 23:40:53 +02003950 /* XXX: Should the base argument of strtol() be explicitly set to 10?
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003951 XXX(avassalotti): Should this uses PyOS_strtol()? */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003952 x = strtol(s, &endptr, 0);
3953
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003954 if (errno || (*endptr != '\n' && *endptr != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003955 /* Hm, maybe we've got something long. Let's try reading
3956 * it as a Python long object. */
3957 errno = 0;
3958 /* XXX: Same thing about the base here. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003959 value = PyLong_FromString(s, NULL, 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003960 if (value == NULL) {
3961 PyErr_SetString(PyExc_ValueError,
3962 "could not convert string to int");
3963 return -1;
3964 }
3965 }
3966 else {
3967 if (len == 3 && (x == 0 || x == 1)) {
3968 if ((value = PyBool_FromLong(x)) == NULL)
3969 return -1;
3970 }
3971 else {
3972 if ((value = PyLong_FromLong(x)) == NULL)
3973 return -1;
3974 }
3975 }
3976
3977 PDATA_PUSH(self->stack, value, -1);
3978 return 0;
3979}
3980
3981static int
3982load_bool(UnpicklerObject *self, PyObject *boolean)
3983{
3984 assert(boolean == Py_True || boolean == Py_False);
3985 PDATA_APPEND(self->stack, boolean, -1);
3986 return 0;
3987}
3988
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003989/* s contains x bytes of an unsigned little-endian integer. Return its value
3990 * as a C Py_ssize_t, or -1 if it's higher than PY_SSIZE_T_MAX.
3991 */
3992static Py_ssize_t
3993calc_binsize(char *bytes, int size)
3994{
3995 unsigned char *s = (unsigned char *)bytes;
3996 size_t x = 0;
3997
3998 assert(size == 4);
3999
4000 x = (size_t) s[0];
4001 x |= (size_t) s[1] << 8;
4002 x |= (size_t) s[2] << 16;
4003 x |= (size_t) s[3] << 24;
4004
4005 if (x > PY_SSIZE_T_MAX)
4006 return -1;
4007 else
4008 return (Py_ssize_t) x;
4009}
4010
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004011/* s contains x bytes of a little-endian integer. Return its value as a
4012 * C int. Obscure: when x is 1 or 2, this is an unsigned little-endian
4013 * int, but when x is 4 it's a signed one. This is an historical source
4014 * of x-platform bugs.
4015 */
4016static long
4017calc_binint(char *bytes, int size)
4018{
4019 unsigned char *s = (unsigned char *)bytes;
4020 int i = size;
4021 long x = 0;
4022
4023 for (i = 0; i < size; i++) {
4024 x |= (long)s[i] << (i * 8);
4025 }
4026
4027 /* Unlike BININT1 and BININT2, BININT (more accurately BININT4)
4028 * is signed, so on a box with longs bigger than 4 bytes we need
4029 * to extend a BININT's sign bit to the full width.
4030 */
4031 if (SIZEOF_LONG > 4 && size == 4) {
4032 x |= -(x & (1L << 31));
4033 }
4034
4035 return x;
4036}
4037
4038static int
4039load_binintx(UnpicklerObject *self, char *s, int size)
4040{
4041 PyObject *value;
4042 long x;
4043
4044 x = calc_binint(s, size);
4045
4046 if ((value = PyLong_FromLong(x)) == NULL)
4047 return -1;
4048
4049 PDATA_PUSH(self->stack, value, -1);
4050 return 0;
4051}
4052
4053static int
4054load_binint(UnpicklerObject *self)
4055{
4056 char *s;
4057
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004058 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004059 return -1;
4060
4061 return load_binintx(self, s, 4);
4062}
4063
4064static int
4065load_binint1(UnpicklerObject *self)
4066{
4067 char *s;
4068
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004069 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004070 return -1;
4071
4072 return load_binintx(self, s, 1);
4073}
4074
4075static int
4076load_binint2(UnpicklerObject *self)
4077{
4078 char *s;
4079
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004080 if (_Unpickler_Read(self, &s, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004081 return -1;
4082
4083 return load_binintx(self, s, 2);
4084}
4085
4086static int
4087load_long(UnpicklerObject *self)
4088{
4089 PyObject *value;
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004090 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004091 Py_ssize_t len;
4092
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004093 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004094 return -1;
4095 if (len < 2)
4096 return bad_readline();
4097
Mark Dickinson8dd05142009-01-20 20:43:58 +00004098 /* s[len-2] will usually be 'L' (and s[len-1] is '\n'); we need to remove
4099 the 'L' before calling PyLong_FromString. In order to maintain
4100 compatibility with Python 3.0.0, we don't actually *require*
4101 the 'L' to be present. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004102 if (s[len-2] == 'L')
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004103 s[len-2] = '\0';
Alexandre Vassalottie4bccb72009-01-24 01:47:57 +00004104 /* XXX: Should the base argument explicitly set to 10? */
4105 value = PyLong_FromString(s, NULL, 0);
Mark Dickinson8dd05142009-01-20 20:43:58 +00004106 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004107 return -1;
4108
4109 PDATA_PUSH(self->stack, value, -1);
4110 return 0;
4111}
4112
4113/* 'size' bytes contain the # of bytes of little-endian 256's-complement
4114 * data following.
4115 */
4116static int
4117load_counted_long(UnpicklerObject *self, int size)
4118{
4119 PyObject *value;
4120 char *nbytes;
4121 char *pdata;
4122
4123 assert(size == 1 || size == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004124 if (_Unpickler_Read(self, &nbytes, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004125 return -1;
4126
4127 size = calc_binint(nbytes, size);
4128 if (size < 0) {
4129 /* Corrupt or hostile pickle -- we never write one like this */
4130 PyErr_SetString(UnpicklingError,
4131 "LONG pickle has negative byte count");
4132 return -1;
4133 }
4134
4135 if (size == 0)
4136 value = PyLong_FromLong(0L);
4137 else {
4138 /* Read the raw little-endian bytes and convert. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004139 if (_Unpickler_Read(self, &pdata, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004140 return -1;
4141 value = _PyLong_FromByteArray((unsigned char *)pdata, (size_t)size,
4142 1 /* little endian */ , 1 /* signed */ );
4143 }
4144 if (value == NULL)
4145 return -1;
4146 PDATA_PUSH(self->stack, value, -1);
4147 return 0;
4148}
4149
4150static int
4151load_float(UnpicklerObject *self)
4152{
4153 PyObject *value;
4154 char *endptr, *s;
4155 Py_ssize_t len;
4156 double d;
4157
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004158 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004159 return -1;
4160 if (len < 2)
4161 return bad_readline();
4162
4163 errno = 0;
Mark Dickinson725bfd82009-05-03 20:33:40 +00004164 d = PyOS_string_to_double(s, &endptr, PyExc_OverflowError);
4165 if (d == -1.0 && PyErr_Occurred())
4166 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004167 if ((endptr[0] != '\n') && (endptr[0] != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004168 PyErr_SetString(PyExc_ValueError, "could not convert string to float");
4169 return -1;
4170 }
Mark Dickinson725bfd82009-05-03 20:33:40 +00004171 value = PyFloat_FromDouble(d);
4172 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004173 return -1;
4174
4175 PDATA_PUSH(self->stack, value, -1);
4176 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004177}
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004178
4179static int
4180load_binfloat(UnpicklerObject *self)
4181{
4182 PyObject *value;
4183 double x;
4184 char *s;
4185
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004186 if (_Unpickler_Read(self, &s, 8) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004187 return -1;
4188
4189 x = _PyFloat_Unpack8((unsigned char *)s, 0);
4190 if (x == -1.0 && PyErr_Occurred())
4191 return -1;
4192
4193 if ((value = PyFloat_FromDouble(x)) == NULL)
4194 return -1;
4195
4196 PDATA_PUSH(self->stack, value, -1);
4197 return 0;
4198}
4199
4200static int
4201load_string(UnpicklerObject *self)
4202{
4203 PyObject *bytes;
4204 PyObject *str = NULL;
4205 Py_ssize_t len;
4206 char *s, *p;
4207
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004208 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004209 return -1;
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004210 /* Strip the newline */
4211 len--;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004212 /* Strip outermost quotes */
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004213 if (len >= 2 && s[0] == s[len - 1] && (s[0] == '\'' || s[0] == '"')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004214 p = s + 1;
4215 len -= 2;
4216 }
4217 else {
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004218 PyErr_SetString(UnpicklingError,
4219 "the STRING opcode argument must be quoted");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004220 return -1;
4221 }
Alexandre Vassalotti7c5e0942013-04-15 23:14:55 -07004222 assert(len >= 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004223
4224 /* Use the PyBytes API to decode the string, since that is what is used
4225 to encode, and then coerce the result to Unicode. */
4226 bytes = PyBytes_DecodeEscape(p, len, NULL, 0, NULL);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004227 if (bytes == NULL)
4228 return -1;
4229 str = PyUnicode_FromEncodedObject(bytes, self->encoding, self->errors);
4230 Py_DECREF(bytes);
4231 if (str == NULL)
4232 return -1;
4233
4234 PDATA_PUSH(self->stack, str, -1);
4235 return 0;
4236}
4237
4238static int
4239load_binbytes(UnpicklerObject *self)
4240{
4241 PyObject *bytes;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004242 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004243 char *s;
4244
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004245 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004246 return -1;
4247
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004248 x = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004249 if (x < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004250 PyErr_Format(PyExc_OverflowError,
4251 "BINBYTES exceeds system's maximum size of %zd bytes",
Alexandre Vassalotticc757172013-04-14 02:25:10 -07004252 PY_SSIZE_T_MAX);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004253 return -1;
4254 }
4255
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004256 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004257 return -1;
4258 bytes = PyBytes_FromStringAndSize(s, x);
4259 if (bytes == NULL)
4260 return -1;
4261
4262 PDATA_PUSH(self->stack, bytes, -1);
4263 return 0;
4264}
4265
4266static int
4267load_short_binbytes(UnpicklerObject *self)
4268{
4269 PyObject *bytes;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004270 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004271 char *s;
4272
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004273 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004274 return -1;
4275
4276 x = (unsigned char)s[0];
4277
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004278 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004279 return -1;
4280
4281 bytes = PyBytes_FromStringAndSize(s, x);
4282 if (bytes == NULL)
4283 return -1;
4284
4285 PDATA_PUSH(self->stack, bytes, -1);
4286 return 0;
4287}
4288
4289static int
4290load_binstring(UnpicklerObject *self)
4291{
4292 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004293 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004294 char *s;
4295
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004296 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004297 return -1;
4298
4299 x = calc_binint(s, 4);
4300 if (x < 0) {
Victor Stinner121aab42011-09-29 23:40:53 +02004301 PyErr_SetString(UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004302 "BINSTRING pickle has negative byte count");
4303 return -1;
4304 }
4305
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004306 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004307 return -1;
4308
4309 /* Convert Python 2.x strings to unicode. */
4310 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4311 if (str == NULL)
4312 return -1;
4313
4314 PDATA_PUSH(self->stack, str, -1);
4315 return 0;
4316}
4317
4318static int
4319load_short_binstring(UnpicklerObject *self)
4320{
4321 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004322 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004323 char *s;
4324
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004325 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004326 return -1;
4327
4328 x = (unsigned char)s[0];
4329
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004330 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004331 return -1;
4332
4333 /* Convert Python 2.x strings to unicode. */
4334 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4335 if (str == NULL)
4336 return -1;
4337
4338 PDATA_PUSH(self->stack, str, -1);
4339 return 0;
4340}
4341
4342static int
4343load_unicode(UnpicklerObject *self)
4344{
4345 PyObject *str;
4346 Py_ssize_t len;
4347 char *s;
4348
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004349 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004350 return -1;
4351 if (len < 1)
4352 return bad_readline();
4353
4354 str = PyUnicode_DecodeRawUnicodeEscape(s, len - 1, NULL);
4355 if (str == NULL)
4356 return -1;
4357
4358 PDATA_PUSH(self->stack, str, -1);
4359 return 0;
4360}
4361
4362static int
4363load_binunicode(UnpicklerObject *self)
4364{
4365 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004366 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004367 char *s;
4368
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004369 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004370 return -1;
4371
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004372 size = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004373 if (size < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004374 PyErr_Format(PyExc_OverflowError,
4375 "BINUNICODE exceeds system's maximum size of %zd bytes",
Alexandre Vassalotticc757172013-04-14 02:25:10 -07004376 PY_SSIZE_T_MAX);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004377 return -1;
4378 }
4379
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004380
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004381 if (_Unpickler_Read(self, &s, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004382 return -1;
4383
Victor Stinner485fb562010-04-13 11:07:24 +00004384 str = PyUnicode_DecodeUTF8(s, size, "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004385 if (str == NULL)
4386 return -1;
4387
4388 PDATA_PUSH(self->stack, str, -1);
4389 return 0;
4390}
4391
4392static int
4393load_tuple(UnpicklerObject *self)
4394{
4395 PyObject *tuple;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004396 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004397
4398 if ((i = marker(self)) < 0)
4399 return -1;
4400
4401 tuple = Pdata_poptuple(self->stack, i);
4402 if (tuple == NULL)
4403 return -1;
4404 PDATA_PUSH(self->stack, tuple, -1);
4405 return 0;
4406}
4407
4408static int
4409load_counted_tuple(UnpicklerObject *self, int len)
4410{
4411 PyObject *tuple;
4412
4413 tuple = PyTuple_New(len);
4414 if (tuple == NULL)
4415 return -1;
4416
4417 while (--len >= 0) {
4418 PyObject *item;
4419
4420 PDATA_POP(self->stack, item);
4421 if (item == NULL)
4422 return -1;
4423 PyTuple_SET_ITEM(tuple, len, item);
4424 }
4425 PDATA_PUSH(self->stack, tuple, -1);
4426 return 0;
4427}
4428
4429static int
4430load_empty_list(UnpicklerObject *self)
4431{
4432 PyObject *list;
4433
4434 if ((list = PyList_New(0)) == NULL)
4435 return -1;
4436 PDATA_PUSH(self->stack, list, -1);
4437 return 0;
4438}
4439
4440static int
4441load_empty_dict(UnpicklerObject *self)
4442{
4443 PyObject *dict;
4444
4445 if ((dict = PyDict_New()) == NULL)
4446 return -1;
4447 PDATA_PUSH(self->stack, dict, -1);
4448 return 0;
4449}
4450
4451static int
4452load_list(UnpicklerObject *self)
4453{
4454 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004455 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004456
4457 if ((i = marker(self)) < 0)
4458 return -1;
4459
4460 list = Pdata_poplist(self->stack, i);
4461 if (list == NULL)
4462 return -1;
4463 PDATA_PUSH(self->stack, list, -1);
4464 return 0;
4465}
4466
4467static int
4468load_dict(UnpicklerObject *self)
4469{
4470 PyObject *dict, *key, *value;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004471 Py_ssize_t i, j, k;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004472
4473 if ((i = marker(self)) < 0)
4474 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004475 j = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004476
4477 if ((dict = PyDict_New()) == NULL)
4478 return -1;
4479
4480 for (k = i + 1; k < j; k += 2) {
4481 key = self->stack->data[k - 1];
4482 value = self->stack->data[k];
4483 if (PyDict_SetItem(dict, key, value) < 0) {
4484 Py_DECREF(dict);
4485 return -1;
4486 }
4487 }
4488 Pdata_clear(self->stack, i);
4489 PDATA_PUSH(self->stack, dict, -1);
4490 return 0;
4491}
4492
4493static PyObject *
4494instantiate(PyObject *cls, PyObject *args)
4495{
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004496 PyObject *result = NULL;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02004497 _Py_IDENTIFIER(__getinitargs__);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004498 /* Caller must assure args are a tuple. Normally, args come from
4499 Pdata_poptuple which packs objects from the top of the stack
4500 into a newly created tuple. */
4501 assert(PyTuple_Check(args));
4502 if (Py_SIZE(args) > 0 || !PyType_Check(cls) ||
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02004503 _PyObject_HasAttrId(cls, &PyId___getinitargs__)) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004504 result = PyObject_CallObject(cls, args);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004505 }
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004506 else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02004507 _Py_IDENTIFIER(__new__);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02004508
4509 result = _PyObject_CallMethodId(cls, &PyId___new__, "O", cls);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004510 }
4511 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004512}
4513
4514static int
4515load_obj(UnpicklerObject *self)
4516{
4517 PyObject *cls, *args, *obj = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004518 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004519
4520 if ((i = marker(self)) < 0)
4521 return -1;
4522
4523 args = Pdata_poptuple(self->stack, i + 1);
4524 if (args == NULL)
4525 return -1;
4526
4527 PDATA_POP(self->stack, cls);
4528 if (cls) {
4529 obj = instantiate(cls, args);
4530 Py_DECREF(cls);
4531 }
4532 Py_DECREF(args);
4533 if (obj == NULL)
4534 return -1;
4535
4536 PDATA_PUSH(self->stack, obj, -1);
4537 return 0;
4538}
4539
4540static int
4541load_inst(UnpicklerObject *self)
4542{
4543 PyObject *cls = NULL;
4544 PyObject *args = NULL;
4545 PyObject *obj = NULL;
4546 PyObject *module_name;
4547 PyObject *class_name;
4548 Py_ssize_t len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004549 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004550 char *s;
4551
4552 if ((i = marker(self)) < 0)
4553 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004554 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004555 return -1;
4556 if (len < 2)
4557 return bad_readline();
4558
4559 /* Here it is safe to use PyUnicode_DecodeASCII(), even though non-ASCII
4560 identifiers are permitted in Python 3.0, since the INST opcode is only
4561 supported by older protocols on Python 2.x. */
4562 module_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
4563 if (module_name == NULL)
4564 return -1;
4565
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004566 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004567 if (len < 2)
4568 return bad_readline();
4569 class_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004570 if (class_name != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004571 cls = find_class(self, module_name, class_name);
4572 Py_DECREF(class_name);
4573 }
4574 }
4575 Py_DECREF(module_name);
4576
4577 if (cls == NULL)
4578 return -1;
4579
4580 if ((args = Pdata_poptuple(self->stack, i)) != NULL) {
4581 obj = instantiate(cls, args);
4582 Py_DECREF(args);
4583 }
4584 Py_DECREF(cls);
4585
4586 if (obj == NULL)
4587 return -1;
4588
4589 PDATA_PUSH(self->stack, obj, -1);
4590 return 0;
4591}
4592
4593static int
4594load_newobj(UnpicklerObject *self)
4595{
4596 PyObject *args = NULL;
4597 PyObject *clsraw = NULL;
4598 PyTypeObject *cls; /* clsraw cast to its true type */
4599 PyObject *obj;
4600
4601 /* Stack is ... cls argtuple, and we want to call
4602 * cls.__new__(cls, *argtuple).
4603 */
4604 PDATA_POP(self->stack, args);
4605 if (args == NULL)
4606 goto error;
4607 if (!PyTuple_Check(args)) {
4608 PyErr_SetString(UnpicklingError, "NEWOBJ expected an arg " "tuple.");
4609 goto error;
4610 }
4611
4612 PDATA_POP(self->stack, clsraw);
4613 cls = (PyTypeObject *)clsraw;
4614 if (cls == NULL)
4615 goto error;
4616 if (!PyType_Check(cls)) {
4617 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4618 "isn't a type object");
4619 goto error;
4620 }
4621 if (cls->tp_new == NULL) {
4622 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4623 "has NULL tp_new");
4624 goto error;
4625 }
4626
4627 /* Call __new__. */
4628 obj = cls->tp_new(cls, args, NULL);
4629 if (obj == NULL)
4630 goto error;
4631
4632 Py_DECREF(args);
4633 Py_DECREF(clsraw);
4634 PDATA_PUSH(self->stack, obj, -1);
4635 return 0;
4636
4637 error:
4638 Py_XDECREF(args);
4639 Py_XDECREF(clsraw);
4640 return -1;
4641}
4642
4643static int
4644load_global(UnpicklerObject *self)
4645{
4646 PyObject *global = NULL;
4647 PyObject *module_name;
4648 PyObject *global_name;
4649 Py_ssize_t len;
4650 char *s;
4651
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004652 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004653 return -1;
4654 if (len < 2)
4655 return bad_readline();
4656 module_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4657 if (!module_name)
4658 return -1;
4659
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004660 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004661 if (len < 2) {
4662 Py_DECREF(module_name);
4663 return bad_readline();
4664 }
4665 global_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4666 if (global_name) {
4667 global = find_class(self, module_name, global_name);
4668 Py_DECREF(global_name);
4669 }
4670 }
4671 Py_DECREF(module_name);
4672
4673 if (global == NULL)
4674 return -1;
4675 PDATA_PUSH(self->stack, global, -1);
4676 return 0;
4677}
4678
4679static int
4680load_persid(UnpicklerObject *self)
4681{
4682 PyObject *pid;
4683 Py_ssize_t len;
4684 char *s;
4685
4686 if (self->pers_func) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004687 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004688 return -1;
4689 if (len < 2)
4690 return bad_readline();
4691
4692 pid = PyBytes_FromStringAndSize(s, len - 1);
4693 if (pid == NULL)
4694 return -1;
4695
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004696 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004697 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004698 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004699 if (pid == NULL)
4700 return -1;
4701
4702 PDATA_PUSH(self->stack, pid, -1);
4703 return 0;
4704 }
4705 else {
4706 PyErr_SetString(UnpicklingError,
4707 "A load persistent id instruction was encountered,\n"
4708 "but no persistent_load function was specified.");
4709 return -1;
4710 }
4711}
4712
4713static int
4714load_binpersid(UnpicklerObject *self)
4715{
4716 PyObject *pid;
4717
4718 if (self->pers_func) {
4719 PDATA_POP(self->stack, pid);
4720 if (pid == NULL)
4721 return -1;
4722
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004723 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004724 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004725 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004726 if (pid == NULL)
4727 return -1;
4728
4729 PDATA_PUSH(self->stack, pid, -1);
4730 return 0;
4731 }
4732 else {
4733 PyErr_SetString(UnpicklingError,
4734 "A load persistent id instruction was encountered,\n"
4735 "but no persistent_load function was specified.");
4736 return -1;
4737 }
4738}
4739
4740static int
4741load_pop(UnpicklerObject *self)
4742{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004743 Py_ssize_t len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004744
4745 /* Note that we split the (pickle.py) stack into two stacks,
4746 * an object stack and a mark stack. We have to be clever and
4747 * pop the right one. We do this by looking at the top of the
Collin Winter8ca69de2009-05-26 16:53:41 +00004748 * mark stack first, and only signalling a stack underflow if
4749 * the object stack is empty and the mark stack doesn't match
4750 * our expectations.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004751 */
Collin Winter8ca69de2009-05-26 16:53:41 +00004752 if (self->num_marks > 0 && self->marks[self->num_marks - 1] == len) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004753 self->num_marks--;
Antoine Pitrou01a15ea2010-01-07 17:57:31 +00004754 } else if (len > 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004755 len--;
4756 Py_DECREF(self->stack->data[len]);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004757 Py_SIZE(self->stack) = len;
Collin Winter8ca69de2009-05-26 16:53:41 +00004758 } else {
4759 return stack_underflow();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004760 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004761 return 0;
4762}
4763
4764static int
4765load_pop_mark(UnpicklerObject *self)
4766{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004767 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004768
4769 if ((i = marker(self)) < 0)
4770 return -1;
4771
4772 Pdata_clear(self->stack, i);
4773
4774 return 0;
4775}
4776
4777static int
4778load_dup(UnpicklerObject *self)
4779{
4780 PyObject *last;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004781 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004782
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004783 if ((len = Py_SIZE(self->stack)) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004784 return stack_underflow();
4785 last = self->stack->data[len - 1];
4786 PDATA_APPEND(self->stack, last, -1);
4787 return 0;
4788}
4789
4790static int
4791load_get(UnpicklerObject *self)
4792{
4793 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004794 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004795 Py_ssize_t len;
4796 char *s;
4797
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004798 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004799 return -1;
4800 if (len < 2)
4801 return bad_readline();
4802
4803 key = PyLong_FromString(s, NULL, 10);
4804 if (key == NULL)
4805 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004806 idx = PyLong_AsSsize_t(key);
4807 if (idx == -1 && PyErr_Occurred()) {
4808 Py_DECREF(key);
4809 return -1;
4810 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004811
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004812 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004813 if (value == NULL) {
4814 if (!PyErr_Occurred())
4815 PyErr_SetObject(PyExc_KeyError, key);
4816 Py_DECREF(key);
4817 return -1;
4818 }
4819 Py_DECREF(key);
4820
4821 PDATA_APPEND(self->stack, value, -1);
4822 return 0;
4823}
4824
4825static int
4826load_binget(UnpicklerObject *self)
4827{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004828 PyObject *value;
4829 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004830 char *s;
4831
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004832 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004833 return -1;
4834
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004835 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004836
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004837 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004838 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004839 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004840 if (!PyErr_Occurred())
4841 PyErr_SetObject(PyExc_KeyError, key);
4842 Py_DECREF(key);
4843 return -1;
4844 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004845
4846 PDATA_APPEND(self->stack, value, -1);
4847 return 0;
4848}
4849
4850static int
4851load_long_binget(UnpicklerObject *self)
4852{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004853 PyObject *value;
4854 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004855 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004856
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004857 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004858 return -1;
4859
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004860 idx = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004861
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004862 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004863 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004864 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004865 if (!PyErr_Occurred())
4866 PyErr_SetObject(PyExc_KeyError, key);
4867 Py_DECREF(key);
4868 return -1;
4869 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004870
4871 PDATA_APPEND(self->stack, value, -1);
4872 return 0;
4873}
4874
4875/* Push an object from the extension registry (EXT[124]). nbytes is
4876 * the number of bytes following the opcode, holding the index (code) value.
4877 */
4878static int
4879load_extension(UnpicklerObject *self, int nbytes)
4880{
4881 char *codebytes; /* the nbytes bytes after the opcode */
4882 long code; /* calc_binint returns long */
4883 PyObject *py_code; /* code as a Python int */
4884 PyObject *obj; /* the object to push */
4885 PyObject *pair; /* (module_name, class_name) */
4886 PyObject *module_name, *class_name;
4887
4888 assert(nbytes == 1 || nbytes == 2 || nbytes == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004889 if (_Unpickler_Read(self, &codebytes, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004890 return -1;
4891 code = calc_binint(codebytes, nbytes);
4892 if (code <= 0) { /* note that 0 is forbidden */
4893 /* Corrupt or hostile pickle. */
4894 PyErr_SetString(UnpicklingError, "EXT specifies code <= 0");
4895 return -1;
4896 }
4897
4898 /* Look for the code in the cache. */
4899 py_code = PyLong_FromLong(code);
4900 if (py_code == NULL)
4901 return -1;
4902 obj = PyDict_GetItem(extension_cache, py_code);
4903 if (obj != NULL) {
4904 /* Bingo. */
4905 Py_DECREF(py_code);
4906 PDATA_APPEND(self->stack, obj, -1);
4907 return 0;
4908 }
4909
4910 /* Look up the (module_name, class_name) pair. */
4911 pair = PyDict_GetItem(inverted_registry, py_code);
4912 if (pair == NULL) {
4913 Py_DECREF(py_code);
4914 PyErr_Format(PyExc_ValueError, "unregistered extension "
4915 "code %ld", code);
4916 return -1;
4917 }
4918 /* Since the extension registry is manipulable via Python code,
4919 * confirm that pair is really a 2-tuple of strings.
4920 */
4921 if (!PyTuple_Check(pair) || PyTuple_Size(pair) != 2 ||
4922 !PyUnicode_Check(module_name = PyTuple_GET_ITEM(pair, 0)) ||
4923 !PyUnicode_Check(class_name = PyTuple_GET_ITEM(pair, 1))) {
4924 Py_DECREF(py_code);
4925 PyErr_Format(PyExc_ValueError, "_inverted_registry[%ld] "
4926 "isn't a 2-tuple of strings", code);
4927 return -1;
4928 }
4929 /* Load the object. */
4930 obj = find_class(self, module_name, class_name);
4931 if (obj == NULL) {
4932 Py_DECREF(py_code);
4933 return -1;
4934 }
4935 /* Cache code -> obj. */
4936 code = PyDict_SetItem(extension_cache, py_code, obj);
4937 Py_DECREF(py_code);
4938 if (code < 0) {
4939 Py_DECREF(obj);
4940 return -1;
4941 }
4942 PDATA_PUSH(self->stack, obj, -1);
4943 return 0;
4944}
4945
4946static int
4947load_put(UnpicklerObject *self)
4948{
4949 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004950 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004951 Py_ssize_t len;
4952 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004953
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004954 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004955 return -1;
4956 if (len < 2)
4957 return bad_readline();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004958 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004959 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004960 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004961
4962 key = PyLong_FromString(s, NULL, 10);
4963 if (key == NULL)
4964 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004965 idx = PyLong_AsSsize_t(key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004966 Py_DECREF(key);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02004967 if (idx < 0) {
4968 if (!PyErr_Occurred())
4969 PyErr_SetString(PyExc_ValueError,
4970 "negative PUT argument");
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004971 return -1;
Antoine Pitrou55549ec2011-08-30 00:27:10 +02004972 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004973
4974 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004975}
4976
4977static int
4978load_binput(UnpicklerObject *self)
4979{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004980 PyObject *value;
4981 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004982 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004983
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004984 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004985 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004986
4987 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004988 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004989 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004990
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004991 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004992
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004993 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004994}
4995
4996static int
4997load_long_binput(UnpicklerObject *self)
4998{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004999 PyObject *value;
5000 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005001 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005002
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005003 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005004 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005005
5006 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005007 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005008 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005009
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005010 idx = calc_binsize(s, 4);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02005011 if (idx < 0) {
5012 PyErr_SetString(PyExc_ValueError,
5013 "negative LONG_BINPUT argument");
5014 return -1;
5015 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005016
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005017 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005018}
5019
5020static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005021do_append(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005022{
5023 PyObject *value;
5024 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005025 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005026
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005027 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005028 if (x > len || x <= 0)
5029 return stack_underflow();
5030 if (len == x) /* nothing to do */
5031 return 0;
5032
5033 list = self->stack->data[x - 1];
5034
5035 if (PyList_Check(list)) {
5036 PyObject *slice;
5037 Py_ssize_t list_len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005038 int ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005039
5040 slice = Pdata_poplist(self->stack, x);
5041 if (!slice)
5042 return -1;
5043 list_len = PyList_GET_SIZE(list);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005044 ret = PyList_SetSlice(list, list_len, list_len, slice);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005045 Py_DECREF(slice);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005046 return ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005047 }
5048 else {
5049 PyObject *append_func;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005050 _Py_IDENTIFIER(append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005051
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005052 append_func = _PyObject_GetAttrId(list, &PyId_append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005053 if (append_func == NULL)
5054 return -1;
5055 for (i = x; i < len; i++) {
5056 PyObject *result;
5057
5058 value = self->stack->data[i];
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005059 result = _Unpickler_FastCall(self, append_func, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005060 if (result == NULL) {
5061 Pdata_clear(self->stack, i + 1);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005062 Py_SIZE(self->stack) = x;
Alexandre Vassalotti637c7c42013-04-20 21:28:21 -07005063 Py_DECREF(append_func);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005064 return -1;
5065 }
5066 Py_DECREF(result);
5067 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005068 Py_SIZE(self->stack) = x;
Alexandre Vassalotti637c7c42013-04-20 21:28:21 -07005069 Py_DECREF(append_func);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005070 }
5071
5072 return 0;
5073}
5074
5075static int
5076load_append(UnpicklerObject *self)
5077{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005078 return do_append(self, Py_SIZE(self->stack) - 1);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005079}
5080
5081static int
5082load_appends(UnpicklerObject *self)
5083{
5084 return do_append(self, marker(self));
5085}
5086
5087static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005088do_setitems(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005089{
5090 PyObject *value, *key;
5091 PyObject *dict;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005092 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005093 int status = 0;
5094
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005095 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005096 if (x > len || x <= 0)
5097 return stack_underflow();
5098 if (len == x) /* nothing to do */
5099 return 0;
Victor Stinner121aab42011-09-29 23:40:53 +02005100 if ((len - x) % 2 != 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005101 /* Currupt or hostile pickle -- we never write one like this. */
5102 PyErr_SetString(UnpicklingError, "odd number of items for SETITEMS");
5103 return -1;
5104 }
5105
5106 /* Here, dict does not actually need to be a PyDict; it could be anything
5107 that supports the __setitem__ attribute. */
5108 dict = self->stack->data[x - 1];
5109
5110 for (i = x + 1; i < len; i += 2) {
5111 key = self->stack->data[i - 1];
5112 value = self->stack->data[i];
5113 if (PyObject_SetItem(dict, key, value) < 0) {
5114 status = -1;
5115 break;
5116 }
5117 }
5118
5119 Pdata_clear(self->stack, x);
5120 return status;
5121}
5122
5123static int
5124load_setitem(UnpicklerObject *self)
5125{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005126 return do_setitems(self, Py_SIZE(self->stack) - 2);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005127}
5128
5129static int
5130load_setitems(UnpicklerObject *self)
5131{
5132 return do_setitems(self, marker(self));
5133}
5134
5135static int
5136load_build(UnpicklerObject *self)
5137{
5138 PyObject *state, *inst, *slotstate;
5139 PyObject *setstate;
5140 int status = 0;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005141 _Py_IDENTIFIER(__setstate__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005142
5143 /* Stack is ... instance, state. We want to leave instance at
5144 * the stack top, possibly mutated via instance.__setstate__(state).
5145 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005146 if (Py_SIZE(self->stack) < 2)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005147 return stack_underflow();
5148
5149 PDATA_POP(self->stack, state);
5150 if (state == NULL)
5151 return -1;
5152
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005153 inst = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005154
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005155 setstate = _PyObject_GetAttrId(inst, &PyId___setstate__);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005156 if (setstate == NULL) {
5157 if (PyErr_ExceptionMatches(PyExc_AttributeError))
5158 PyErr_Clear();
Antoine Pitroud79dc622008-09-05 00:03:33 +00005159 else {
5160 Py_DECREF(state);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005161 return -1;
Antoine Pitroud79dc622008-09-05 00:03:33 +00005162 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005163 }
5164 else {
5165 PyObject *result;
5166
5167 /* The explicit __setstate__ is responsible for everything. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005168 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Antoine Pitroud79dc622008-09-05 00:03:33 +00005169 reference to state first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005170 result = _Unpickler_FastCall(self, setstate, state);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005171 Py_DECREF(setstate);
5172 if (result == NULL)
5173 return -1;
5174 Py_DECREF(result);
5175 return 0;
5176 }
5177
5178 /* A default __setstate__. First see whether state embeds a
5179 * slot state dict too (a proto 2 addition).
5180 */
5181 if (PyTuple_Check(state) && Py_SIZE(state) == 2) {
5182 PyObject *tmp = state;
5183
5184 state = PyTuple_GET_ITEM(tmp, 0);
5185 slotstate = PyTuple_GET_ITEM(tmp, 1);
5186 Py_INCREF(state);
5187 Py_INCREF(slotstate);
5188 Py_DECREF(tmp);
5189 }
5190 else
5191 slotstate = NULL;
5192
5193 /* Set inst.__dict__ from the state dict (if any). */
5194 if (state != Py_None) {
5195 PyObject *dict;
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005196 PyObject *d_key, *d_value;
5197 Py_ssize_t i;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005198 _Py_IDENTIFIER(__dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005199
5200 if (!PyDict_Check(state)) {
5201 PyErr_SetString(UnpicklingError, "state is not a dictionary");
5202 goto error;
5203 }
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005204 dict = _PyObject_GetAttrId(inst, &PyId___dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005205 if (dict == NULL)
5206 goto error;
5207
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005208 i = 0;
5209 while (PyDict_Next(state, &i, &d_key, &d_value)) {
5210 /* normally the keys for instance attributes are
5211 interned. we should try to do that here. */
5212 Py_INCREF(d_key);
5213 if (PyUnicode_CheckExact(d_key))
5214 PyUnicode_InternInPlace(&d_key);
5215 if (PyObject_SetItem(dict, d_key, d_value) < 0) {
5216 Py_DECREF(d_key);
5217 goto error;
5218 }
5219 Py_DECREF(d_key);
5220 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005221 Py_DECREF(dict);
5222 }
5223
5224 /* Also set instance attributes from the slotstate dict (if any). */
5225 if (slotstate != NULL) {
5226 PyObject *d_key, *d_value;
5227 Py_ssize_t i;
5228
5229 if (!PyDict_Check(slotstate)) {
5230 PyErr_SetString(UnpicklingError,
5231 "slot state is not a dictionary");
5232 goto error;
5233 }
5234 i = 0;
5235 while (PyDict_Next(slotstate, &i, &d_key, &d_value)) {
5236 if (PyObject_SetAttr(inst, d_key, d_value) < 0)
5237 goto error;
5238 }
5239 }
5240
5241 if (0) {
5242 error:
5243 status = -1;
5244 }
5245
5246 Py_DECREF(state);
5247 Py_XDECREF(slotstate);
5248 return status;
5249}
5250
5251static int
5252load_mark(UnpicklerObject *self)
5253{
5254
5255 /* Note that we split the (pickle.py) stack into two stacks, an
5256 * object stack and a mark stack. Here we push a mark onto the
5257 * mark stack.
5258 */
5259
5260 if ((self->num_marks + 1) >= self->marks_size) {
5261 size_t alloc;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005262 Py_ssize_t *marks;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005263
5264 /* Use the size_t type to check for overflow. */
5265 alloc = ((size_t)self->num_marks << 1) + 20;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005266 if (alloc > (PY_SSIZE_T_MAX / sizeof(Py_ssize_t)) ||
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00005267 alloc <= ((size_t)self->num_marks + 1)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005268 PyErr_NoMemory();
5269 return -1;
5270 }
5271
5272 if (self->marks == NULL)
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005273 marks = (Py_ssize_t *) PyMem_Malloc(alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005274 else
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005275 marks = (Py_ssize_t *) PyMem_Realloc(self->marks,
5276 alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005277 if (marks == NULL) {
5278 PyErr_NoMemory();
5279 return -1;
5280 }
5281 self->marks = marks;
5282 self->marks_size = (Py_ssize_t)alloc;
5283 }
5284
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005285 self->marks[self->num_marks++] = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005286
5287 return 0;
5288}
5289
5290static int
5291load_reduce(UnpicklerObject *self)
5292{
5293 PyObject *callable = NULL;
5294 PyObject *argtup = NULL;
5295 PyObject *obj = NULL;
5296
5297 PDATA_POP(self->stack, argtup);
5298 if (argtup == NULL)
5299 return -1;
5300 PDATA_POP(self->stack, callable);
5301 if (callable) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005302 obj = PyObject_CallObject(callable, argtup);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005303 Py_DECREF(callable);
5304 }
5305 Py_DECREF(argtup);
5306
5307 if (obj == NULL)
5308 return -1;
5309
5310 PDATA_PUSH(self->stack, obj, -1);
5311 return 0;
5312}
5313
5314/* Just raises an error if we don't know the protocol specified. PROTO
5315 * is the first opcode for protocols >= 2.
5316 */
5317static int
5318load_proto(UnpicklerObject *self)
5319{
5320 char *s;
5321 int i;
5322
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005323 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005324 return -1;
5325
5326 i = (unsigned char)s[0];
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005327 if (i <= HIGHEST_PROTOCOL) {
5328 self->proto = i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005329 return 0;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005330 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005331
5332 PyErr_Format(PyExc_ValueError, "unsupported pickle protocol: %d", i);
5333 return -1;
5334}
5335
5336static PyObject *
5337load(UnpicklerObject *self)
5338{
5339 PyObject *err;
5340 PyObject *value = NULL;
5341 char *s;
5342
5343 self->num_marks = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005344 if (Py_SIZE(self->stack))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005345 Pdata_clear(self->stack, 0);
5346
5347 /* Convenient macros for the dispatch while-switch loop just below. */
5348#define OP(opcode, load_func) \
5349 case opcode: if (load_func(self) < 0) break; continue;
5350
5351#define OP_ARG(opcode, load_func, arg) \
5352 case opcode: if (load_func(self, (arg)) < 0) break; continue;
5353
5354 while (1) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005355 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005356 break;
5357
5358 switch ((enum opcode)s[0]) {
5359 OP(NONE, load_none)
5360 OP(BININT, load_binint)
5361 OP(BININT1, load_binint1)
5362 OP(BININT2, load_binint2)
5363 OP(INT, load_int)
5364 OP(LONG, load_long)
5365 OP_ARG(LONG1, load_counted_long, 1)
5366 OP_ARG(LONG4, load_counted_long, 4)
5367 OP(FLOAT, load_float)
5368 OP(BINFLOAT, load_binfloat)
5369 OP(BINBYTES, load_binbytes)
5370 OP(SHORT_BINBYTES, load_short_binbytes)
5371 OP(BINSTRING, load_binstring)
5372 OP(SHORT_BINSTRING, load_short_binstring)
5373 OP(STRING, load_string)
5374 OP(UNICODE, load_unicode)
5375 OP(BINUNICODE, load_binunicode)
5376 OP_ARG(EMPTY_TUPLE, load_counted_tuple, 0)
5377 OP_ARG(TUPLE1, load_counted_tuple, 1)
5378 OP_ARG(TUPLE2, load_counted_tuple, 2)
5379 OP_ARG(TUPLE3, load_counted_tuple, 3)
5380 OP(TUPLE, load_tuple)
5381 OP(EMPTY_LIST, load_empty_list)
5382 OP(LIST, load_list)
5383 OP(EMPTY_DICT, load_empty_dict)
5384 OP(DICT, load_dict)
5385 OP(OBJ, load_obj)
5386 OP(INST, load_inst)
5387 OP(NEWOBJ, load_newobj)
5388 OP(GLOBAL, load_global)
5389 OP(APPEND, load_append)
5390 OP(APPENDS, load_appends)
5391 OP(BUILD, load_build)
5392 OP(DUP, load_dup)
5393 OP(BINGET, load_binget)
5394 OP(LONG_BINGET, load_long_binget)
5395 OP(GET, load_get)
5396 OP(MARK, load_mark)
5397 OP(BINPUT, load_binput)
5398 OP(LONG_BINPUT, load_long_binput)
5399 OP(PUT, load_put)
5400 OP(POP, load_pop)
5401 OP(POP_MARK, load_pop_mark)
5402 OP(SETITEM, load_setitem)
5403 OP(SETITEMS, load_setitems)
5404 OP(PERSID, load_persid)
5405 OP(BINPERSID, load_binpersid)
5406 OP(REDUCE, load_reduce)
5407 OP(PROTO, load_proto)
5408 OP_ARG(EXT1, load_extension, 1)
5409 OP_ARG(EXT2, load_extension, 2)
5410 OP_ARG(EXT4, load_extension, 4)
5411 OP_ARG(NEWTRUE, load_bool, Py_True)
5412 OP_ARG(NEWFALSE, load_bool, Py_False)
5413
5414 case STOP:
5415 break;
5416
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005417 default:
Benjamin Petersonadde86d2011-09-23 13:41:41 -04005418 if (s[0] == '\0')
5419 PyErr_SetNone(PyExc_EOFError);
5420 else
5421 PyErr_Format(UnpicklingError,
5422 "invalid load key, '%c'.", s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005423 return NULL;
5424 }
5425
5426 break; /* and we are done! */
5427 }
5428
Antoine Pitrou04248a82010-10-12 20:51:21 +00005429 if (_Unpickler_SkipConsumed(self) < 0)
5430 return NULL;
5431
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005432 /* XXX: It is not clear what this is actually for. */
5433 if ((err = PyErr_Occurred())) {
5434 if (err == PyExc_EOFError) {
5435 PyErr_SetNone(PyExc_EOFError);
5436 }
5437 return NULL;
5438 }
5439
5440 PDATA_POP(self->stack, value);
5441 return value;
5442}
5443
5444PyDoc_STRVAR(Unpickler_load_doc,
5445"load() -> object. Load a pickle."
5446"\n"
5447"Read a pickled object representation from the open file object given in\n"
5448"the constructor, and return the reconstituted object hierarchy specified\n"
5449"therein.\n");
5450
5451static PyObject *
5452Unpickler_load(UnpicklerObject *self)
5453{
5454 /* Check whether the Unpickler was initialized correctly. This prevents
5455 segfaulting if a subclass overridden __init__ with a function that does
5456 not call Unpickler.__init__(). Here, we simply ensure that self->read
5457 is not NULL. */
5458 if (self->read == NULL) {
Victor Stinner121aab42011-09-29 23:40:53 +02005459 PyErr_Format(UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005460 "Unpickler.__init__() was not called by %s.__init__()",
5461 Py_TYPE(self)->tp_name);
5462 return NULL;
5463 }
5464
5465 return load(self);
5466}
5467
5468/* The name of find_class() is misleading. In newer pickle protocols, this
5469 function is used for loading any global (i.e., functions), not just
5470 classes. The name is kept only for backward compatibility. */
5471
5472PyDoc_STRVAR(Unpickler_find_class_doc,
5473"find_class(module_name, global_name) -> object.\n"
5474"\n"
5475"Return an object from a specified module, importing the module if\n"
5476"necessary. Subclasses may override this method (e.g. to restrict\n"
5477"unpickling of arbitrary classes and functions).\n"
5478"\n"
5479"This method is called whenever a class or a function object is\n"
5480"needed. Both arguments passed are str objects.\n");
5481
5482static PyObject *
5483Unpickler_find_class(UnpicklerObject *self, PyObject *args)
5484{
5485 PyObject *global;
5486 PyObject *modules_dict;
5487 PyObject *module;
5488 PyObject *module_name, *global_name;
5489
5490 if (!PyArg_UnpackTuple(args, "find_class", 2, 2,
5491 &module_name, &global_name))
5492 return NULL;
5493
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005494 /* Try to map the old names used in Python 2.x to the new ones used in
5495 Python 3.x. We do this only with old pickle protocols and when the
5496 user has not disabled the feature. */
5497 if (self->proto < 3 && self->fix_imports) {
5498 PyObject *key;
5499 PyObject *item;
5500
5501 /* Check if the global (i.e., a function or a class) was renamed
5502 or moved to another module. */
5503 key = PyTuple_Pack(2, module_name, global_name);
5504 if (key == NULL)
5505 return NULL;
5506 item = PyDict_GetItemWithError(name_mapping_2to3, key);
5507 Py_DECREF(key);
5508 if (item) {
5509 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
5510 PyErr_Format(PyExc_RuntimeError,
5511 "_compat_pickle.NAME_MAPPING values should be "
5512 "2-tuples, not %.200s", Py_TYPE(item)->tp_name);
5513 return NULL;
5514 }
5515 module_name = PyTuple_GET_ITEM(item, 0);
5516 global_name = PyTuple_GET_ITEM(item, 1);
5517 if (!PyUnicode_Check(module_name) ||
5518 !PyUnicode_Check(global_name)) {
5519 PyErr_Format(PyExc_RuntimeError,
5520 "_compat_pickle.NAME_MAPPING values should be "
5521 "pairs of str, not (%.200s, %.200s)",
5522 Py_TYPE(module_name)->tp_name,
5523 Py_TYPE(global_name)->tp_name);
5524 return NULL;
5525 }
5526 }
5527 else if (PyErr_Occurred()) {
5528 return NULL;
5529 }
5530
5531 /* Check if the module was renamed. */
5532 item = PyDict_GetItemWithError(import_mapping_2to3, module_name);
5533 if (item) {
5534 if (!PyUnicode_Check(item)) {
5535 PyErr_Format(PyExc_RuntimeError,
5536 "_compat_pickle.IMPORT_MAPPING values should be "
5537 "strings, not %.200s", Py_TYPE(item)->tp_name);
5538 return NULL;
5539 }
5540 module_name = item;
5541 }
5542 else if (PyErr_Occurred()) {
5543 return NULL;
5544 }
5545 }
5546
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005547 modules_dict = PySys_GetObject("modules");
5548 if (modules_dict == NULL)
5549 return NULL;
5550
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005551 module = PyDict_GetItemWithError(modules_dict, module_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005552 if (module == NULL) {
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005553 if (PyErr_Occurred())
5554 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005555 module = PyImport_Import(module_name);
5556 if (module == NULL)
5557 return NULL;
5558 global = PyObject_GetAttr(module, global_name);
5559 Py_DECREF(module);
5560 }
Victor Stinner121aab42011-09-29 23:40:53 +02005561 else {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005562 global = PyObject_GetAttr(module, global_name);
5563 }
5564 return global;
5565}
5566
5567static struct PyMethodDef Unpickler_methods[] = {
5568 {"load", (PyCFunction)Unpickler_load, METH_NOARGS,
5569 Unpickler_load_doc},
5570 {"find_class", (PyCFunction)Unpickler_find_class, METH_VARARGS,
5571 Unpickler_find_class_doc},
5572 {NULL, NULL} /* sentinel */
5573};
5574
5575static void
5576Unpickler_dealloc(UnpicklerObject *self)
5577{
5578 PyObject_GC_UnTrack((PyObject *)self);
5579 Py_XDECREF(self->readline);
5580 Py_XDECREF(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005581 Py_XDECREF(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005582 Py_XDECREF(self->stack);
5583 Py_XDECREF(self->pers_func);
5584 Py_XDECREF(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005585 if (self->buffer.buf != NULL) {
5586 PyBuffer_Release(&self->buffer);
5587 self->buffer.buf = NULL;
5588 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005589
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005590 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005591 PyMem_Free(self->marks);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005592 PyMem_Free(self->input_line);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005593 free(self->encoding);
5594 free(self->errors);
5595
5596 Py_TYPE(self)->tp_free((PyObject *)self);
5597}
5598
5599static int
5600Unpickler_traverse(UnpicklerObject *self, visitproc visit, void *arg)
5601{
5602 Py_VISIT(self->readline);
5603 Py_VISIT(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005604 Py_VISIT(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005605 Py_VISIT(self->stack);
5606 Py_VISIT(self->pers_func);
5607 Py_VISIT(self->arg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005608 return 0;
5609}
5610
5611static int
5612Unpickler_clear(UnpicklerObject *self)
5613{
5614 Py_CLEAR(self->readline);
5615 Py_CLEAR(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005616 Py_CLEAR(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005617 Py_CLEAR(self->stack);
5618 Py_CLEAR(self->pers_func);
5619 Py_CLEAR(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005620 if (self->buffer.buf != NULL) {
5621 PyBuffer_Release(&self->buffer);
5622 self->buffer.buf = NULL;
5623 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005624
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005625 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005626 PyMem_Free(self->marks);
5627 self->marks = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005628 PyMem_Free(self->input_line);
5629 self->input_line = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005630 free(self->encoding);
5631 self->encoding = NULL;
5632 free(self->errors);
5633 self->errors = NULL;
5634
5635 return 0;
5636}
5637
5638PyDoc_STRVAR(Unpickler_doc,
5639"Unpickler(file, *, encoding='ASCII', errors='strict')"
5640"\n"
5641"This takes a binary file for reading a pickle data stream.\n"
5642"\n"
5643"The protocol version of the pickle is detected automatically, so no\n"
5644"proto argument is needed.\n"
5645"\n"
5646"The file-like object must have two methods, a read() method\n"
5647"that takes an integer argument, and a readline() method that\n"
5648"requires no arguments. Both methods should return bytes.\n"
5649"Thus file-like object can be a binary file object opened for\n"
5650"reading, a BytesIO object, or any other custom object that\n"
5651"meets this interface.\n"
5652"\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005653"Optional keyword arguments are *fix_imports*, *encoding* and *errors*,\n"
5654"which are used to control compatiblity support for pickle stream\n"
5655"generated by Python 2.x. If *fix_imports* is True, pickle will try to\n"
5656"map the old Python 2.x names to the new names used in Python 3.x. The\n"
5657"*encoding* and *errors* tell pickle how to decode 8-bit string\n"
5658"instances pickled by Python 2.x; these default to 'ASCII' and\n"
5659"'strict', respectively.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005660
5661static int
5662Unpickler_init(UnpicklerObject *self, PyObject *args, PyObject *kwds)
5663{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005664 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005665 PyObject *file;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005666 PyObject *fix_imports = Py_True;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005667 char *encoding = NULL;
5668 char *errors = NULL;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005669 _Py_IDENTIFIER(persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005670
5671 /* XXX: That is an horrible error message. But, I don't know how to do
5672 better... */
5673 if (Py_SIZE(args) != 1) {
5674 PyErr_Format(PyExc_TypeError,
5675 "%s takes exactly one positional argument (%zd given)",
5676 Py_TYPE(self)->tp_name, Py_SIZE(args));
5677 return -1;
5678 }
5679
5680 /* Arguments parsing needs to be done in the __init__() method to allow
5681 subclasses to define their own __init__() method, which may (or may
5682 not) support Unpickler arguments. However, this means we need to be
5683 extra careful in the other Unpickler methods, since a subclass could
5684 forget to call Unpickler.__init__() thus breaking our internal
5685 invariants. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005686 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:Unpickler", kwlist,
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005687 &file, &fix_imports, &encoding, &errors))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005688 return -1;
5689
5690 /* In case of multiple __init__() calls, clear previous content. */
5691 if (self->read != NULL)
5692 (void)Unpickler_clear(self);
5693
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005694 if (_Unpickler_SetInputStream(self, file) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005695 return -1;
5696
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005697 if (_Unpickler_SetInputEncoding(self, encoding, errors) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005698 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005699
5700 self->fix_imports = PyObject_IsTrue(fix_imports);
5701 if (self->fix_imports == -1)
5702 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005703
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005704 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_load)) {
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005705 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
5706 &PyId_persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005707 if (self->pers_func == NULL)
5708 return -1;
5709 }
5710 else {
5711 self->pers_func = NULL;
5712 }
5713
5714 self->stack = (Pdata *)Pdata_New();
5715 if (self->stack == NULL)
5716 return -1;
5717
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005718 self->memo_size = 32;
5719 self->memo = _Unpickler_NewMemo(self->memo_size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005720 if (self->memo == NULL)
5721 return -1;
5722
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005723 self->arg = NULL;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005724 self->proto = 0;
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005725
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005726 return 0;
5727}
5728
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005729/* Define a proxy object for the Unpickler's internal memo object. This is to
5730 * avoid breaking code like:
5731 * unpickler.memo.clear()
5732 * and
5733 * unpickler.memo = saved_memo
5734 * Is this a good idea? Not really, but we don't want to break code that uses
5735 * it. Note that we don't implement the entire mapping API here. This is
5736 * intentional, as these should be treated as black-box implementation details.
5737 *
5738 * We do, however, have to implement pickling/unpickling support because of
Victor Stinner121aab42011-09-29 23:40:53 +02005739 * real-world code like cvs2svn.
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005740 */
5741
5742typedef struct {
5743 PyObject_HEAD
5744 UnpicklerObject *unpickler;
5745} UnpicklerMemoProxyObject;
5746
5747PyDoc_STRVAR(ump_clear_doc,
5748"memo.clear() -> None. Remove all items from memo.");
5749
5750static PyObject *
5751ump_clear(UnpicklerMemoProxyObject *self)
5752{
5753 _Unpickler_MemoCleanup(self->unpickler);
5754 self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size);
5755 if (self->unpickler->memo == NULL)
5756 return NULL;
5757 Py_RETURN_NONE;
5758}
5759
5760PyDoc_STRVAR(ump_copy_doc,
5761"memo.copy() -> new_memo. Copy the memo to a new object.");
5762
5763static PyObject *
5764ump_copy(UnpicklerMemoProxyObject *self)
5765{
5766 Py_ssize_t i;
5767 PyObject *new_memo = PyDict_New();
5768 if (new_memo == NULL)
5769 return NULL;
5770
5771 for (i = 0; i < self->unpickler->memo_size; i++) {
5772 int status;
5773 PyObject *key, *value;
5774
5775 value = self->unpickler->memo[i];
5776 if (value == NULL)
5777 continue;
5778
5779 key = PyLong_FromSsize_t(i);
5780 if (key == NULL)
5781 goto error;
5782 status = PyDict_SetItem(new_memo, key, value);
5783 Py_DECREF(key);
5784 if (status < 0)
5785 goto error;
5786 }
5787 return new_memo;
5788
5789error:
5790 Py_DECREF(new_memo);
5791 return NULL;
5792}
5793
5794PyDoc_STRVAR(ump_reduce_doc,
5795"memo.__reduce__(). Pickling support.");
5796
5797static PyObject *
5798ump_reduce(UnpicklerMemoProxyObject *self, PyObject *args)
5799{
5800 PyObject *reduce_value;
5801 PyObject *constructor_args;
5802 PyObject *contents = ump_copy(self);
5803 if (contents == NULL)
5804 return NULL;
5805
5806 reduce_value = PyTuple_New(2);
5807 if (reduce_value == NULL) {
5808 Py_DECREF(contents);
5809 return NULL;
5810 }
5811 constructor_args = PyTuple_New(1);
5812 if (constructor_args == NULL) {
5813 Py_DECREF(contents);
5814 Py_DECREF(reduce_value);
5815 return NULL;
5816 }
5817 PyTuple_SET_ITEM(constructor_args, 0, contents);
5818 Py_INCREF((PyObject *)&PyDict_Type);
5819 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
5820 PyTuple_SET_ITEM(reduce_value, 1, constructor_args);
5821 return reduce_value;
5822}
5823
5824static PyMethodDef unpicklerproxy_methods[] = {
5825 {"clear", (PyCFunction)ump_clear, METH_NOARGS, ump_clear_doc},
5826 {"copy", (PyCFunction)ump_copy, METH_NOARGS, ump_copy_doc},
5827 {"__reduce__", (PyCFunction)ump_reduce, METH_VARARGS, ump_reduce_doc},
5828 {NULL, NULL} /* sentinel */
5829};
5830
5831static void
5832UnpicklerMemoProxy_dealloc(UnpicklerMemoProxyObject *self)
5833{
5834 PyObject_GC_UnTrack(self);
5835 Py_XDECREF(self->unpickler);
5836 PyObject_GC_Del((PyObject *)self);
5837}
5838
5839static int
5840UnpicklerMemoProxy_traverse(UnpicklerMemoProxyObject *self,
5841 visitproc visit, void *arg)
5842{
5843 Py_VISIT(self->unpickler);
5844 return 0;
5845}
5846
5847static int
5848UnpicklerMemoProxy_clear(UnpicklerMemoProxyObject *self)
5849{
5850 Py_CLEAR(self->unpickler);
5851 return 0;
5852}
5853
5854static PyTypeObject UnpicklerMemoProxyType = {
5855 PyVarObject_HEAD_INIT(NULL, 0)
5856 "_pickle.UnpicklerMemoProxy", /*tp_name*/
5857 sizeof(UnpicklerMemoProxyObject), /*tp_basicsize*/
5858 0,
5859 (destructor)UnpicklerMemoProxy_dealloc, /* tp_dealloc */
5860 0, /* tp_print */
5861 0, /* tp_getattr */
5862 0, /* tp_setattr */
5863 0, /* tp_compare */
5864 0, /* tp_repr */
5865 0, /* tp_as_number */
5866 0, /* tp_as_sequence */
5867 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00005868 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005869 0, /* tp_call */
5870 0, /* tp_str */
5871 PyObject_GenericGetAttr, /* tp_getattro */
5872 PyObject_GenericSetAttr, /* tp_setattro */
5873 0, /* tp_as_buffer */
5874 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5875 0, /* tp_doc */
5876 (traverseproc)UnpicklerMemoProxy_traverse, /* tp_traverse */
5877 (inquiry)UnpicklerMemoProxy_clear, /* tp_clear */
5878 0, /* tp_richcompare */
5879 0, /* tp_weaklistoffset */
5880 0, /* tp_iter */
5881 0, /* tp_iternext */
5882 unpicklerproxy_methods, /* tp_methods */
5883};
5884
5885static PyObject *
5886UnpicklerMemoProxy_New(UnpicklerObject *unpickler)
5887{
5888 UnpicklerMemoProxyObject *self;
5889
5890 self = PyObject_GC_New(UnpicklerMemoProxyObject,
5891 &UnpicklerMemoProxyType);
5892 if (self == NULL)
5893 return NULL;
5894 Py_INCREF(unpickler);
5895 self->unpickler = unpickler;
5896 PyObject_GC_Track(self);
5897 return (PyObject *)self;
5898}
5899
5900/*****************************************************************************/
5901
5902
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005903static PyObject *
5904Unpickler_get_memo(UnpicklerObject *self)
5905{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005906 return UnpicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005907}
5908
5909static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005910Unpickler_set_memo(UnpicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005911{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005912 PyObject **new_memo;
5913 Py_ssize_t new_memo_size = 0;
5914 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005915
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005916 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005917 PyErr_SetString(PyExc_TypeError,
5918 "attribute deletion is not supported");
5919 return -1;
5920 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005921
5922 if (Py_TYPE(obj) == &UnpicklerMemoProxyType) {
5923 UnpicklerObject *unpickler =
5924 ((UnpicklerMemoProxyObject *)obj)->unpickler;
5925
5926 new_memo_size = unpickler->memo_size;
5927 new_memo = _Unpickler_NewMemo(new_memo_size);
5928 if (new_memo == NULL)
5929 return -1;
5930
5931 for (i = 0; i < new_memo_size; i++) {
5932 Py_XINCREF(unpickler->memo[i]);
5933 new_memo[i] = unpickler->memo[i];
5934 }
5935 }
5936 else if (PyDict_Check(obj)) {
5937 Py_ssize_t i = 0;
5938 PyObject *key, *value;
5939
5940 new_memo_size = PyDict_Size(obj);
5941 new_memo = _Unpickler_NewMemo(new_memo_size);
5942 if (new_memo == NULL)
5943 return -1;
5944
5945 while (PyDict_Next(obj, &i, &key, &value)) {
5946 Py_ssize_t idx;
5947 if (!PyLong_Check(key)) {
5948 PyErr_SetString(PyExc_TypeError,
5949 "memo key must be integers");
5950 goto error;
5951 }
5952 idx = PyLong_AsSsize_t(key);
5953 if (idx == -1 && PyErr_Occurred())
5954 goto error;
Christian Heimesa24b4d22013-07-01 15:17:45 +02005955 if (idx < 0) {
5956 PyErr_SetString(PyExc_ValueError,
Christian Heimes80878792013-07-01 15:23:39 +02005957 "memo key must be positive integers.");
Christian Heimesa24b4d22013-07-01 15:17:45 +02005958 goto error;
5959 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005960 if (_Unpickler_MemoPut(self, idx, value) < 0)
5961 goto error;
5962 }
5963 }
5964 else {
5965 PyErr_Format(PyExc_TypeError,
5966 "'memo' attribute must be an UnpicklerMemoProxy object"
5967 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005968 return -1;
5969 }
5970
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005971 _Unpickler_MemoCleanup(self);
5972 self->memo_size = new_memo_size;
5973 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005974
5975 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005976
5977 error:
5978 if (new_memo_size) {
5979 i = new_memo_size;
5980 while (--i >= 0) {
5981 Py_XDECREF(new_memo[i]);
5982 }
5983 PyMem_FREE(new_memo);
5984 }
5985 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005986}
5987
5988static PyObject *
5989Unpickler_get_persload(UnpicklerObject *self)
5990{
5991 if (self->pers_func == NULL)
5992 PyErr_SetString(PyExc_AttributeError, "persistent_load");
5993 else
5994 Py_INCREF(self->pers_func);
5995 return self->pers_func;
5996}
5997
5998static int
5999Unpickler_set_persload(UnpicklerObject *self, PyObject *value)
6000{
6001 PyObject *tmp;
6002
6003 if (value == NULL) {
6004 PyErr_SetString(PyExc_TypeError,
6005 "attribute deletion is not supported");
6006 return -1;
6007 }
6008 if (!PyCallable_Check(value)) {
6009 PyErr_SetString(PyExc_TypeError,
6010 "persistent_load must be a callable taking "
6011 "one argument");
6012 return -1;
6013 }
6014
6015 tmp = self->pers_func;
6016 Py_INCREF(value);
6017 self->pers_func = value;
6018 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
6019
6020 return 0;
6021}
6022
6023static PyGetSetDef Unpickler_getsets[] = {
6024 {"memo", (getter)Unpickler_get_memo, (setter)Unpickler_set_memo},
6025 {"persistent_load", (getter)Unpickler_get_persload,
6026 (setter)Unpickler_set_persload},
6027 {NULL}
6028};
6029
6030static PyTypeObject Unpickler_Type = {
6031 PyVarObject_HEAD_INIT(NULL, 0)
6032 "_pickle.Unpickler", /*tp_name*/
6033 sizeof(UnpicklerObject), /*tp_basicsize*/
6034 0, /*tp_itemsize*/
6035 (destructor)Unpickler_dealloc, /*tp_dealloc*/
6036 0, /*tp_print*/
6037 0, /*tp_getattr*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006038 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00006039 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006040 0, /*tp_repr*/
6041 0, /*tp_as_number*/
6042 0, /*tp_as_sequence*/
6043 0, /*tp_as_mapping*/
6044 0, /*tp_hash*/
6045 0, /*tp_call*/
6046 0, /*tp_str*/
6047 0, /*tp_getattro*/
6048 0, /*tp_setattro*/
6049 0, /*tp_as_buffer*/
6050 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
6051 Unpickler_doc, /*tp_doc*/
6052 (traverseproc)Unpickler_traverse, /*tp_traverse*/
6053 (inquiry)Unpickler_clear, /*tp_clear*/
6054 0, /*tp_richcompare*/
6055 0, /*tp_weaklistoffset*/
6056 0, /*tp_iter*/
6057 0, /*tp_iternext*/
6058 Unpickler_methods, /*tp_methods*/
6059 0, /*tp_members*/
6060 Unpickler_getsets, /*tp_getset*/
6061 0, /*tp_base*/
6062 0, /*tp_dict*/
6063 0, /*tp_descr_get*/
6064 0, /*tp_descr_set*/
6065 0, /*tp_dictoffset*/
6066 (initproc)Unpickler_init, /*tp_init*/
6067 PyType_GenericAlloc, /*tp_alloc*/
6068 PyType_GenericNew, /*tp_new*/
6069 PyObject_GC_Del, /*tp_free*/
6070 0, /*tp_is_gc*/
6071};
6072
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006073PyDoc_STRVAR(pickle_dump_doc,
6074"dump(obj, file, protocol=None, *, fix_imports=True) -> None\n"
6075"\n"
6076"Write a pickled representation of obj to the open file object file. This\n"
6077"is equivalent to ``Pickler(file, protocol).dump(obj)``, but may be more\n"
6078"efficient.\n"
6079"\n"
6080"The optional protocol argument tells the pickler to use the given protocol;\n"
6081"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
6082"backward-incompatible protocol designed for Python 3.0.\n"
6083"\n"
6084"Specifying a negative protocol version selects the highest protocol version\n"
6085"supported. The higher the protocol used, the more recent the version of\n"
6086"Python needed to read the pickle produced.\n"
6087"\n"
6088"The file argument must have a write() method that accepts a single bytes\n"
6089"argument. It can thus be a file object opened for binary writing, a\n"
6090"io.BytesIO instance, or any other custom object that meets this interface.\n"
6091"\n"
6092"If fix_imports is True and protocol is less than 3, pickle will try to\n"
6093"map the new Python 3.x names to the old module names used in Python 2.x,\n"
6094"so that the pickle data stream is readable with Python 2.x.\n");
6095
6096static PyObject *
6097pickle_dump(PyObject *self, PyObject *args, PyObject *kwds)
6098{
6099 static char *kwlist[] = {"obj", "file", "protocol", "fix_imports", 0};
6100 PyObject *obj;
6101 PyObject *file;
6102 PyObject *proto = NULL;
6103 PyObject *fix_imports = Py_True;
6104 PicklerObject *pickler;
6105
6106 /* fix_imports is a keyword-only argument. */
6107 if (Py_SIZE(args) > 3) {
6108 PyErr_Format(PyExc_TypeError,
6109 "pickle.dump() takes at most 3 positional "
6110 "argument (%zd given)", Py_SIZE(args));
6111 return NULL;
6112 }
6113
6114 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|OO:dump", kwlist,
6115 &obj, &file, &proto, &fix_imports))
6116 return NULL;
6117
6118 pickler = _Pickler_New();
6119 if (pickler == NULL)
6120 return NULL;
6121
6122 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6123 goto error;
6124
6125 if (_Pickler_SetOutputStream(pickler, file) < 0)
6126 goto error;
6127
6128 if (dump(pickler, obj) < 0)
6129 goto error;
6130
6131 if (_Pickler_FlushToFile(pickler) < 0)
6132 goto error;
6133
6134 Py_DECREF(pickler);
6135 Py_RETURN_NONE;
6136
6137 error:
6138 Py_XDECREF(pickler);
6139 return NULL;
6140}
6141
6142PyDoc_STRVAR(pickle_dumps_doc,
6143"dumps(obj, protocol=None, *, fix_imports=True) -> bytes\n"
6144"\n"
6145"Return the pickled representation of the object as a bytes\n"
6146"object, instead of writing it to a file.\n"
6147"\n"
6148"The optional protocol argument tells the pickler to use the given protocol;\n"
6149"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
6150"backward-incompatible protocol designed for Python 3.0.\n"
6151"\n"
6152"Specifying a negative protocol version selects the highest protocol version\n"
6153"supported. The higher the protocol used, the more recent the version of\n"
6154"Python needed to read the pickle produced.\n"
6155"\n"
6156"If fix_imports is True and *protocol* is less than 3, pickle will try to\n"
6157"map the new Python 3.x names to the old module names used in Python 2.x,\n"
6158"so that the pickle data stream is readable with Python 2.x.\n");
6159
6160static PyObject *
6161pickle_dumps(PyObject *self, PyObject *args, PyObject *kwds)
6162{
6163 static char *kwlist[] = {"obj", "protocol", "fix_imports", 0};
6164 PyObject *obj;
6165 PyObject *proto = NULL;
6166 PyObject *result;
6167 PyObject *fix_imports = Py_True;
6168 PicklerObject *pickler;
6169
6170 /* fix_imports is a keyword-only argument. */
6171 if (Py_SIZE(args) > 2) {
6172 PyErr_Format(PyExc_TypeError,
6173 "pickle.dumps() takes at most 2 positional "
6174 "argument (%zd given)", Py_SIZE(args));
6175 return NULL;
6176 }
6177
6178 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:dumps", kwlist,
6179 &obj, &proto, &fix_imports))
6180 return NULL;
6181
6182 pickler = _Pickler_New();
6183 if (pickler == NULL)
6184 return NULL;
6185
6186 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6187 goto error;
6188
6189 if (dump(pickler, obj) < 0)
6190 goto error;
6191
6192 result = _Pickler_GetString(pickler);
6193 Py_DECREF(pickler);
6194 return result;
6195
6196 error:
6197 Py_XDECREF(pickler);
6198 return NULL;
6199}
6200
6201PyDoc_STRVAR(pickle_load_doc,
6202"load(file, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6203"\n"
6204"Read a pickled object representation from the open file object file and\n"
6205"return the reconstituted object hierarchy specified therein. This is\n"
6206"equivalent to ``Unpickler(file).load()``, but may be more efficient.\n"
6207"\n"
6208"The protocol version of the pickle is detected automatically, so no protocol\n"
6209"argument is needed. Bytes past the pickled object's representation are\n"
6210"ignored.\n"
6211"\n"
6212"The argument file must have two methods, a read() method that takes an\n"
6213"integer argument, and a readline() method that requires no arguments. Both\n"
6214"methods should return bytes. Thus *file* can be a binary file object opened\n"
6215"for reading, a BytesIO object, or any other custom object that meets this\n"
6216"interface.\n"
6217"\n"
6218"Optional keyword arguments are fix_imports, encoding and errors,\n"
6219"which are used to control compatiblity support for pickle stream generated\n"
6220"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6221"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6222"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6223"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6224
6225static PyObject *
6226pickle_load(PyObject *self, PyObject *args, PyObject *kwds)
6227{
6228 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
6229 PyObject *file;
6230 PyObject *fix_imports = Py_True;
6231 PyObject *result;
6232 char *encoding = NULL;
6233 char *errors = NULL;
6234 UnpicklerObject *unpickler;
6235
6236 /* fix_imports, encoding and errors are a keyword-only argument. */
6237 if (Py_SIZE(args) != 1) {
6238 PyErr_Format(PyExc_TypeError,
6239 "pickle.load() takes exactly one positional "
6240 "argument (%zd given)", Py_SIZE(args));
6241 return NULL;
6242 }
6243
6244 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:load", kwlist,
6245 &file, &fix_imports, &encoding, &errors))
6246 return NULL;
6247
6248 unpickler = _Unpickler_New();
6249 if (unpickler == NULL)
6250 return NULL;
6251
6252 if (_Unpickler_SetInputStream(unpickler, file) < 0)
6253 goto error;
6254
6255 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6256 goto error;
6257
6258 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6259 if (unpickler->fix_imports == -1)
6260 goto error;
6261
6262 result = load(unpickler);
6263 Py_DECREF(unpickler);
6264 return result;
6265
6266 error:
6267 Py_XDECREF(unpickler);
6268 return NULL;
6269}
6270
6271PyDoc_STRVAR(pickle_loads_doc,
6272"loads(input, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6273"\n"
6274"Read a pickled object hierarchy from a bytes object and return the\n"
6275"reconstituted object hierarchy specified therein\n"
6276"\n"
6277"The protocol version of the pickle is detected automatically, so no protocol\n"
6278"argument is needed. Bytes past the pickled object's representation are\n"
6279"ignored.\n"
6280"\n"
6281"Optional keyword arguments are fix_imports, encoding and errors, which\n"
6282"are used to control compatiblity support for pickle stream generated\n"
6283"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6284"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6285"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6286"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6287
6288static PyObject *
6289pickle_loads(PyObject *self, PyObject *args, PyObject *kwds)
6290{
6291 static char *kwlist[] = {"input", "fix_imports", "encoding", "errors", 0};
6292 PyObject *input;
6293 PyObject *fix_imports = Py_True;
6294 PyObject *result;
6295 char *encoding = NULL;
6296 char *errors = NULL;
6297 UnpicklerObject *unpickler;
6298
6299 /* fix_imports, encoding and errors are a keyword-only argument. */
6300 if (Py_SIZE(args) != 1) {
6301 PyErr_Format(PyExc_TypeError,
6302 "pickle.loads() takes exactly one positional "
6303 "argument (%zd given)", Py_SIZE(args));
6304 return NULL;
6305 }
6306
6307 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:loads", kwlist,
6308 &input, &fix_imports, &encoding, &errors))
6309 return NULL;
6310
6311 unpickler = _Unpickler_New();
6312 if (unpickler == NULL)
6313 return NULL;
6314
6315 if (_Unpickler_SetStringInput(unpickler, input) < 0)
6316 goto error;
6317
6318 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6319 goto error;
6320
6321 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6322 if (unpickler->fix_imports == -1)
6323 goto error;
6324
6325 result = load(unpickler);
6326 Py_DECREF(unpickler);
6327 return result;
6328
6329 error:
6330 Py_XDECREF(unpickler);
6331 return NULL;
6332}
6333
6334
6335static struct PyMethodDef pickle_methods[] = {
6336 {"dump", (PyCFunction)pickle_dump, METH_VARARGS|METH_KEYWORDS,
6337 pickle_dump_doc},
6338 {"dumps", (PyCFunction)pickle_dumps, METH_VARARGS|METH_KEYWORDS,
6339 pickle_dumps_doc},
6340 {"load", (PyCFunction)pickle_load, METH_VARARGS|METH_KEYWORDS,
6341 pickle_load_doc},
6342 {"loads", (PyCFunction)pickle_loads, METH_VARARGS|METH_KEYWORDS,
6343 pickle_loads_doc},
6344 {NULL, NULL} /* sentinel */
6345};
6346
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006347static int
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006348initmodule(void)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006349{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006350 PyObject *copyreg = NULL;
6351 PyObject *compat_pickle = NULL;
6352
6353 /* XXX: We should ensure that the types of the dictionaries imported are
6354 exactly PyDict objects. Otherwise, it is possible to crash the pickle
6355 since we use the PyDict API directly to access these dictionaries. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006356
6357 copyreg = PyImport_ImportModule("copyreg");
6358 if (!copyreg)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006359 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006360 dispatch_table = PyObject_GetAttrString(copyreg, "dispatch_table");
6361 if (!dispatch_table)
6362 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006363 extension_registry = \
6364 PyObject_GetAttrString(copyreg, "_extension_registry");
6365 if (!extension_registry)
6366 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006367 inverted_registry = PyObject_GetAttrString(copyreg, "_inverted_registry");
6368 if (!inverted_registry)
6369 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006370 extension_cache = PyObject_GetAttrString(copyreg, "_extension_cache");
6371 if (!extension_cache)
6372 goto error;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006373 Py_CLEAR(copyreg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006374
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006375 /* Load the 2.x -> 3.x stdlib module mapping tables */
6376 compat_pickle = PyImport_ImportModule("_compat_pickle");
6377 if (!compat_pickle)
6378 goto error;
6379 name_mapping_2to3 = PyObject_GetAttrString(compat_pickle, "NAME_MAPPING");
6380 if (!name_mapping_2to3)
6381 goto error;
6382 if (!PyDict_CheckExact(name_mapping_2to3)) {
6383 PyErr_Format(PyExc_RuntimeError,
6384 "_compat_pickle.NAME_MAPPING should be a dict, not %.200s",
6385 Py_TYPE(name_mapping_2to3)->tp_name);
6386 goto error;
6387 }
6388 import_mapping_2to3 = PyObject_GetAttrString(compat_pickle,
6389 "IMPORT_MAPPING");
6390 if (!import_mapping_2to3)
6391 goto error;
6392 if (!PyDict_CheckExact(import_mapping_2to3)) {
6393 PyErr_Format(PyExc_RuntimeError,
6394 "_compat_pickle.IMPORT_MAPPING should be a dict, "
6395 "not %.200s", Py_TYPE(import_mapping_2to3)->tp_name);
6396 goto error;
6397 }
6398 /* ... and the 3.x -> 2.x mapping tables */
6399 name_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6400 "REVERSE_NAME_MAPPING");
6401 if (!name_mapping_3to2)
6402 goto error;
6403 if (!PyDict_CheckExact(name_mapping_3to2)) {
6404 PyErr_Format(PyExc_RuntimeError,
Ezio Melotti13925002011-03-16 11:05:33 +02006405 "_compat_pickle.REVERSE_NAME_MAPPING should be a dict, "
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006406 "not %.200s", Py_TYPE(name_mapping_3to2)->tp_name);
6407 goto error;
6408 }
6409 import_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6410 "REVERSE_IMPORT_MAPPING");
6411 if (!import_mapping_3to2)
6412 goto error;
6413 if (!PyDict_CheckExact(import_mapping_3to2)) {
6414 PyErr_Format(PyExc_RuntimeError,
6415 "_compat_pickle.REVERSE_IMPORT_MAPPING should be a dict, "
6416 "not %.200s", Py_TYPE(import_mapping_3to2)->tp_name);
6417 goto error;
6418 }
6419 Py_CLEAR(compat_pickle);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006420
6421 empty_tuple = PyTuple_New(0);
6422 if (empty_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006423 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006424 two_tuple = PyTuple_New(2);
6425 if (two_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006426 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006427 /* We use this temp container with no regard to refcounts, or to
6428 * keeping containees alive. Exempt from GC, because we don't
6429 * want anything looking at two_tuple() by magic.
6430 */
6431 PyObject_GC_UnTrack(two_tuple);
6432
6433 return 0;
6434
6435 error:
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006436 Py_CLEAR(copyreg);
6437 Py_CLEAR(dispatch_table);
6438 Py_CLEAR(extension_registry);
6439 Py_CLEAR(inverted_registry);
6440 Py_CLEAR(extension_cache);
6441 Py_CLEAR(compat_pickle);
6442 Py_CLEAR(name_mapping_2to3);
6443 Py_CLEAR(import_mapping_2to3);
6444 Py_CLEAR(name_mapping_3to2);
6445 Py_CLEAR(import_mapping_3to2);
6446 Py_CLEAR(empty_tuple);
6447 Py_CLEAR(two_tuple);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006448 return -1;
6449}
6450
6451static struct PyModuleDef _picklemodule = {
6452 PyModuleDef_HEAD_INIT,
6453 "_pickle",
6454 pickle_module_doc,
6455 -1,
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006456 pickle_methods,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006457 NULL,
6458 NULL,
6459 NULL,
6460 NULL
6461};
6462
6463PyMODINIT_FUNC
6464PyInit__pickle(void)
6465{
6466 PyObject *m;
6467
6468 if (PyType_Ready(&Unpickler_Type) < 0)
6469 return NULL;
6470 if (PyType_Ready(&Pickler_Type) < 0)
6471 return NULL;
6472 if (PyType_Ready(&Pdata_Type) < 0)
6473 return NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006474 if (PyType_Ready(&PicklerMemoProxyType) < 0)
6475 return NULL;
6476 if (PyType_Ready(&UnpicklerMemoProxyType) < 0)
6477 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006478
6479 /* Create the module and add the functions. */
6480 m = PyModule_Create(&_picklemodule);
6481 if (m == NULL)
6482 return NULL;
6483
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006484 Py_INCREF(&Pickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006485 if (PyModule_AddObject(m, "Pickler", (PyObject *)&Pickler_Type) < 0)
6486 return NULL;
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006487 Py_INCREF(&Unpickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006488 if (PyModule_AddObject(m, "Unpickler", (PyObject *)&Unpickler_Type) < 0)
6489 return NULL;
6490
6491 /* Initialize the exceptions. */
6492 PickleError = PyErr_NewException("_pickle.PickleError", NULL, NULL);
6493 if (PickleError == NULL)
6494 return NULL;
6495 PicklingError = \
6496 PyErr_NewException("_pickle.PicklingError", PickleError, NULL);
6497 if (PicklingError == NULL)
6498 return NULL;
6499 UnpicklingError = \
6500 PyErr_NewException("_pickle.UnpicklingError", PickleError, NULL);
6501 if (UnpicklingError == NULL)
6502 return NULL;
6503
6504 if (PyModule_AddObject(m, "PickleError", PickleError) < 0)
6505 return NULL;
6506 if (PyModule_AddObject(m, "PicklingError", PicklingError) < 0)
6507 return NULL;
6508 if (PyModule_AddObject(m, "UnpicklingError", UnpicklingError) < 0)
6509 return NULL;
6510
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006511 if (initmodule() < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006512 return NULL;
6513
6514 return m;
6515}