blob: f0d3e7928ccaa03e73928da9a94194fbca3944c0 [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");
1751 if (latin1 == NULL)
1752 return -1;
1753 }
1754 reduce_value = Py_BuildValue("(O(OO))",
1755 codecs_encode, unicode_str, latin1);
1756 Py_DECREF(unicode_str);
1757 }
1758
1759 if (reduce_value == NULL)
1760 return -1;
1761
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001762 /* save_reduce() will memoize the object automatically. */
1763 status = save_reduce(self, reduce_value, obj);
1764 Py_DECREF(reduce_value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001765 return status;
1766 }
1767 else {
1768 Py_ssize_t size;
1769 char header[5];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001770 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001771
Alexandre Vassalotti3bfc65a2011-12-13 13:08:09 -05001772 size = PyBytes_GET_SIZE(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001773 if (size < 0)
1774 return -1;
1775
1776 if (size < 256) {
1777 header[0] = SHORT_BINBYTES;
1778 header[1] = (unsigned char)size;
1779 len = 2;
1780 }
1781 else if (size <= 0xffffffffL) {
1782 header[0] = BINBYTES;
1783 header[1] = (unsigned char)(size & 0xff);
1784 header[2] = (unsigned char)((size >> 8) & 0xff);
1785 header[3] = (unsigned char)((size >> 16) & 0xff);
1786 header[4] = (unsigned char)((size >> 24) & 0xff);
1787 len = 5;
1788 }
1789 else {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001790 PyErr_SetString(PyExc_OverflowError,
Serhiy Storchakaf8def282013-02-16 17:29:56 +02001791 "cannot serialize a bytes object larger than 4 GiB");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001792 return -1; /* string too large */
1793 }
1794
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001795 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001796 return -1;
1797
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001798 if (_Pickler_Write(self, PyBytes_AS_STRING(obj), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001799 return -1;
1800
1801 if (memo_put(self, obj) < 0)
1802 return -1;
1803
1804 return 0;
1805 }
1806}
1807
1808/* A copy of PyUnicode_EncodeRawUnicodeEscape() that also translates
1809 backslash and newline characters to \uXXXX escapes. */
1810static PyObject *
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001811raw_unicode_escape(PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001812{
1813 PyObject *repr, *result;
1814 char *p;
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001815 Py_ssize_t i, size, expandsize;
1816 void *data;
1817 unsigned int kind;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001818
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001819 if (PyUnicode_READY(obj))
1820 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001821
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001822 size = PyUnicode_GET_LENGTH(obj);
1823 data = PyUnicode_DATA(obj);
1824 kind = PyUnicode_KIND(obj);
1825 if (kind == PyUnicode_4BYTE_KIND)
1826 expandsize = 10;
1827 else
1828 expandsize = 6;
Victor Stinner121aab42011-09-29 23:40:53 +02001829
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001830 if (size > PY_SSIZE_T_MAX / expandsize)
1831 return PyErr_NoMemory();
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001832 repr = PyByteArray_FromStringAndSize(NULL, expandsize * size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001833 if (repr == NULL)
1834 return NULL;
1835 if (size == 0)
1836 goto done;
1837
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001838 p = PyByteArray_AS_STRING(repr);
1839 for (i=0; i < size; i++) {
1840 Py_UCS4 ch = PyUnicode_READ(kind, data, i);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001841 /* Map 32-bit characters to '\Uxxxxxxxx' */
1842 if (ch >= 0x10000) {
1843 *p++ = '\\';
1844 *p++ = 'U';
Victor Stinnerf5cff562011-10-14 02:13:11 +02001845 *p++ = Py_hexdigits[(ch >> 28) & 0xf];
1846 *p++ = Py_hexdigits[(ch >> 24) & 0xf];
1847 *p++ = Py_hexdigits[(ch >> 20) & 0xf];
1848 *p++ = Py_hexdigits[(ch >> 16) & 0xf];
1849 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
1850 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
1851 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
1852 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001853 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001854 /* Map 16-bit characters to '\uxxxx' */
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001855 else if (ch >= 256 || ch == '\\' || ch == '\n') {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001856 *p++ = '\\';
1857 *p++ = 'u';
Victor Stinnerf5cff562011-10-14 02:13:11 +02001858 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
1859 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
1860 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
1861 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001862 }
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001863 /* Copy everything else as-is */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001864 else
1865 *p++ = (char) ch;
1866 }
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001867 size = p - PyByteArray_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001868
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001869done:
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001870 result = PyBytes_FromStringAndSize(PyByteArray_AS_STRING(repr), size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001871 Py_DECREF(repr);
1872 return result;
1873}
1874
1875static int
Antoine Pitrou299978d2013-04-07 17:38:11 +02001876write_utf8(PicklerObject *self, char *data, Py_ssize_t size)
1877{
1878 char pdata[5];
1879
1880#if SIZEOF_SIZE_T > 4
1881 if (size > 0xffffffffUL) {
1882 /* string too large */
1883 PyErr_SetString(PyExc_OverflowError,
1884 "cannot serialize a string larger than 4GB");
1885 return -1;
1886 }
1887#endif
1888
1889 pdata[0] = BINUNICODE;
1890 pdata[1] = (unsigned char)(size & 0xff);
1891 pdata[2] = (unsigned char)((size >> 8) & 0xff);
1892 pdata[3] = (unsigned char)((size >> 16) & 0xff);
1893 pdata[4] = (unsigned char)((size >> 24) & 0xff);
1894
1895 if (_Pickler_Write(self, pdata, sizeof(pdata)) < 0)
1896 return -1;
1897
1898 if (_Pickler_Write(self, data, size) < 0)
1899 return -1;
1900
1901 return 0;
1902}
1903
1904static int
1905write_unicode_binary(PicklerObject *self, PyObject *obj)
1906{
1907 PyObject *encoded = NULL;
1908 Py_ssize_t size;
1909 char *data;
1910 int r;
1911
1912 if (PyUnicode_READY(obj))
1913 return -1;
1914
1915 data = PyUnicode_AsUTF8AndSize(obj, &size);
1916 if (data != NULL)
1917 return write_utf8(self, data, size);
1918
1919 /* Issue #8383: for strings with lone surrogates, fallback on the
1920 "surrogatepass" error handler. */
1921 PyErr_Clear();
1922 encoded = PyUnicode_AsEncodedString(obj, "utf-8", "surrogatepass");
1923 if (encoded == NULL)
1924 return -1;
1925
1926 r = write_utf8(self, PyBytes_AS_STRING(encoded),
1927 PyBytes_GET_SIZE(encoded));
1928 Py_DECREF(encoded);
1929 return r;
1930}
1931
1932static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001933save_unicode(PicklerObject *self, PyObject *obj)
1934{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001935 if (self->bin) {
Antoine Pitrou299978d2013-04-07 17:38:11 +02001936 if (write_unicode_binary(self, obj) < 0)
1937 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001938 }
1939 else {
Antoine Pitrou299978d2013-04-07 17:38:11 +02001940 PyObject *encoded;
1941 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001942 const char unicode_op = UNICODE;
1943
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001944 encoded = raw_unicode_escape(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001945 if (encoded == NULL)
Antoine Pitrou299978d2013-04-07 17:38:11 +02001946 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001947
Antoine Pitrou299978d2013-04-07 17:38:11 +02001948 if (_Pickler_Write(self, &unicode_op, 1) < 0) {
1949 Py_DECREF(encoded);
1950 return -1;
1951 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001952
1953 size = PyBytes_GET_SIZE(encoded);
Antoine Pitrou299978d2013-04-07 17:38:11 +02001954 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0) {
1955 Py_DECREF(encoded);
1956 return -1;
1957 }
1958 Py_DECREF(encoded);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001959
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001960 if (_Pickler_Write(self, "\n", 1) < 0)
Antoine Pitrou299978d2013-04-07 17:38:11 +02001961 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001962 }
1963 if (memo_put(self, obj) < 0)
Antoine Pitrou299978d2013-04-07 17:38:11 +02001964 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001965
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001966 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001967}
1968
1969/* A helper for save_tuple. Push the len elements in tuple t on the stack. */
1970static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001971store_tuple_elements(PicklerObject *self, PyObject *t, Py_ssize_t len)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001972{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001973 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001974
1975 assert(PyTuple_Size(t) == len);
1976
1977 for (i = 0; i < len; i++) {
1978 PyObject *element = PyTuple_GET_ITEM(t, i);
1979
1980 if (element == NULL)
1981 return -1;
1982 if (save(self, element, 0) < 0)
1983 return -1;
1984 }
1985
1986 return 0;
1987}
1988
1989/* Tuples are ubiquitous in the pickle protocols, so many techniques are
1990 * used across protocols to minimize the space needed to pickle them.
1991 * Tuples are also the only builtin immutable type that can be recursive
1992 * (a tuple can be reached from itself), and that requires some subtle
1993 * magic so that it works in all cases. IOW, this is a long routine.
1994 */
1995static int
1996save_tuple(PicklerObject *self, PyObject *obj)
1997{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001998 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001999
2000 const char mark_op = MARK;
2001 const char tuple_op = TUPLE;
2002 const char pop_op = POP;
2003 const char pop_mark_op = POP_MARK;
2004 const char len2opcode[] = {EMPTY_TUPLE, TUPLE1, TUPLE2, TUPLE3};
2005
2006 if ((len = PyTuple_Size(obj)) < 0)
2007 return -1;
2008
2009 if (len == 0) {
2010 char pdata[2];
2011
2012 if (self->proto) {
2013 pdata[0] = EMPTY_TUPLE;
2014 len = 1;
2015 }
2016 else {
2017 pdata[0] = MARK;
2018 pdata[1] = TUPLE;
2019 len = 2;
2020 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002021 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002022 return -1;
2023 return 0;
2024 }
2025
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002026 /* The tuple isn't in the memo now. If it shows up there after
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002027 * saving the tuple elements, the tuple must be recursive, in
2028 * which case we'll pop everything we put on the stack, and fetch
2029 * its value from the memo.
2030 */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002031 if (len <= 3 && self->proto >= 2) {
2032 /* Use TUPLE{1,2,3} opcodes. */
2033 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002034 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002035
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002036 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002037 /* pop the len elements */
2038 for (i = 0; i < len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002039 if (_Pickler_Write(self, &pop_op, 1) < 0)
2040 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002041 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002042 if (memo_get(self, obj) < 0)
2043 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002044
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002045 return 0;
2046 }
2047 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002048 if (_Pickler_Write(self, len2opcode + len, 1) < 0)
2049 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002050 }
2051 goto memoize;
2052 }
2053
2054 /* proto < 2 and len > 0, or proto >= 2 and len > 3.
2055 * Generate MARK e1 e2 ... TUPLE
2056 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002057 if (_Pickler_Write(self, &mark_op, 1) < 0)
2058 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002059
2060 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002061 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002062
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002063 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002064 /* pop the stack stuff we pushed */
2065 if (self->bin) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002066 if (_Pickler_Write(self, &pop_mark_op, 1) < 0)
2067 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002068 }
2069 else {
2070 /* Note that we pop one more than len, to remove
2071 * the MARK too.
2072 */
2073 for (i = 0; i <= len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002074 if (_Pickler_Write(self, &pop_op, 1) < 0)
2075 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002076 }
2077 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002078 if (memo_get(self, obj) < 0)
2079 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002080
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002081 return 0;
2082 }
2083 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002084 if (_Pickler_Write(self, &tuple_op, 1) < 0)
2085 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002086 }
2087
2088 memoize:
2089 if (memo_put(self, obj) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002090 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002091
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002092 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002093}
2094
2095/* iter is an iterator giving items, and we batch up chunks of
2096 * MARK item item ... item APPENDS
2097 * opcode sequences. Calling code should have arranged to first create an
2098 * empty list, or list-like object, for the APPENDS to operate on.
2099 * Returns 0 on success, <0 on error.
2100 */
2101static int
2102batch_list(PicklerObject *self, PyObject *iter)
2103{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002104 PyObject *obj = NULL;
2105 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002106 int i, n;
2107
2108 const char mark_op = MARK;
2109 const char append_op = APPEND;
2110 const char appends_op = APPENDS;
2111
2112 assert(iter != NULL);
2113
2114 /* XXX: I think this function could be made faster by avoiding the
2115 iterator interface and fetching objects directly from list using
2116 PyList_GET_ITEM.
2117 */
2118
2119 if (self->proto == 0) {
2120 /* APPENDS isn't available; do one at a time. */
2121 for (;;) {
2122 obj = PyIter_Next(iter);
2123 if (obj == NULL) {
2124 if (PyErr_Occurred())
2125 return -1;
2126 break;
2127 }
2128 i = save(self, obj, 0);
2129 Py_DECREF(obj);
2130 if (i < 0)
2131 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002132 if (_Pickler_Write(self, &append_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002133 return -1;
2134 }
2135 return 0;
2136 }
2137
2138 /* proto > 0: write in batches of BATCHSIZE. */
2139 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002140 /* Get first item */
2141 firstitem = PyIter_Next(iter);
2142 if (firstitem == NULL) {
2143 if (PyErr_Occurred())
2144 goto error;
2145
2146 /* nothing more to add */
2147 break;
2148 }
2149
2150 /* Try to get a second item */
2151 obj = PyIter_Next(iter);
2152 if (obj == NULL) {
2153 if (PyErr_Occurred())
2154 goto error;
2155
2156 /* Only one item to write */
2157 if (save(self, firstitem, 0) < 0)
2158 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002159 if (_Pickler_Write(self, &append_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002160 goto error;
2161 Py_CLEAR(firstitem);
2162 break;
2163 }
2164
2165 /* More than one item to write */
2166
2167 /* Pump out MARK, items, APPENDS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002168 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002169 goto error;
2170
2171 if (save(self, firstitem, 0) < 0)
2172 goto error;
2173 Py_CLEAR(firstitem);
2174 n = 1;
2175
2176 /* Fetch and save up to BATCHSIZE items */
2177 while (obj) {
2178 if (save(self, obj, 0) < 0)
2179 goto error;
2180 Py_CLEAR(obj);
2181 n += 1;
2182
2183 if (n == BATCHSIZE)
2184 break;
2185
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002186 obj = PyIter_Next(iter);
2187 if (obj == NULL) {
2188 if (PyErr_Occurred())
2189 goto error;
2190 break;
2191 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002192 }
2193
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002194 if (_Pickler_Write(self, &appends_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002195 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002196
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002197 } while (n == BATCHSIZE);
2198 return 0;
2199
2200 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002201 Py_XDECREF(firstitem);
2202 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002203 return -1;
2204}
2205
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002206/* This is a variant of batch_list() above, specialized for lists (with no
2207 * support for list subclasses). Like batch_list(), we batch up chunks of
2208 * MARK item item ... item APPENDS
2209 * opcode sequences. Calling code should have arranged to first create an
2210 * empty list, or list-like object, for the APPENDS to operate on.
2211 * Returns 0 on success, -1 on error.
2212 *
2213 * This version is considerably faster than batch_list(), if less general.
2214 *
2215 * Note that this only works for protocols > 0.
2216 */
2217static int
2218batch_list_exact(PicklerObject *self, PyObject *obj)
2219{
2220 PyObject *item = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002221 Py_ssize_t this_batch, total;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002222
2223 const char append_op = APPEND;
2224 const char appends_op = APPENDS;
2225 const char mark_op = MARK;
2226
2227 assert(obj != NULL);
2228 assert(self->proto > 0);
2229 assert(PyList_CheckExact(obj));
2230
2231 if (PyList_GET_SIZE(obj) == 1) {
2232 item = PyList_GET_ITEM(obj, 0);
2233 if (save(self, item, 0) < 0)
2234 return -1;
2235 if (_Pickler_Write(self, &append_op, 1) < 0)
2236 return -1;
2237 return 0;
2238 }
2239
2240 /* Write in batches of BATCHSIZE. */
2241 total = 0;
2242 do {
2243 this_batch = 0;
2244 if (_Pickler_Write(self, &mark_op, 1) < 0)
2245 return -1;
2246 while (total < PyList_GET_SIZE(obj)) {
2247 item = PyList_GET_ITEM(obj, total);
2248 if (save(self, item, 0) < 0)
2249 return -1;
2250 total++;
2251 if (++this_batch == BATCHSIZE)
2252 break;
2253 }
2254 if (_Pickler_Write(self, &appends_op, 1) < 0)
2255 return -1;
2256
2257 } while (total < PyList_GET_SIZE(obj));
2258
2259 return 0;
2260}
2261
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002262static int
2263save_list(PicklerObject *self, PyObject *obj)
2264{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002265 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002266 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002267 int status = 0;
2268
2269 if (self->fast && !fast_save_enter(self, obj))
2270 goto error;
2271
2272 /* Create an empty list. */
2273 if (self->bin) {
2274 header[0] = EMPTY_LIST;
2275 len = 1;
2276 }
2277 else {
2278 header[0] = MARK;
2279 header[1] = LIST;
2280 len = 2;
2281 }
2282
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002283 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002284 goto error;
2285
2286 /* Get list length, and bow out early if empty. */
2287 if ((len = PyList_Size(obj)) < 0)
2288 goto error;
2289
2290 if (memo_put(self, obj) < 0)
2291 goto error;
2292
2293 if (len != 0) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002294 /* Materialize the list elements. */
2295 if (PyList_CheckExact(obj) && self->proto > 0) {
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002296 if (Py_EnterRecursiveCall(" while pickling an object"))
2297 goto error;
2298 status = batch_list_exact(self, obj);
2299 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002300 } else {
2301 PyObject *iter = PyObject_GetIter(obj);
2302 if (iter == NULL)
2303 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002304
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002305 if (Py_EnterRecursiveCall(" while pickling an object")) {
2306 Py_DECREF(iter);
2307 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002308 }
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002309 status = batch_list(self, iter);
2310 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002311 Py_DECREF(iter);
2312 }
2313 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002314 if (0) {
2315 error:
2316 status = -1;
2317 }
2318
2319 if (self->fast && !fast_save_leave(self, obj))
2320 status = -1;
2321
2322 return status;
2323}
2324
2325/* iter is an iterator giving (key, value) pairs, and we batch up chunks of
2326 * MARK key value ... key value SETITEMS
2327 * opcode sequences. Calling code should have arranged to first create an
2328 * empty dict, or dict-like object, for the SETITEMS to operate on.
2329 * Returns 0 on success, <0 on error.
2330 *
2331 * This is very much like batch_list(). The difference between saving
2332 * elements directly, and picking apart two-tuples, is so long-winded at
2333 * the C level, though, that attempts to combine these routines were too
2334 * ugly to bear.
2335 */
2336static int
2337batch_dict(PicklerObject *self, PyObject *iter)
2338{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002339 PyObject *obj = NULL;
2340 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002341 int i, n;
2342
2343 const char mark_op = MARK;
2344 const char setitem_op = SETITEM;
2345 const char setitems_op = SETITEMS;
2346
2347 assert(iter != NULL);
2348
2349 if (self->proto == 0) {
2350 /* SETITEMS isn't available; do one at a time. */
2351 for (;;) {
2352 obj = PyIter_Next(iter);
2353 if (obj == NULL) {
2354 if (PyErr_Occurred())
2355 return -1;
2356 break;
2357 }
2358 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2359 PyErr_SetString(PyExc_TypeError, "dict items "
2360 "iterator must return 2-tuples");
2361 return -1;
2362 }
2363 i = save(self, PyTuple_GET_ITEM(obj, 0), 0);
2364 if (i >= 0)
2365 i = save(self, PyTuple_GET_ITEM(obj, 1), 0);
2366 Py_DECREF(obj);
2367 if (i < 0)
2368 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002369 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002370 return -1;
2371 }
2372 return 0;
2373 }
2374
2375 /* proto > 0: write in batches of BATCHSIZE. */
2376 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002377 /* Get first item */
2378 firstitem = PyIter_Next(iter);
2379 if (firstitem == NULL) {
2380 if (PyErr_Occurred())
2381 goto error;
2382
2383 /* nothing more to add */
2384 break;
2385 }
2386 if (!PyTuple_Check(firstitem) || PyTuple_Size(firstitem) != 2) {
2387 PyErr_SetString(PyExc_TypeError, "dict items "
2388 "iterator must return 2-tuples");
2389 goto error;
2390 }
2391
2392 /* Try to get a second item */
2393 obj = PyIter_Next(iter);
2394 if (obj == NULL) {
2395 if (PyErr_Occurred())
2396 goto error;
2397
2398 /* Only one item to write */
2399 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2400 goto error;
2401 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2402 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002403 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002404 goto error;
2405 Py_CLEAR(firstitem);
2406 break;
2407 }
2408
2409 /* More than one item to write */
2410
2411 /* Pump out MARK, items, SETITEMS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002412 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002413 goto error;
2414
2415 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2416 goto error;
2417 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2418 goto error;
2419 Py_CLEAR(firstitem);
2420 n = 1;
2421
2422 /* Fetch and save up to BATCHSIZE items */
2423 while (obj) {
2424 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2425 PyErr_SetString(PyExc_TypeError, "dict items "
2426 "iterator must return 2-tuples");
2427 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002428 }
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002429 if (save(self, PyTuple_GET_ITEM(obj, 0), 0) < 0 ||
2430 save(self, PyTuple_GET_ITEM(obj, 1), 0) < 0)
2431 goto error;
2432 Py_CLEAR(obj);
2433 n += 1;
2434
2435 if (n == BATCHSIZE)
2436 break;
2437
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002438 obj = PyIter_Next(iter);
2439 if (obj == NULL) {
2440 if (PyErr_Occurred())
2441 goto error;
2442 break;
2443 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002444 }
2445
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002446 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002447 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002448
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002449 } while (n == BATCHSIZE);
2450 return 0;
2451
2452 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002453 Py_XDECREF(firstitem);
2454 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002455 return -1;
2456}
2457
Collin Winter5c9b02d2009-05-25 05:43:30 +00002458/* This is a variant of batch_dict() above that specializes for dicts, with no
2459 * support for dict subclasses. Like batch_dict(), we batch up chunks of
2460 * MARK key value ... key value SETITEMS
2461 * opcode sequences. Calling code should have arranged to first create an
2462 * empty dict, or dict-like object, for the SETITEMS to operate on.
2463 * Returns 0 on success, -1 on error.
2464 *
2465 * Note that this currently doesn't work for protocol 0.
2466 */
2467static int
2468batch_dict_exact(PicklerObject *self, PyObject *obj)
2469{
2470 PyObject *key = NULL, *value = NULL;
2471 int i;
2472 Py_ssize_t dict_size, ppos = 0;
2473
Alexandre Vassalottif70b1292009-05-25 18:00:52 +00002474 const char mark_op = MARK;
2475 const char setitem_op = SETITEM;
2476 const char setitems_op = SETITEMS;
Collin Winter5c9b02d2009-05-25 05:43:30 +00002477
2478 assert(obj != NULL);
2479 assert(self->proto > 0);
2480
2481 dict_size = PyDict_Size(obj);
2482
2483 /* Special-case len(d) == 1 to save space. */
2484 if (dict_size == 1) {
2485 PyDict_Next(obj, &ppos, &key, &value);
2486 if (save(self, key, 0) < 0)
2487 return -1;
2488 if (save(self, value, 0) < 0)
2489 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002490 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002491 return -1;
2492 return 0;
2493 }
2494
2495 /* Write in batches of BATCHSIZE. */
2496 do {
2497 i = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002498 if (_Pickler_Write(self, &mark_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002499 return -1;
2500 while (PyDict_Next(obj, &ppos, &key, &value)) {
2501 if (save(self, key, 0) < 0)
2502 return -1;
2503 if (save(self, value, 0) < 0)
2504 return -1;
2505 if (++i == BATCHSIZE)
2506 break;
2507 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002508 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002509 return -1;
2510 if (PyDict_Size(obj) != dict_size) {
2511 PyErr_Format(
2512 PyExc_RuntimeError,
2513 "dictionary changed size during iteration");
2514 return -1;
2515 }
2516
2517 } while (i == BATCHSIZE);
2518 return 0;
2519}
2520
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002521static int
2522save_dict(PicklerObject *self, PyObject *obj)
2523{
2524 PyObject *items, *iter;
2525 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002526 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002527 int status = 0;
2528
2529 if (self->fast && !fast_save_enter(self, obj))
2530 goto error;
2531
2532 /* Create an empty dict. */
2533 if (self->bin) {
2534 header[0] = EMPTY_DICT;
2535 len = 1;
2536 }
2537 else {
2538 header[0] = MARK;
2539 header[1] = DICT;
2540 len = 2;
2541 }
2542
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002543 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002544 goto error;
2545
2546 /* Get dict size, and bow out early if empty. */
2547 if ((len = PyDict_Size(obj)) < 0)
2548 goto error;
2549
2550 if (memo_put(self, obj) < 0)
2551 goto error;
2552
2553 if (len != 0) {
2554 /* Save the dict items. */
Collin Winter5c9b02d2009-05-25 05:43:30 +00002555 if (PyDict_CheckExact(obj) && self->proto > 0) {
2556 /* We can take certain shortcuts if we know this is a dict and
2557 not a dict subclass. */
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002558 if (Py_EnterRecursiveCall(" while pickling an object"))
2559 goto error;
2560 status = batch_dict_exact(self, obj);
2561 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002562 } else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02002563 _Py_IDENTIFIER(items);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02002564
2565 items = _PyObject_CallMethodId(obj, &PyId_items, "()");
Collin Winter5c9b02d2009-05-25 05:43:30 +00002566 if (items == NULL)
2567 goto error;
2568 iter = PyObject_GetIter(items);
2569 Py_DECREF(items);
2570 if (iter == NULL)
2571 goto error;
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002572 if (Py_EnterRecursiveCall(" while pickling an object")) {
2573 Py_DECREF(iter);
2574 goto error;
2575 }
Collin Winter5c9b02d2009-05-25 05:43:30 +00002576 status = batch_dict(self, iter);
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002577 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002578 Py_DECREF(iter);
2579 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002580 }
2581
2582 if (0) {
2583 error:
2584 status = -1;
2585 }
2586
2587 if (self->fast && !fast_save_leave(self, obj))
2588 status = -1;
2589
2590 return status;
2591}
2592
2593static int
2594save_global(PicklerObject *self, PyObject *obj, PyObject *name)
2595{
2596 static PyObject *name_str = NULL;
2597 PyObject *global_name = NULL;
2598 PyObject *module_name = NULL;
2599 PyObject *module = NULL;
2600 PyObject *cls;
2601 int status = 0;
2602
2603 const char global_op = GLOBAL;
2604
2605 if (name_str == NULL) {
2606 name_str = PyUnicode_InternFromString("__name__");
2607 if (name_str == NULL)
2608 goto error;
2609 }
2610
2611 if (name) {
2612 global_name = name;
2613 Py_INCREF(global_name);
2614 }
2615 else {
2616 global_name = PyObject_GetAttr(obj, name_str);
2617 if (global_name == NULL)
2618 goto error;
2619 }
2620
2621 module_name = whichmodule(obj, global_name);
2622 if (module_name == NULL)
2623 goto error;
2624
2625 /* XXX: Change to use the import C API directly with level=0 to disallow
2626 relative imports.
2627
2628 XXX: PyImport_ImportModuleLevel could be used. However, this bypasses
2629 builtins.__import__. Therefore, _pickle, unlike pickle.py, will ignore
2630 custom import functions (IMHO, this would be a nice security
2631 feature). The import C API would need to be extended to support the
2632 extra parameters of __import__ to fix that. */
2633 module = PyImport_Import(module_name);
2634 if (module == NULL) {
2635 PyErr_Format(PicklingError,
2636 "Can't pickle %R: import of module %R failed",
2637 obj, module_name);
2638 goto error;
2639 }
2640 cls = PyObject_GetAttr(module, global_name);
2641 if (cls == NULL) {
2642 PyErr_Format(PicklingError,
2643 "Can't pickle %R: attribute lookup %S.%S failed",
2644 obj, module_name, global_name);
2645 goto error;
2646 }
2647 if (cls != obj) {
2648 Py_DECREF(cls);
2649 PyErr_Format(PicklingError,
2650 "Can't pickle %R: it's not the same object as %S.%S",
2651 obj, module_name, global_name);
2652 goto error;
2653 }
2654 Py_DECREF(cls);
2655
2656 if (self->proto >= 2) {
2657 /* See whether this is in the extension registry, and if
2658 * so generate an EXT opcode.
2659 */
2660 PyObject *code_obj; /* extension code as Python object */
2661 long code; /* extension code as C value */
2662 char pdata[5];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002663 Py_ssize_t n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002664
2665 PyTuple_SET_ITEM(two_tuple, 0, module_name);
2666 PyTuple_SET_ITEM(two_tuple, 1, global_name);
2667 code_obj = PyDict_GetItem(extension_registry, two_tuple);
2668 /* The object is not registered in the extension registry.
2669 This is the most likely code path. */
2670 if (code_obj == NULL)
2671 goto gen_global;
2672
2673 /* XXX: pickle.py doesn't check neither the type, nor the range
2674 of the value returned by the extension_registry. It should for
2675 consistency. */
2676
2677 /* Verify code_obj has the right type and value. */
2678 if (!PyLong_Check(code_obj)) {
2679 PyErr_Format(PicklingError,
2680 "Can't pickle %R: extension code %R isn't an integer",
2681 obj, code_obj);
2682 goto error;
2683 }
2684 code = PyLong_AS_LONG(code_obj);
2685 if (code <= 0 || code > 0x7fffffffL) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002686 if (!PyErr_Occurred())
2687 PyErr_Format(PicklingError,
2688 "Can't pickle %R: extension code %ld is out of range",
2689 obj, code);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002690 goto error;
2691 }
2692
2693 /* Generate an EXT opcode. */
2694 if (code <= 0xff) {
2695 pdata[0] = EXT1;
2696 pdata[1] = (unsigned char)code;
2697 n = 2;
2698 }
2699 else if (code <= 0xffff) {
2700 pdata[0] = EXT2;
2701 pdata[1] = (unsigned char)(code & 0xff);
2702 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2703 n = 3;
2704 }
2705 else {
2706 pdata[0] = EXT4;
2707 pdata[1] = (unsigned char)(code & 0xff);
2708 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2709 pdata[3] = (unsigned char)((code >> 16) & 0xff);
2710 pdata[4] = (unsigned char)((code >> 24) & 0xff);
2711 n = 5;
2712 }
2713
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002714 if (_Pickler_Write(self, pdata, n) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002715 goto error;
2716 }
2717 else {
2718 /* Generate a normal global opcode if we are using a pickle
2719 protocol <= 2, or if the object is not registered in the
2720 extension registry. */
2721 PyObject *encoded;
2722 PyObject *(*unicode_encoder)(PyObject *);
2723
2724 gen_global:
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002725 if (_Pickler_Write(self, &global_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002726 goto error;
2727
2728 /* Since Python 3.0 now supports non-ASCII identifiers, we encode both
2729 the module name and the global name using UTF-8. We do so only when
2730 we are using the pickle protocol newer than version 3. This is to
2731 ensure compatibility with older Unpickler running on Python 2.x. */
2732 if (self->proto >= 3) {
2733 unicode_encoder = PyUnicode_AsUTF8String;
2734 }
2735 else {
2736 unicode_encoder = PyUnicode_AsASCIIString;
2737 }
2738
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00002739 /* For protocol < 3 and if the user didn't request against doing so,
2740 we convert module names to the old 2.x module names. */
2741 if (self->fix_imports) {
2742 PyObject *key;
2743 PyObject *item;
2744
2745 key = PyTuple_Pack(2, module_name, global_name);
2746 if (key == NULL)
2747 goto error;
2748 item = PyDict_GetItemWithError(name_mapping_3to2, key);
2749 Py_DECREF(key);
2750 if (item) {
2751 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
2752 PyErr_Format(PyExc_RuntimeError,
2753 "_compat_pickle.REVERSE_NAME_MAPPING values "
2754 "should be 2-tuples, not %.200s",
2755 Py_TYPE(item)->tp_name);
2756 goto error;
2757 }
2758 Py_CLEAR(module_name);
2759 Py_CLEAR(global_name);
2760 module_name = PyTuple_GET_ITEM(item, 0);
2761 global_name = PyTuple_GET_ITEM(item, 1);
2762 if (!PyUnicode_Check(module_name) ||
2763 !PyUnicode_Check(global_name)) {
2764 PyErr_Format(PyExc_RuntimeError,
2765 "_compat_pickle.REVERSE_NAME_MAPPING values "
2766 "should be pairs of str, not (%.200s, %.200s)",
2767 Py_TYPE(module_name)->tp_name,
2768 Py_TYPE(global_name)->tp_name);
2769 goto error;
2770 }
2771 Py_INCREF(module_name);
2772 Py_INCREF(global_name);
2773 }
2774 else if (PyErr_Occurred()) {
2775 goto error;
2776 }
2777
2778 item = PyDict_GetItemWithError(import_mapping_3to2, module_name);
2779 if (item) {
2780 if (!PyUnicode_Check(item)) {
2781 PyErr_Format(PyExc_RuntimeError,
2782 "_compat_pickle.REVERSE_IMPORT_MAPPING values "
2783 "should be strings, not %.200s",
2784 Py_TYPE(item)->tp_name);
2785 goto error;
2786 }
2787 Py_CLEAR(module_name);
2788 module_name = item;
2789 Py_INCREF(module_name);
2790 }
2791 else if (PyErr_Occurred()) {
2792 goto error;
2793 }
2794 }
2795
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002796 /* Save the name of the module. */
2797 encoded = unicode_encoder(module_name);
2798 if (encoded == NULL) {
2799 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2800 PyErr_Format(PicklingError,
2801 "can't pickle module identifier '%S' using "
2802 "pickle protocol %i", module_name, self->proto);
2803 goto error;
2804 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002805 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002806 PyBytes_GET_SIZE(encoded)) < 0) {
2807 Py_DECREF(encoded);
2808 goto error;
2809 }
2810 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002811 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002812 goto error;
2813
2814 /* Save the name of the module. */
2815 encoded = unicode_encoder(global_name);
2816 if (encoded == NULL) {
2817 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2818 PyErr_Format(PicklingError,
2819 "can't pickle global identifier '%S' using "
2820 "pickle protocol %i", global_name, self->proto);
2821 goto error;
2822 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002823 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002824 PyBytes_GET_SIZE(encoded)) < 0) {
2825 Py_DECREF(encoded);
2826 goto error;
2827 }
2828 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002829 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002830 goto error;
2831
2832 /* Memoize the object. */
2833 if (memo_put(self, obj) < 0)
2834 goto error;
2835 }
2836
2837 if (0) {
2838 error:
2839 status = -1;
2840 }
2841 Py_XDECREF(module_name);
2842 Py_XDECREF(global_name);
2843 Py_XDECREF(module);
2844
2845 return status;
2846}
2847
2848static int
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002849save_ellipsis(PicklerObject *self, PyObject *obj)
2850{
Łukasz Langadbd78252012-03-12 22:59:11 +01002851 PyObject *str = PyUnicode_FromString("Ellipsis");
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002852 int res;
Łukasz Langadbd78252012-03-12 22:59:11 +01002853 if (str == NULL)
Łukasz Langacad1a072012-03-12 23:41:07 +01002854 return -1;
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002855 res = save_global(self, Py_Ellipsis, str);
2856 Py_DECREF(str);
2857 return res;
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002858}
2859
2860static int
2861save_notimplemented(PicklerObject *self, PyObject *obj)
2862{
Łukasz Langadbd78252012-03-12 22:59:11 +01002863 PyObject *str = PyUnicode_FromString("NotImplemented");
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002864 int res;
Łukasz Langadbd78252012-03-12 22:59:11 +01002865 if (str == NULL)
Łukasz Langacad1a072012-03-12 23:41:07 +01002866 return -1;
Benjamin Petersone80b29b2012-03-16 18:45:31 -05002867 res = save_global(self, Py_NotImplemented, str);
2868 Py_DECREF(str);
2869 return res;
Łukasz Langaf3078fb2012-03-12 19:46:12 +01002870}
2871
2872static int
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002873save_pers(PicklerObject *self, PyObject *obj, PyObject *func)
2874{
2875 PyObject *pid = NULL;
2876 int status = 0;
2877
2878 const char persid_op = PERSID;
2879 const char binpersid_op = BINPERSID;
2880
2881 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002882 pid = _Pickler_FastCall(self, func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002883 if (pid == NULL)
2884 return -1;
2885
2886 if (pid != Py_None) {
2887 if (self->bin) {
2888 if (save(self, pid, 1) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002889 _Pickler_Write(self, &binpersid_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002890 goto error;
2891 }
2892 else {
2893 PyObject *pid_str = NULL;
2894 char *pid_ascii_bytes;
2895 Py_ssize_t size;
2896
2897 pid_str = PyObject_Str(pid);
2898 if (pid_str == NULL)
2899 goto error;
2900
2901 /* XXX: Should it check whether the persistent id only contains
2902 ASCII characters? And what if the pid contains embedded
2903 newlines? */
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00002904 pid_ascii_bytes = _PyUnicode_AsStringAndSize(pid_str, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002905 Py_DECREF(pid_str);
2906 if (pid_ascii_bytes == NULL)
2907 goto error;
2908
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002909 if (_Pickler_Write(self, &persid_op, 1) < 0 ||
2910 _Pickler_Write(self, pid_ascii_bytes, size) < 0 ||
2911 _Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002912 goto error;
2913 }
2914 status = 1;
2915 }
2916
2917 if (0) {
2918 error:
2919 status = -1;
2920 }
2921 Py_XDECREF(pid);
2922
2923 return status;
2924}
2925
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002926static PyObject *
2927get_class(PyObject *obj)
2928{
2929 PyObject *cls;
2930 static PyObject *str_class;
2931
2932 if (str_class == NULL) {
2933 str_class = PyUnicode_InternFromString("__class__");
2934 if (str_class == NULL)
2935 return NULL;
2936 }
2937 cls = PyObject_GetAttr(obj, str_class);
2938 if (cls == NULL) {
2939 if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
2940 PyErr_Clear();
2941 cls = (PyObject *) Py_TYPE(obj);
2942 Py_INCREF(cls);
2943 }
2944 }
2945 return cls;
2946}
2947
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002948/* We're saving obj, and args is the 2-thru-5 tuple returned by the
2949 * appropriate __reduce__ method for obj.
2950 */
2951static int
2952save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
2953{
2954 PyObject *callable;
2955 PyObject *argtup;
2956 PyObject *state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002957 PyObject *listitems = Py_None;
2958 PyObject *dictitems = Py_None;
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002959 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002960
2961 int use_newobj = self->proto >= 2;
2962
2963 const char reduce_op = REDUCE;
2964 const char build_op = BUILD;
2965 const char newobj_op = NEWOBJ;
2966
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002967 size = PyTuple_Size(args);
2968 if (size < 2 || size > 5) {
2969 PyErr_SetString(PicklingError, "tuple returned by "
2970 "__reduce__ must contain 2 through 5 elements");
2971 return -1;
2972 }
2973
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002974 if (!PyArg_UnpackTuple(args, "save_reduce", 2, 5,
2975 &callable, &argtup, &state, &listitems, &dictitems))
2976 return -1;
2977
2978 if (!PyCallable_Check(callable)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002979 PyErr_SetString(PicklingError, "first item of the tuple "
2980 "returned by __reduce__ must be callable");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002981 return -1;
2982 }
2983 if (!PyTuple_Check(argtup)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002984 PyErr_SetString(PicklingError, "second item of the tuple "
2985 "returned by __reduce__ must be a tuple");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002986 return -1;
2987 }
2988
2989 if (state == Py_None)
2990 state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002991
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002992 if (listitems == Py_None)
2993 listitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002994 else if (!PyIter_Check(listitems)) {
2995 PyErr_Format(PicklingError, "Fourth element of tuple"
2996 "returned by __reduce__ must be an iterator, not %s",
2997 Py_TYPE(listitems)->tp_name);
2998 return -1;
2999 }
3000
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003001 if (dictitems == Py_None)
3002 dictitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00003003 else if (!PyIter_Check(dictitems)) {
3004 PyErr_Format(PicklingError, "Fifth element of tuple"
3005 "returned by __reduce__ must be an iterator, not %s",
3006 Py_TYPE(dictitems)->tp_name);
3007 return -1;
3008 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003009
3010 /* Protocol 2 special case: if callable's name is __newobj__, use
3011 NEWOBJ. */
3012 if (use_newobj) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003013 static PyObject *newobj_str = NULL, *name_str = NULL;
3014 PyObject *name;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003015
3016 if (newobj_str == NULL) {
3017 newobj_str = PyUnicode_InternFromString("__newobj__");
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003018 name_str = PyUnicode_InternFromString("__name__");
3019 if (newobj_str == NULL || name_str == NULL)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003020 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003021 }
3022
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003023 name = PyObject_GetAttr(callable, name_str);
3024 if (name == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003025 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3026 PyErr_Clear();
3027 else
3028 return -1;
3029 use_newobj = 0;
3030 }
3031 else {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003032 use_newobj = PyUnicode_Check(name) &&
3033 PyUnicode_Compare(name, newobj_str) == 0;
3034 Py_DECREF(name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003035 }
3036 }
3037 if (use_newobj) {
3038 PyObject *cls;
3039 PyObject *newargtup;
3040 PyObject *obj_class;
3041 int p;
3042
3043 /* Sanity checks. */
3044 if (Py_SIZE(argtup) < 1) {
3045 PyErr_SetString(PicklingError, "__newobj__ arglist is empty");
3046 return -1;
3047 }
3048
3049 cls = PyTuple_GET_ITEM(argtup, 0);
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003050 if (!PyType_Check(cls)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003051 PyErr_SetString(PicklingError, "args[0] from "
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003052 "__newobj__ args is not a type");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003053 return -1;
3054 }
3055
3056 if (obj != NULL) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01003057 obj_class = get_class(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003058 p = obj_class != cls; /* true iff a problem */
3059 Py_DECREF(obj_class);
3060 if (p) {
3061 PyErr_SetString(PicklingError, "args[0] from "
3062 "__newobj__ args has the wrong class");
3063 return -1;
3064 }
3065 }
3066 /* XXX: These calls save() are prone to infinite recursion. Imagine
3067 what happen if the value returned by the __reduce__() method of
3068 some extension type contains another object of the same type. Ouch!
3069
3070 Here is a quick example, that I ran into, to illustrate what I
3071 mean:
3072
3073 >>> import pickle, copyreg
3074 >>> copyreg.dispatch_table.pop(complex)
3075 >>> pickle.dumps(1+2j)
3076 Traceback (most recent call last):
3077 ...
3078 RuntimeError: maximum recursion depth exceeded
3079
3080 Removing the complex class from copyreg.dispatch_table made the
3081 __reduce_ex__() method emit another complex object:
3082
3083 >>> (1+1j).__reduce_ex__(2)
3084 (<function __newobj__ at 0xb7b71c3c>,
3085 (<class 'complex'>, (1+1j)), None, None, None)
3086
3087 Thus when save() was called on newargstup (the 2nd item) recursion
3088 ensued. Of course, the bug was in the complex class which had a
3089 broken __getnewargs__() that emitted another complex object. But,
3090 the point, here, is it is quite easy to end up with a broken reduce
3091 function. */
3092
3093 /* Save the class and its __new__ arguments. */
3094 if (save(self, cls, 0) < 0)
3095 return -1;
3096
3097 newargtup = PyTuple_GetSlice(argtup, 1, Py_SIZE(argtup));
3098 if (newargtup == NULL)
3099 return -1;
3100
3101 p = save(self, newargtup, 0);
3102 Py_DECREF(newargtup);
3103 if (p < 0)
3104 return -1;
3105
3106 /* Add NEWOBJ opcode. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003107 if (_Pickler_Write(self, &newobj_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003108 return -1;
3109 }
3110 else { /* Not using NEWOBJ. */
3111 if (save(self, callable, 0) < 0 ||
3112 save(self, argtup, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003113 _Pickler_Write(self, &reduce_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003114 return -1;
3115 }
3116
3117 /* obj can be NULL when save_reduce() is used directly. A NULL obj means
3118 the caller do not want to memoize the object. Not particularly useful,
3119 but that is to mimic the behavior save_reduce() in pickle.py when
3120 obj is None. */
3121 if (obj && memo_put(self, obj) < 0)
3122 return -1;
3123
3124 if (listitems && batch_list(self, listitems) < 0)
3125 return -1;
3126
3127 if (dictitems && batch_dict(self, dictitems) < 0)
3128 return -1;
3129
3130 if (state) {
Victor Stinner121aab42011-09-29 23:40:53 +02003131 if (save(self, state, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003132 _Pickler_Write(self, &build_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003133 return -1;
3134 }
3135
3136 return 0;
3137}
3138
3139static int
3140save(PicklerObject *self, PyObject *obj, int pers_save)
3141{
3142 PyTypeObject *type;
3143 PyObject *reduce_func = NULL;
3144 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003145 int status = 0;
3146
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00003147 if (Py_EnterRecursiveCall(" while pickling an object"))
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003148 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003149
3150 /* The extra pers_save argument is necessary to avoid calling save_pers()
3151 on its returned object. */
3152 if (!pers_save && self->pers_func) {
3153 /* save_pers() returns:
3154 -1 to signal an error;
3155 0 if it did nothing successfully;
3156 1 if a persistent id was saved.
3157 */
3158 if ((status = save_pers(self, obj, self->pers_func)) != 0)
3159 goto done;
3160 }
3161
3162 type = Py_TYPE(obj);
3163
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003164 /* The old cPickle had an optimization that used switch-case statement
3165 dispatching on the first letter of the type name. This has was removed
3166 since benchmarks shown that this optimization was actually slowing
3167 things down. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003168
3169 /* Atom types; these aren't memoized, so don't check the memo. */
3170
3171 if (obj == Py_None) {
3172 status = save_none(self, obj);
3173 goto done;
3174 }
Łukasz Langaf3078fb2012-03-12 19:46:12 +01003175 else if (obj == Py_Ellipsis) {
3176 status = save_ellipsis(self, obj);
3177 goto done;
3178 }
3179 else if (obj == Py_NotImplemented) {
3180 status = save_notimplemented(self, obj);
3181 goto done;
3182 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003183 else if (obj == Py_False || obj == Py_True) {
3184 status = save_bool(self, obj);
3185 goto done;
3186 }
3187 else if (type == &PyLong_Type) {
3188 status = save_long(self, obj);
3189 goto done;
3190 }
3191 else if (type == &PyFloat_Type) {
3192 status = save_float(self, obj);
3193 goto done;
3194 }
3195
3196 /* Check the memo to see if it has the object. If so, generate
3197 a GET (or BINGET) opcode, instead of pickling the object
3198 once again. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003199 if (PyMemoTable_Get(self->memo, obj)) {
3200 if (memo_get(self, obj) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003201 goto error;
3202 goto done;
3203 }
3204
3205 if (type == &PyBytes_Type) {
3206 status = save_bytes(self, obj);
3207 goto done;
3208 }
3209 else if (type == &PyUnicode_Type) {
3210 status = save_unicode(self, obj);
3211 goto done;
3212 }
3213 else if (type == &PyDict_Type) {
3214 status = save_dict(self, obj);
3215 goto done;
3216 }
3217 else if (type == &PyList_Type) {
3218 status = save_list(self, obj);
3219 goto done;
3220 }
3221 else if (type == &PyTuple_Type) {
3222 status = save_tuple(self, obj);
3223 goto done;
3224 }
3225 else if (type == &PyType_Type) {
3226 status = save_global(self, obj, NULL);
3227 goto done;
3228 }
3229 else if (type == &PyFunction_Type) {
3230 status = save_global(self, obj, NULL);
3231 if (status < 0 && PyErr_ExceptionMatches(PickleError)) {
3232 /* fall back to reduce */
3233 PyErr_Clear();
3234 }
3235 else {
3236 goto done;
3237 }
3238 }
3239 else if (type == &PyCFunction_Type) {
3240 status = save_global(self, obj, NULL);
3241 goto done;
3242 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003243
3244 /* XXX: This part needs some unit tests. */
3245
3246 /* Get a reduction callable, and call it. This may come from
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003247 * self.dispatch_table, copyreg.dispatch_table, the object's
3248 * __reduce_ex__ method, or the object's __reduce__ method.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003249 */
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003250 if (self->dispatch_table == NULL) {
3251 reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
3252 /* PyDict_GetItem() unlike PyObject_GetItem() and
3253 PyObject_GetAttr() returns a borrowed ref */
3254 Py_XINCREF(reduce_func);
3255 } else {
3256 reduce_func = PyObject_GetItem(self->dispatch_table, (PyObject *)type);
3257 if (reduce_func == NULL) {
3258 if (PyErr_ExceptionMatches(PyExc_KeyError))
3259 PyErr_Clear();
3260 else
3261 goto error;
3262 }
3263 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003264 if (reduce_func != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003265 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003266 reduce_value = _Pickler_FastCall(self, reduce_func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003267 }
Antoine Pitrouffd41d92011-10-04 09:23:04 +02003268 else if (PyType_IsSubtype(type, &PyType_Type)) {
3269 status = save_global(self, obj, NULL);
3270 goto done;
3271 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003272 else {
3273 static PyObject *reduce_str = NULL;
3274 static PyObject *reduce_ex_str = NULL;
3275
3276 /* Cache the name of the reduce methods. */
3277 if (reduce_str == NULL) {
3278 reduce_str = PyUnicode_InternFromString("__reduce__");
3279 if (reduce_str == NULL)
3280 goto error;
3281 reduce_ex_str = PyUnicode_InternFromString("__reduce_ex__");
3282 if (reduce_ex_str == NULL)
3283 goto error;
3284 }
3285
3286 /* XXX: If the __reduce__ method is defined, __reduce_ex__ is
3287 automatically defined as __reduce__. While this is convenient, this
3288 make it impossible to know which method was actually called. Of
3289 course, this is not a big deal. But still, it would be nice to let
3290 the user know which method was called when something go
3291 wrong. Incidentally, this means if __reduce_ex__ is not defined, we
3292 don't actually have to check for a __reduce__ method. */
3293
3294 /* Check for a __reduce_ex__ method. */
3295 reduce_func = PyObject_GetAttr(obj, reduce_ex_str);
3296 if (reduce_func != NULL) {
3297 PyObject *proto;
3298 proto = PyLong_FromLong(self->proto);
3299 if (proto != NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003300 reduce_value = _Pickler_FastCall(self, reduce_func, proto);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003301 }
3302 }
3303 else {
3304 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3305 PyErr_Clear();
3306 else
3307 goto error;
3308 /* Check for a __reduce__ method. */
3309 reduce_func = PyObject_GetAttr(obj, reduce_str);
3310 if (reduce_func != NULL) {
3311 reduce_value = PyObject_Call(reduce_func, empty_tuple, NULL);
3312 }
3313 else {
3314 PyErr_Format(PicklingError, "can't pickle '%.200s' object: %R",
3315 type->tp_name, obj);
3316 goto error;
3317 }
3318 }
3319 }
3320
3321 if (reduce_value == NULL)
3322 goto error;
3323
3324 if (PyUnicode_Check(reduce_value)) {
3325 status = save_global(self, obj, reduce_value);
3326 goto done;
3327 }
3328
3329 if (!PyTuple_Check(reduce_value)) {
3330 PyErr_SetString(PicklingError,
3331 "__reduce__ must return a string or tuple");
3332 goto error;
3333 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003334
3335 status = save_reduce(self, reduce_value, obj);
3336
3337 if (0) {
3338 error:
3339 status = -1;
3340 }
3341 done:
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003342 Py_LeaveRecursiveCall();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003343 Py_XDECREF(reduce_func);
3344 Py_XDECREF(reduce_value);
3345
3346 return status;
3347}
3348
3349static int
3350dump(PicklerObject *self, PyObject *obj)
3351{
3352 const char stop_op = STOP;
3353
3354 if (self->proto >= 2) {
3355 char header[2];
3356
3357 header[0] = PROTO;
3358 assert(self->proto >= 0 && self->proto < 256);
3359 header[1] = (unsigned char)self->proto;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003360 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003361 return -1;
3362 }
3363
3364 if (save(self, obj, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003365 _Pickler_Write(self, &stop_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003366 return -1;
3367
3368 return 0;
3369}
3370
3371PyDoc_STRVAR(Pickler_clear_memo_doc,
3372"clear_memo() -> None. Clears the pickler's \"memo\"."
3373"\n"
3374"The memo is the data structure that remembers which objects the\n"
3375"pickler has already seen, so that shared or recursive objects are\n"
3376"pickled by reference and not by value. This method is useful when\n"
3377"re-using picklers.");
3378
3379static PyObject *
3380Pickler_clear_memo(PicklerObject *self)
3381{
3382 if (self->memo)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003383 PyMemoTable_Clear(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003384
3385 Py_RETURN_NONE;
3386}
3387
3388PyDoc_STRVAR(Pickler_dump_doc,
3389"dump(obj) -> None. Write a pickled representation of obj to the open file.");
3390
3391static PyObject *
3392Pickler_dump(PicklerObject *self, PyObject *args)
3393{
3394 PyObject *obj;
3395
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003396 /* Check whether the Pickler was initialized correctly (issue3664).
3397 Developers often forget to call __init__() in their subclasses, which
3398 would trigger a segfault without this check. */
3399 if (self->write == NULL) {
Victor Stinner121aab42011-09-29 23:40:53 +02003400 PyErr_Format(PicklingError,
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003401 "Pickler.__init__() was not called by %s.__init__()",
3402 Py_TYPE(self)->tp_name);
3403 return NULL;
3404 }
3405
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003406 if (!PyArg_ParseTuple(args, "O:dump", &obj))
3407 return NULL;
3408
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003409 if (_Pickler_ClearBuffer(self) < 0)
3410 return NULL;
3411
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003412 if (dump(self, obj) < 0)
3413 return NULL;
3414
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003415 if (_Pickler_FlushToFile(self) < 0)
3416 return NULL;
3417
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003418 Py_RETURN_NONE;
3419}
3420
3421static struct PyMethodDef Pickler_methods[] = {
3422 {"dump", (PyCFunction)Pickler_dump, METH_VARARGS,
3423 Pickler_dump_doc},
3424 {"clear_memo", (PyCFunction)Pickler_clear_memo, METH_NOARGS,
3425 Pickler_clear_memo_doc},
3426 {NULL, NULL} /* sentinel */
3427};
3428
3429static void
3430Pickler_dealloc(PicklerObject *self)
3431{
3432 PyObject_GC_UnTrack(self);
3433
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003434 Py_XDECREF(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003435 Py_XDECREF(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003436 Py_XDECREF(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003437 Py_XDECREF(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003438 Py_XDECREF(self->arg);
3439 Py_XDECREF(self->fast_memo);
3440
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003441 PyMemoTable_Del(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003442
3443 Py_TYPE(self)->tp_free((PyObject *)self);
3444}
3445
3446static int
3447Pickler_traverse(PicklerObject *self, visitproc visit, void *arg)
3448{
3449 Py_VISIT(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003450 Py_VISIT(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003451 Py_VISIT(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003452 Py_VISIT(self->arg);
3453 Py_VISIT(self->fast_memo);
3454 return 0;
3455}
3456
3457static int
3458Pickler_clear(PicklerObject *self)
3459{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003460 Py_CLEAR(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003461 Py_CLEAR(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003462 Py_CLEAR(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003463 Py_CLEAR(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003464 Py_CLEAR(self->arg);
3465 Py_CLEAR(self->fast_memo);
3466
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003467 if (self->memo != NULL) {
3468 PyMemoTable *memo = self->memo;
3469 self->memo = NULL;
3470 PyMemoTable_Del(memo);
3471 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003472 return 0;
3473}
3474
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003475
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003476PyDoc_STRVAR(Pickler_doc,
3477"Pickler(file, protocol=None)"
3478"\n"
3479"This takes a binary file for writing a pickle data stream.\n"
3480"\n"
3481"The optional protocol argument tells the pickler to use the\n"
3482"given protocol; supported protocols are 0, 1, 2, 3. The default\n"
3483"protocol is 3; a backward-incompatible protocol designed for\n"
3484"Python 3.0.\n"
3485"\n"
3486"Specifying a negative protocol version selects the highest\n"
3487"protocol version supported. The higher the protocol used, the\n"
3488"more recent the version of Python needed to read the pickle\n"
3489"produced.\n"
3490"\n"
3491"The file argument must have a write() method that accepts a single\n"
3492"bytes argument. It can thus be a file object opened for binary\n"
3493"writing, a io.BytesIO instance, or any other custom object that\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003494"meets this interface.\n"
3495"\n"
3496"If fix_imports is True and protocol is less than 3, pickle will try to\n"
3497"map the new Python 3.x names to the old module names used in Python\n"
3498"2.x, so that the pickle data stream is readable with Python 2.x.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003499
3500static int
3501Pickler_init(PicklerObject *self, PyObject *args, PyObject *kwds)
3502{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003503 static char *kwlist[] = {"file", "protocol", "fix_imports", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003504 PyObject *file;
3505 PyObject *proto_obj = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003506 PyObject *fix_imports = Py_True;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02003507 _Py_IDENTIFIER(persistent_id);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003508 _Py_IDENTIFIER(dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003509
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003510 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:Pickler",
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003511 kwlist, &file, &proto_obj, &fix_imports))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003512 return -1;
3513
3514 /* In case of multiple __init__() calls, clear previous content. */
3515 if (self->write != NULL)
3516 (void)Pickler_clear(self);
3517
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003518 if (_Pickler_SetProtocol(self, proto_obj, fix_imports) < 0)
3519 return -1;
3520
3521 if (_Pickler_SetOutputStream(self, file) < 0)
3522 return -1;
3523
3524 /* memo and output_buffer may have already been created in _Pickler_New */
3525 if (self->memo == NULL) {
3526 self->memo = PyMemoTable_New();
3527 if (self->memo == NULL)
3528 return -1;
3529 }
3530 self->output_len = 0;
3531 if (self->output_buffer == NULL) {
3532 self->max_output_len = WRITE_BUF_SIZE;
3533 self->output_buffer = PyBytes_FromStringAndSize(NULL,
3534 self->max_output_len);
3535 if (self->output_buffer == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003536 return -1;
3537 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003538
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003539 self->arg = NULL;
3540 self->fast = 0;
3541 self->fast_nesting = 0;
3542 self->fast_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003543 self->pers_func = NULL;
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02003544 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_id)) {
3545 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
3546 &PyId_persistent_id);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003547 if (self->pers_func == NULL)
3548 return -1;
3549 }
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003550 self->dispatch_table = NULL;
3551 if (_PyObject_HasAttrId((PyObject *)self, &PyId_dispatch_table)) {
3552 self->dispatch_table = _PyObject_GetAttrId((PyObject *)self,
3553 &PyId_dispatch_table);
3554 if (self->dispatch_table == NULL)
3555 return -1;
3556 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003557 return 0;
3558}
3559
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003560/* Define a proxy object for the Pickler's internal memo object. This is to
3561 * avoid breaking code like:
3562 * pickler.memo.clear()
3563 * and
3564 * pickler.memo = saved_memo
3565 * Is this a good idea? Not really, but we don't want to break code that uses
3566 * it. Note that we don't implement the entire mapping API here. This is
3567 * intentional, as these should be treated as black-box implementation details.
3568 */
3569
3570typedef struct {
3571 PyObject_HEAD
3572 PicklerObject *pickler; /* Pickler whose memo table we're proxying. */
3573} PicklerMemoProxyObject;
3574
3575PyDoc_STRVAR(pmp_clear_doc,
3576"memo.clear() -> None. Remove all items from memo.");
3577
3578static PyObject *
3579pmp_clear(PicklerMemoProxyObject *self)
3580{
3581 if (self->pickler->memo)
3582 PyMemoTable_Clear(self->pickler->memo);
3583 Py_RETURN_NONE;
3584}
3585
3586PyDoc_STRVAR(pmp_copy_doc,
3587"memo.copy() -> new_memo. Copy the memo to a new object.");
3588
3589static PyObject *
3590pmp_copy(PicklerMemoProxyObject *self)
3591{
3592 Py_ssize_t i;
3593 PyMemoTable *memo;
3594 PyObject *new_memo = PyDict_New();
3595 if (new_memo == NULL)
3596 return NULL;
3597
3598 memo = self->pickler->memo;
3599 for (i = 0; i < memo->mt_allocated; ++i) {
3600 PyMemoEntry entry = memo->mt_table[i];
3601 if (entry.me_key != NULL) {
3602 int status;
3603 PyObject *key, *value;
3604
3605 key = PyLong_FromVoidPtr(entry.me_key);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003606 value = Py_BuildValue("nO", entry.me_value, entry.me_key);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003607
3608 if (key == NULL || value == NULL) {
3609 Py_XDECREF(key);
3610 Py_XDECREF(value);
3611 goto error;
3612 }
3613 status = PyDict_SetItem(new_memo, key, value);
3614 Py_DECREF(key);
3615 Py_DECREF(value);
3616 if (status < 0)
3617 goto error;
3618 }
3619 }
3620 return new_memo;
3621
3622 error:
3623 Py_XDECREF(new_memo);
3624 return NULL;
3625}
3626
3627PyDoc_STRVAR(pmp_reduce_doc,
3628"memo.__reduce__(). Pickling support.");
3629
3630static PyObject *
3631pmp_reduce(PicklerMemoProxyObject *self, PyObject *args)
3632{
3633 PyObject *reduce_value, *dict_args;
3634 PyObject *contents = pmp_copy(self);
3635 if (contents == NULL)
3636 return NULL;
3637
3638 reduce_value = PyTuple_New(2);
3639 if (reduce_value == NULL) {
3640 Py_DECREF(contents);
3641 return NULL;
3642 }
3643 dict_args = PyTuple_New(1);
3644 if (dict_args == NULL) {
3645 Py_DECREF(contents);
3646 Py_DECREF(reduce_value);
3647 return NULL;
3648 }
3649 PyTuple_SET_ITEM(dict_args, 0, contents);
3650 Py_INCREF((PyObject *)&PyDict_Type);
3651 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
3652 PyTuple_SET_ITEM(reduce_value, 1, dict_args);
3653 return reduce_value;
3654}
3655
3656static PyMethodDef picklerproxy_methods[] = {
3657 {"clear", (PyCFunction)pmp_clear, METH_NOARGS, pmp_clear_doc},
3658 {"copy", (PyCFunction)pmp_copy, METH_NOARGS, pmp_copy_doc},
3659 {"__reduce__", (PyCFunction)pmp_reduce, METH_VARARGS, pmp_reduce_doc},
3660 {NULL, NULL} /* sentinel */
3661};
3662
3663static void
3664PicklerMemoProxy_dealloc(PicklerMemoProxyObject *self)
3665{
3666 PyObject_GC_UnTrack(self);
3667 Py_XDECREF(self->pickler);
3668 PyObject_GC_Del((PyObject *)self);
3669}
3670
3671static int
3672PicklerMemoProxy_traverse(PicklerMemoProxyObject *self,
3673 visitproc visit, void *arg)
3674{
3675 Py_VISIT(self->pickler);
3676 return 0;
3677}
3678
3679static int
3680PicklerMemoProxy_clear(PicklerMemoProxyObject *self)
3681{
3682 Py_CLEAR(self->pickler);
3683 return 0;
3684}
3685
3686static PyTypeObject PicklerMemoProxyType = {
3687 PyVarObject_HEAD_INIT(NULL, 0)
3688 "_pickle.PicklerMemoProxy", /*tp_name*/
3689 sizeof(PicklerMemoProxyObject), /*tp_basicsize*/
3690 0,
3691 (destructor)PicklerMemoProxy_dealloc, /* tp_dealloc */
3692 0, /* tp_print */
3693 0, /* tp_getattr */
3694 0, /* tp_setattr */
3695 0, /* tp_compare */
3696 0, /* tp_repr */
3697 0, /* tp_as_number */
3698 0, /* tp_as_sequence */
3699 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00003700 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003701 0, /* tp_call */
3702 0, /* tp_str */
3703 PyObject_GenericGetAttr, /* tp_getattro */
3704 PyObject_GenericSetAttr, /* tp_setattro */
3705 0, /* tp_as_buffer */
3706 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3707 0, /* tp_doc */
3708 (traverseproc)PicklerMemoProxy_traverse, /* tp_traverse */
3709 (inquiry)PicklerMemoProxy_clear, /* tp_clear */
3710 0, /* tp_richcompare */
3711 0, /* tp_weaklistoffset */
3712 0, /* tp_iter */
3713 0, /* tp_iternext */
3714 picklerproxy_methods, /* tp_methods */
3715};
3716
3717static PyObject *
3718PicklerMemoProxy_New(PicklerObject *pickler)
3719{
3720 PicklerMemoProxyObject *self;
3721
3722 self = PyObject_GC_New(PicklerMemoProxyObject, &PicklerMemoProxyType);
3723 if (self == NULL)
3724 return NULL;
3725 Py_INCREF(pickler);
3726 self->pickler = pickler;
3727 PyObject_GC_Track(self);
3728 return (PyObject *)self;
3729}
3730
3731/*****************************************************************************/
3732
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003733static PyObject *
3734Pickler_get_memo(PicklerObject *self)
3735{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003736 return PicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003737}
3738
3739static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003740Pickler_set_memo(PicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003741{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003742 PyMemoTable *new_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003743
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003744 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003745 PyErr_SetString(PyExc_TypeError,
3746 "attribute deletion is not supported");
3747 return -1;
3748 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003749
3750 if (Py_TYPE(obj) == &PicklerMemoProxyType) {
3751 PicklerObject *pickler =
3752 ((PicklerMemoProxyObject *)obj)->pickler;
3753
3754 new_memo = PyMemoTable_Copy(pickler->memo);
3755 if (new_memo == NULL)
3756 return -1;
3757 }
3758 else if (PyDict_Check(obj)) {
3759 Py_ssize_t i = 0;
3760 PyObject *key, *value;
3761
3762 new_memo = PyMemoTable_New();
3763 if (new_memo == NULL)
3764 return -1;
3765
3766 while (PyDict_Next(obj, &i, &key, &value)) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003767 Py_ssize_t memo_id;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003768 PyObject *memo_obj;
3769
3770 if (!PyTuple_Check(value) || Py_SIZE(value) != 2) {
3771 PyErr_SetString(PyExc_TypeError,
3772 "'memo' values must be 2-item tuples");
3773 goto error;
3774 }
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003775 memo_id = PyLong_AsSsize_t(PyTuple_GET_ITEM(value, 0));
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003776 if (memo_id == -1 && PyErr_Occurred())
3777 goto error;
3778 memo_obj = PyTuple_GET_ITEM(value, 1);
3779 if (PyMemoTable_Set(new_memo, memo_obj, memo_id) < 0)
3780 goto error;
3781 }
3782 }
3783 else {
3784 PyErr_Format(PyExc_TypeError,
3785 "'memo' attribute must be an PicklerMemoProxy object"
3786 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003787 return -1;
3788 }
3789
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003790 PyMemoTable_Del(self->memo);
3791 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003792
3793 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003794
3795 error:
3796 if (new_memo)
3797 PyMemoTable_Del(new_memo);
3798 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003799}
3800
3801static PyObject *
3802Pickler_get_persid(PicklerObject *self)
3803{
3804 if (self->pers_func == NULL)
3805 PyErr_SetString(PyExc_AttributeError, "persistent_id");
3806 else
3807 Py_INCREF(self->pers_func);
3808 return self->pers_func;
3809}
3810
3811static int
3812Pickler_set_persid(PicklerObject *self, PyObject *value)
3813{
3814 PyObject *tmp;
3815
3816 if (value == NULL) {
3817 PyErr_SetString(PyExc_TypeError,
3818 "attribute deletion is not supported");
3819 return -1;
3820 }
3821 if (!PyCallable_Check(value)) {
3822 PyErr_SetString(PyExc_TypeError,
3823 "persistent_id must be a callable taking one argument");
3824 return -1;
3825 }
3826
3827 tmp = self->pers_func;
3828 Py_INCREF(value);
3829 self->pers_func = value;
3830 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
3831
3832 return 0;
3833}
3834
3835static PyMemberDef Pickler_members[] = {
3836 {"bin", T_INT, offsetof(PicklerObject, bin)},
3837 {"fast", T_INT, offsetof(PicklerObject, fast)},
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003838 {"dispatch_table", T_OBJECT_EX, offsetof(PicklerObject, dispatch_table)},
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003839 {NULL}
3840};
3841
3842static PyGetSetDef Pickler_getsets[] = {
3843 {"memo", (getter)Pickler_get_memo,
3844 (setter)Pickler_set_memo},
3845 {"persistent_id", (getter)Pickler_get_persid,
3846 (setter)Pickler_set_persid},
3847 {NULL}
3848};
3849
3850static PyTypeObject Pickler_Type = {
3851 PyVarObject_HEAD_INIT(NULL, 0)
3852 "_pickle.Pickler" , /*tp_name*/
3853 sizeof(PicklerObject), /*tp_basicsize*/
3854 0, /*tp_itemsize*/
3855 (destructor)Pickler_dealloc, /*tp_dealloc*/
3856 0, /*tp_print*/
3857 0, /*tp_getattr*/
3858 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00003859 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003860 0, /*tp_repr*/
3861 0, /*tp_as_number*/
3862 0, /*tp_as_sequence*/
3863 0, /*tp_as_mapping*/
3864 0, /*tp_hash*/
3865 0, /*tp_call*/
3866 0, /*tp_str*/
3867 0, /*tp_getattro*/
3868 0, /*tp_setattro*/
3869 0, /*tp_as_buffer*/
3870 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3871 Pickler_doc, /*tp_doc*/
3872 (traverseproc)Pickler_traverse, /*tp_traverse*/
3873 (inquiry)Pickler_clear, /*tp_clear*/
3874 0, /*tp_richcompare*/
3875 0, /*tp_weaklistoffset*/
3876 0, /*tp_iter*/
3877 0, /*tp_iternext*/
3878 Pickler_methods, /*tp_methods*/
3879 Pickler_members, /*tp_members*/
3880 Pickler_getsets, /*tp_getset*/
3881 0, /*tp_base*/
3882 0, /*tp_dict*/
3883 0, /*tp_descr_get*/
3884 0, /*tp_descr_set*/
3885 0, /*tp_dictoffset*/
3886 (initproc)Pickler_init, /*tp_init*/
3887 PyType_GenericAlloc, /*tp_alloc*/
3888 PyType_GenericNew, /*tp_new*/
3889 PyObject_GC_Del, /*tp_free*/
3890 0, /*tp_is_gc*/
3891};
3892
Victor Stinner121aab42011-09-29 23:40:53 +02003893/* Temporary helper for calling self.find_class().
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003894
3895 XXX: It would be nice to able to avoid Python function call overhead, by
3896 using directly the C version of find_class(), when find_class() is not
3897 overridden by a subclass. Although, this could become rather hackish. A
3898 simpler optimization would be to call the C function when self is not a
3899 subclass instance. */
3900static PyObject *
3901find_class(UnpicklerObject *self, PyObject *module_name, PyObject *global_name)
3902{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02003903 _Py_IDENTIFIER(find_class);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02003904
3905 return _PyObject_CallMethodId((PyObject *)self, &PyId_find_class, "OO",
3906 module_name, global_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003907}
3908
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003909static Py_ssize_t
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003910marker(UnpicklerObject *self)
3911{
3912 if (self->num_marks < 1) {
3913 PyErr_SetString(UnpicklingError, "could not find MARK");
3914 return -1;
3915 }
3916
3917 return self->marks[--self->num_marks];
3918}
3919
3920static int
3921load_none(UnpicklerObject *self)
3922{
3923 PDATA_APPEND(self->stack, Py_None, -1);
3924 return 0;
3925}
3926
3927static int
3928bad_readline(void)
3929{
3930 PyErr_SetString(UnpicklingError, "pickle data was truncated");
3931 return -1;
3932}
3933
3934static int
3935load_int(UnpicklerObject *self)
3936{
3937 PyObject *value;
3938 char *endptr, *s;
3939 Py_ssize_t len;
3940 long x;
3941
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003942 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003943 return -1;
3944 if (len < 2)
3945 return bad_readline();
3946
3947 errno = 0;
Victor Stinner121aab42011-09-29 23:40:53 +02003948 /* XXX: Should the base argument of strtol() be explicitly set to 10?
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003949 XXX(avassalotti): Should this uses PyOS_strtol()? */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003950 x = strtol(s, &endptr, 0);
3951
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003952 if (errno || (*endptr != '\n' && *endptr != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003953 /* Hm, maybe we've got something long. Let's try reading
3954 * it as a Python long object. */
3955 errno = 0;
3956 /* XXX: Same thing about the base here. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003957 value = PyLong_FromString(s, NULL, 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003958 if (value == NULL) {
3959 PyErr_SetString(PyExc_ValueError,
3960 "could not convert string to int");
3961 return -1;
3962 }
3963 }
3964 else {
3965 if (len == 3 && (x == 0 || x == 1)) {
3966 if ((value = PyBool_FromLong(x)) == NULL)
3967 return -1;
3968 }
3969 else {
3970 if ((value = PyLong_FromLong(x)) == NULL)
3971 return -1;
3972 }
3973 }
3974
3975 PDATA_PUSH(self->stack, value, -1);
3976 return 0;
3977}
3978
3979static int
3980load_bool(UnpicklerObject *self, PyObject *boolean)
3981{
3982 assert(boolean == Py_True || boolean == Py_False);
3983 PDATA_APPEND(self->stack, boolean, -1);
3984 return 0;
3985}
3986
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003987/* s contains x bytes of an unsigned little-endian integer. Return its value
3988 * as a C Py_ssize_t, or -1 if it's higher than PY_SSIZE_T_MAX.
3989 */
3990static Py_ssize_t
3991calc_binsize(char *bytes, int size)
3992{
3993 unsigned char *s = (unsigned char *)bytes;
3994 size_t x = 0;
3995
3996 assert(size == 4);
3997
3998 x = (size_t) s[0];
3999 x |= (size_t) s[1] << 8;
4000 x |= (size_t) s[2] << 16;
4001 x |= (size_t) s[3] << 24;
4002
4003 if (x > PY_SSIZE_T_MAX)
4004 return -1;
4005 else
4006 return (Py_ssize_t) x;
4007}
4008
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004009/* s contains x bytes of a little-endian integer. Return its value as a
4010 * C int. Obscure: when x is 1 or 2, this is an unsigned little-endian
4011 * int, but when x is 4 it's a signed one. This is an historical source
4012 * of x-platform bugs.
4013 */
4014static long
4015calc_binint(char *bytes, int size)
4016{
4017 unsigned char *s = (unsigned char *)bytes;
4018 int i = size;
4019 long x = 0;
4020
4021 for (i = 0; i < size; i++) {
4022 x |= (long)s[i] << (i * 8);
4023 }
4024
4025 /* Unlike BININT1 and BININT2, BININT (more accurately BININT4)
4026 * is signed, so on a box with longs bigger than 4 bytes we need
4027 * to extend a BININT's sign bit to the full width.
4028 */
4029 if (SIZEOF_LONG > 4 && size == 4) {
4030 x |= -(x & (1L << 31));
4031 }
4032
4033 return x;
4034}
4035
4036static int
4037load_binintx(UnpicklerObject *self, char *s, int size)
4038{
4039 PyObject *value;
4040 long x;
4041
4042 x = calc_binint(s, size);
4043
4044 if ((value = PyLong_FromLong(x)) == NULL)
4045 return -1;
4046
4047 PDATA_PUSH(self->stack, value, -1);
4048 return 0;
4049}
4050
4051static int
4052load_binint(UnpicklerObject *self)
4053{
4054 char *s;
4055
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004056 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004057 return -1;
4058
4059 return load_binintx(self, s, 4);
4060}
4061
4062static int
4063load_binint1(UnpicklerObject *self)
4064{
4065 char *s;
4066
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004067 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004068 return -1;
4069
4070 return load_binintx(self, s, 1);
4071}
4072
4073static int
4074load_binint2(UnpicklerObject *self)
4075{
4076 char *s;
4077
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004078 if (_Unpickler_Read(self, &s, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004079 return -1;
4080
4081 return load_binintx(self, s, 2);
4082}
4083
4084static int
4085load_long(UnpicklerObject *self)
4086{
4087 PyObject *value;
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004088 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004089 Py_ssize_t len;
4090
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004091 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004092 return -1;
4093 if (len < 2)
4094 return bad_readline();
4095
Mark Dickinson8dd05142009-01-20 20:43:58 +00004096 /* s[len-2] will usually be 'L' (and s[len-1] is '\n'); we need to remove
4097 the 'L' before calling PyLong_FromString. In order to maintain
4098 compatibility with Python 3.0.0, we don't actually *require*
4099 the 'L' to be present. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004100 if (s[len-2] == 'L')
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004101 s[len-2] = '\0';
Alexandre Vassalottie4bccb72009-01-24 01:47:57 +00004102 /* XXX: Should the base argument explicitly set to 10? */
4103 value = PyLong_FromString(s, NULL, 0);
Mark Dickinson8dd05142009-01-20 20:43:58 +00004104 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004105 return -1;
4106
4107 PDATA_PUSH(self->stack, value, -1);
4108 return 0;
4109}
4110
4111/* 'size' bytes contain the # of bytes of little-endian 256's-complement
4112 * data following.
4113 */
4114static int
4115load_counted_long(UnpicklerObject *self, int size)
4116{
4117 PyObject *value;
4118 char *nbytes;
4119 char *pdata;
4120
4121 assert(size == 1 || size == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004122 if (_Unpickler_Read(self, &nbytes, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004123 return -1;
4124
4125 size = calc_binint(nbytes, size);
4126 if (size < 0) {
4127 /* Corrupt or hostile pickle -- we never write one like this */
4128 PyErr_SetString(UnpicklingError,
4129 "LONG pickle has negative byte count");
4130 return -1;
4131 }
4132
4133 if (size == 0)
4134 value = PyLong_FromLong(0L);
4135 else {
4136 /* Read the raw little-endian bytes and convert. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004137 if (_Unpickler_Read(self, &pdata, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004138 return -1;
4139 value = _PyLong_FromByteArray((unsigned char *)pdata, (size_t)size,
4140 1 /* little endian */ , 1 /* signed */ );
4141 }
4142 if (value == NULL)
4143 return -1;
4144 PDATA_PUSH(self->stack, value, -1);
4145 return 0;
4146}
4147
4148static int
4149load_float(UnpicklerObject *self)
4150{
4151 PyObject *value;
4152 char *endptr, *s;
4153 Py_ssize_t len;
4154 double d;
4155
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004156 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004157 return -1;
4158 if (len < 2)
4159 return bad_readline();
4160
4161 errno = 0;
Mark Dickinson725bfd82009-05-03 20:33:40 +00004162 d = PyOS_string_to_double(s, &endptr, PyExc_OverflowError);
4163 if (d == -1.0 && PyErr_Occurred())
4164 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004165 if ((endptr[0] != '\n') && (endptr[0] != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004166 PyErr_SetString(PyExc_ValueError, "could not convert string to float");
4167 return -1;
4168 }
Mark Dickinson725bfd82009-05-03 20:33:40 +00004169 value = PyFloat_FromDouble(d);
4170 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004171 return -1;
4172
4173 PDATA_PUSH(self->stack, value, -1);
4174 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004175}
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004176
4177static int
4178load_binfloat(UnpicklerObject *self)
4179{
4180 PyObject *value;
4181 double x;
4182 char *s;
4183
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004184 if (_Unpickler_Read(self, &s, 8) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004185 return -1;
4186
4187 x = _PyFloat_Unpack8((unsigned char *)s, 0);
4188 if (x == -1.0 && PyErr_Occurred())
4189 return -1;
4190
4191 if ((value = PyFloat_FromDouble(x)) == NULL)
4192 return -1;
4193
4194 PDATA_PUSH(self->stack, value, -1);
4195 return 0;
4196}
4197
4198static int
4199load_string(UnpicklerObject *self)
4200{
4201 PyObject *bytes;
4202 PyObject *str = NULL;
4203 Py_ssize_t len;
4204 char *s, *p;
4205
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004206 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004207 return -1;
4208 if (len < 3)
4209 return bad_readline();
4210 if ((s = strdup(s)) == NULL) {
4211 PyErr_NoMemory();
4212 return -1;
4213 }
4214
4215 /* Strip outermost quotes */
4216 while (s[len - 1] <= ' ')
4217 len--;
4218 if (s[0] == '"' && s[len - 1] == '"') {
4219 s[len - 1] = '\0';
4220 p = s + 1;
4221 len -= 2;
4222 }
4223 else if (s[0] == '\'' && s[len - 1] == '\'') {
4224 s[len - 1] = '\0';
4225 p = s + 1;
4226 len -= 2;
4227 }
4228 else {
4229 free(s);
4230 PyErr_SetString(PyExc_ValueError, "insecure string pickle");
4231 return -1;
4232 }
4233
4234 /* Use the PyBytes API to decode the string, since that is what is used
4235 to encode, and then coerce the result to Unicode. */
4236 bytes = PyBytes_DecodeEscape(p, len, NULL, 0, NULL);
4237 free(s);
4238 if (bytes == NULL)
4239 return -1;
4240 str = PyUnicode_FromEncodedObject(bytes, self->encoding, self->errors);
4241 Py_DECREF(bytes);
4242 if (str == NULL)
4243 return -1;
4244
4245 PDATA_PUSH(self->stack, str, -1);
4246 return 0;
4247}
4248
4249static int
4250load_binbytes(UnpicklerObject *self)
4251{
4252 PyObject *bytes;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004253 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004254 char *s;
4255
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004256 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004257 return -1;
4258
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004259 x = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004260 if (x < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004261 PyErr_Format(PyExc_OverflowError,
4262 "BINBYTES exceeds system's maximum size of %zd bytes",
4263 PY_SSIZE_T_MAX
4264 );
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004265 return -1;
4266 }
4267
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004268 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004269 return -1;
4270 bytes = PyBytes_FromStringAndSize(s, x);
4271 if (bytes == NULL)
4272 return -1;
4273
4274 PDATA_PUSH(self->stack, bytes, -1);
4275 return 0;
4276}
4277
4278static int
4279load_short_binbytes(UnpicklerObject *self)
4280{
4281 PyObject *bytes;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004282 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004283 char *s;
4284
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004285 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004286 return -1;
4287
4288 x = (unsigned char)s[0];
4289
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004290 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004291 return -1;
4292
4293 bytes = PyBytes_FromStringAndSize(s, x);
4294 if (bytes == NULL)
4295 return -1;
4296
4297 PDATA_PUSH(self->stack, bytes, -1);
4298 return 0;
4299}
4300
4301static int
4302load_binstring(UnpicklerObject *self)
4303{
4304 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004305 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004306 char *s;
4307
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004308 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004309 return -1;
4310
4311 x = calc_binint(s, 4);
4312 if (x < 0) {
Victor Stinner121aab42011-09-29 23:40:53 +02004313 PyErr_SetString(UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004314 "BINSTRING pickle has negative byte count");
4315 return -1;
4316 }
4317
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004318 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004319 return -1;
4320
4321 /* Convert Python 2.x strings to unicode. */
4322 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4323 if (str == NULL)
4324 return -1;
4325
4326 PDATA_PUSH(self->stack, str, -1);
4327 return 0;
4328}
4329
4330static int
4331load_short_binstring(UnpicklerObject *self)
4332{
4333 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004334 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004335 char *s;
4336
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004337 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004338 return -1;
4339
4340 x = (unsigned char)s[0];
4341
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004342 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004343 return -1;
4344
4345 /* Convert Python 2.x strings to unicode. */
4346 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4347 if (str == NULL)
4348 return -1;
4349
4350 PDATA_PUSH(self->stack, str, -1);
4351 return 0;
4352}
4353
4354static int
4355load_unicode(UnpicklerObject *self)
4356{
4357 PyObject *str;
4358 Py_ssize_t len;
4359 char *s;
4360
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004361 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004362 return -1;
4363 if (len < 1)
4364 return bad_readline();
4365
4366 str = PyUnicode_DecodeRawUnicodeEscape(s, len - 1, NULL);
4367 if (str == NULL)
4368 return -1;
4369
4370 PDATA_PUSH(self->stack, str, -1);
4371 return 0;
4372}
4373
4374static int
4375load_binunicode(UnpicklerObject *self)
4376{
4377 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004378 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004379 char *s;
4380
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004381 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004382 return -1;
4383
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004384 size = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004385 if (size < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004386 PyErr_Format(PyExc_OverflowError,
4387 "BINUNICODE exceeds system's maximum size of %zd bytes",
4388 PY_SSIZE_T_MAX
4389 );
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004390 return -1;
4391 }
4392
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004393
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004394 if (_Unpickler_Read(self, &s, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004395 return -1;
4396
Victor Stinner485fb562010-04-13 11:07:24 +00004397 str = PyUnicode_DecodeUTF8(s, size, "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004398 if (str == NULL)
4399 return -1;
4400
4401 PDATA_PUSH(self->stack, str, -1);
4402 return 0;
4403}
4404
4405static int
4406load_tuple(UnpicklerObject *self)
4407{
4408 PyObject *tuple;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004409 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004410
4411 if ((i = marker(self)) < 0)
4412 return -1;
4413
4414 tuple = Pdata_poptuple(self->stack, i);
4415 if (tuple == NULL)
4416 return -1;
4417 PDATA_PUSH(self->stack, tuple, -1);
4418 return 0;
4419}
4420
4421static int
4422load_counted_tuple(UnpicklerObject *self, int len)
4423{
4424 PyObject *tuple;
4425
4426 tuple = PyTuple_New(len);
4427 if (tuple == NULL)
4428 return -1;
4429
4430 while (--len >= 0) {
4431 PyObject *item;
4432
4433 PDATA_POP(self->stack, item);
4434 if (item == NULL)
4435 return -1;
4436 PyTuple_SET_ITEM(tuple, len, item);
4437 }
4438 PDATA_PUSH(self->stack, tuple, -1);
4439 return 0;
4440}
4441
4442static int
4443load_empty_list(UnpicklerObject *self)
4444{
4445 PyObject *list;
4446
4447 if ((list = PyList_New(0)) == NULL)
4448 return -1;
4449 PDATA_PUSH(self->stack, list, -1);
4450 return 0;
4451}
4452
4453static int
4454load_empty_dict(UnpicklerObject *self)
4455{
4456 PyObject *dict;
4457
4458 if ((dict = PyDict_New()) == NULL)
4459 return -1;
4460 PDATA_PUSH(self->stack, dict, -1);
4461 return 0;
4462}
4463
4464static int
4465load_list(UnpicklerObject *self)
4466{
4467 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004468 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004469
4470 if ((i = marker(self)) < 0)
4471 return -1;
4472
4473 list = Pdata_poplist(self->stack, i);
4474 if (list == NULL)
4475 return -1;
4476 PDATA_PUSH(self->stack, list, -1);
4477 return 0;
4478}
4479
4480static int
4481load_dict(UnpicklerObject *self)
4482{
4483 PyObject *dict, *key, *value;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004484 Py_ssize_t i, j, k;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004485
4486 if ((i = marker(self)) < 0)
4487 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004488 j = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004489
4490 if ((dict = PyDict_New()) == NULL)
4491 return -1;
4492
4493 for (k = i + 1; k < j; k += 2) {
4494 key = self->stack->data[k - 1];
4495 value = self->stack->data[k];
4496 if (PyDict_SetItem(dict, key, value) < 0) {
4497 Py_DECREF(dict);
4498 return -1;
4499 }
4500 }
4501 Pdata_clear(self->stack, i);
4502 PDATA_PUSH(self->stack, dict, -1);
4503 return 0;
4504}
4505
4506static PyObject *
4507instantiate(PyObject *cls, PyObject *args)
4508{
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004509 PyObject *result = NULL;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02004510 _Py_IDENTIFIER(__getinitargs__);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004511 /* Caller must assure args are a tuple. Normally, args come from
4512 Pdata_poptuple which packs objects from the top of the stack
4513 into a newly created tuple. */
4514 assert(PyTuple_Check(args));
4515 if (Py_SIZE(args) > 0 || !PyType_Check(cls) ||
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02004516 _PyObject_HasAttrId(cls, &PyId___getinitargs__)) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004517 result = PyObject_CallObject(cls, args);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004518 }
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004519 else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02004520 _Py_IDENTIFIER(__new__);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02004521
4522 result = _PyObject_CallMethodId(cls, &PyId___new__, "O", cls);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004523 }
4524 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004525}
4526
4527static int
4528load_obj(UnpicklerObject *self)
4529{
4530 PyObject *cls, *args, *obj = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004531 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004532
4533 if ((i = marker(self)) < 0)
4534 return -1;
4535
4536 args = Pdata_poptuple(self->stack, i + 1);
4537 if (args == NULL)
4538 return -1;
4539
4540 PDATA_POP(self->stack, cls);
4541 if (cls) {
4542 obj = instantiate(cls, args);
4543 Py_DECREF(cls);
4544 }
4545 Py_DECREF(args);
4546 if (obj == NULL)
4547 return -1;
4548
4549 PDATA_PUSH(self->stack, obj, -1);
4550 return 0;
4551}
4552
4553static int
4554load_inst(UnpicklerObject *self)
4555{
4556 PyObject *cls = NULL;
4557 PyObject *args = NULL;
4558 PyObject *obj = NULL;
4559 PyObject *module_name;
4560 PyObject *class_name;
4561 Py_ssize_t len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004562 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004563 char *s;
4564
4565 if ((i = marker(self)) < 0)
4566 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004567 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004568 return -1;
4569 if (len < 2)
4570 return bad_readline();
4571
4572 /* Here it is safe to use PyUnicode_DecodeASCII(), even though non-ASCII
4573 identifiers are permitted in Python 3.0, since the INST opcode is only
4574 supported by older protocols on Python 2.x. */
4575 module_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
4576 if (module_name == NULL)
4577 return -1;
4578
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004579 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004580 if (len < 2)
4581 return bad_readline();
4582 class_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004583 if (class_name != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004584 cls = find_class(self, module_name, class_name);
4585 Py_DECREF(class_name);
4586 }
4587 }
4588 Py_DECREF(module_name);
4589
4590 if (cls == NULL)
4591 return -1;
4592
4593 if ((args = Pdata_poptuple(self->stack, i)) != NULL) {
4594 obj = instantiate(cls, args);
4595 Py_DECREF(args);
4596 }
4597 Py_DECREF(cls);
4598
4599 if (obj == NULL)
4600 return -1;
4601
4602 PDATA_PUSH(self->stack, obj, -1);
4603 return 0;
4604}
4605
4606static int
4607load_newobj(UnpicklerObject *self)
4608{
4609 PyObject *args = NULL;
4610 PyObject *clsraw = NULL;
4611 PyTypeObject *cls; /* clsraw cast to its true type */
4612 PyObject *obj;
4613
4614 /* Stack is ... cls argtuple, and we want to call
4615 * cls.__new__(cls, *argtuple).
4616 */
4617 PDATA_POP(self->stack, args);
4618 if (args == NULL)
4619 goto error;
4620 if (!PyTuple_Check(args)) {
4621 PyErr_SetString(UnpicklingError, "NEWOBJ expected an arg " "tuple.");
4622 goto error;
4623 }
4624
4625 PDATA_POP(self->stack, clsraw);
4626 cls = (PyTypeObject *)clsraw;
4627 if (cls == NULL)
4628 goto error;
4629 if (!PyType_Check(cls)) {
4630 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4631 "isn't a type object");
4632 goto error;
4633 }
4634 if (cls->tp_new == NULL) {
4635 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4636 "has NULL tp_new");
4637 goto error;
4638 }
4639
4640 /* Call __new__. */
4641 obj = cls->tp_new(cls, args, NULL);
4642 if (obj == NULL)
4643 goto error;
4644
4645 Py_DECREF(args);
4646 Py_DECREF(clsraw);
4647 PDATA_PUSH(self->stack, obj, -1);
4648 return 0;
4649
4650 error:
4651 Py_XDECREF(args);
4652 Py_XDECREF(clsraw);
4653 return -1;
4654}
4655
4656static int
4657load_global(UnpicklerObject *self)
4658{
4659 PyObject *global = NULL;
4660 PyObject *module_name;
4661 PyObject *global_name;
4662 Py_ssize_t len;
4663 char *s;
4664
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004665 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004666 return -1;
4667 if (len < 2)
4668 return bad_readline();
4669 module_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4670 if (!module_name)
4671 return -1;
4672
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004673 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004674 if (len < 2) {
4675 Py_DECREF(module_name);
4676 return bad_readline();
4677 }
4678 global_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4679 if (global_name) {
4680 global = find_class(self, module_name, global_name);
4681 Py_DECREF(global_name);
4682 }
4683 }
4684 Py_DECREF(module_name);
4685
4686 if (global == NULL)
4687 return -1;
4688 PDATA_PUSH(self->stack, global, -1);
4689 return 0;
4690}
4691
4692static int
4693load_persid(UnpicklerObject *self)
4694{
4695 PyObject *pid;
4696 Py_ssize_t len;
4697 char *s;
4698
4699 if (self->pers_func) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004700 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004701 return -1;
4702 if (len < 2)
4703 return bad_readline();
4704
4705 pid = PyBytes_FromStringAndSize(s, len - 1);
4706 if (pid == NULL)
4707 return -1;
4708
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004709 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004710 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004711 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004712 if (pid == NULL)
4713 return -1;
4714
4715 PDATA_PUSH(self->stack, pid, -1);
4716 return 0;
4717 }
4718 else {
4719 PyErr_SetString(UnpicklingError,
4720 "A load persistent id instruction was encountered,\n"
4721 "but no persistent_load function was specified.");
4722 return -1;
4723 }
4724}
4725
4726static int
4727load_binpersid(UnpicklerObject *self)
4728{
4729 PyObject *pid;
4730
4731 if (self->pers_func) {
4732 PDATA_POP(self->stack, pid);
4733 if (pid == NULL)
4734 return -1;
4735
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004736 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004737 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004738 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004739 if (pid == NULL)
4740 return -1;
4741
4742 PDATA_PUSH(self->stack, pid, -1);
4743 return 0;
4744 }
4745 else {
4746 PyErr_SetString(UnpicklingError,
4747 "A load persistent id instruction was encountered,\n"
4748 "but no persistent_load function was specified.");
4749 return -1;
4750 }
4751}
4752
4753static int
4754load_pop(UnpicklerObject *self)
4755{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004756 Py_ssize_t len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004757
4758 /* Note that we split the (pickle.py) stack into two stacks,
4759 * an object stack and a mark stack. We have to be clever and
4760 * pop the right one. We do this by looking at the top of the
Collin Winter8ca69de2009-05-26 16:53:41 +00004761 * mark stack first, and only signalling a stack underflow if
4762 * the object stack is empty and the mark stack doesn't match
4763 * our expectations.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004764 */
Collin Winter8ca69de2009-05-26 16:53:41 +00004765 if (self->num_marks > 0 && self->marks[self->num_marks - 1] == len) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004766 self->num_marks--;
Antoine Pitrou01a15ea2010-01-07 17:57:31 +00004767 } else if (len > 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004768 len--;
4769 Py_DECREF(self->stack->data[len]);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004770 Py_SIZE(self->stack) = len;
Collin Winter8ca69de2009-05-26 16:53:41 +00004771 } else {
4772 return stack_underflow();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004773 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004774 return 0;
4775}
4776
4777static int
4778load_pop_mark(UnpicklerObject *self)
4779{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004780 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004781
4782 if ((i = marker(self)) < 0)
4783 return -1;
4784
4785 Pdata_clear(self->stack, i);
4786
4787 return 0;
4788}
4789
4790static int
4791load_dup(UnpicklerObject *self)
4792{
4793 PyObject *last;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004794 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004795
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004796 if ((len = Py_SIZE(self->stack)) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004797 return stack_underflow();
4798 last = self->stack->data[len - 1];
4799 PDATA_APPEND(self->stack, last, -1);
4800 return 0;
4801}
4802
4803static int
4804load_get(UnpicklerObject *self)
4805{
4806 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004807 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004808 Py_ssize_t len;
4809 char *s;
4810
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004811 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004812 return -1;
4813 if (len < 2)
4814 return bad_readline();
4815
4816 key = PyLong_FromString(s, NULL, 10);
4817 if (key == NULL)
4818 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004819 idx = PyLong_AsSsize_t(key);
4820 if (idx == -1 && PyErr_Occurred()) {
4821 Py_DECREF(key);
4822 return -1;
4823 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004824
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004825 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004826 if (value == NULL) {
4827 if (!PyErr_Occurred())
4828 PyErr_SetObject(PyExc_KeyError, key);
4829 Py_DECREF(key);
4830 return -1;
4831 }
4832 Py_DECREF(key);
4833
4834 PDATA_APPEND(self->stack, value, -1);
4835 return 0;
4836}
4837
4838static int
4839load_binget(UnpicklerObject *self)
4840{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004841 PyObject *value;
4842 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004843 char *s;
4844
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004845 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004846 return -1;
4847
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004848 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004849
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004850 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004851 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004852 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004853 if (!PyErr_Occurred())
4854 PyErr_SetObject(PyExc_KeyError, key);
4855 Py_DECREF(key);
4856 return -1;
4857 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004858
4859 PDATA_APPEND(self->stack, value, -1);
4860 return 0;
4861}
4862
4863static int
4864load_long_binget(UnpicklerObject *self)
4865{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004866 PyObject *value;
4867 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004868 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004869
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004870 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004871 return -1;
4872
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004873 idx = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004874
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004875 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004876 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004877 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004878 if (!PyErr_Occurred())
4879 PyErr_SetObject(PyExc_KeyError, key);
4880 Py_DECREF(key);
4881 return -1;
4882 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004883
4884 PDATA_APPEND(self->stack, value, -1);
4885 return 0;
4886}
4887
4888/* Push an object from the extension registry (EXT[124]). nbytes is
4889 * the number of bytes following the opcode, holding the index (code) value.
4890 */
4891static int
4892load_extension(UnpicklerObject *self, int nbytes)
4893{
4894 char *codebytes; /* the nbytes bytes after the opcode */
4895 long code; /* calc_binint returns long */
4896 PyObject *py_code; /* code as a Python int */
4897 PyObject *obj; /* the object to push */
4898 PyObject *pair; /* (module_name, class_name) */
4899 PyObject *module_name, *class_name;
4900
4901 assert(nbytes == 1 || nbytes == 2 || nbytes == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004902 if (_Unpickler_Read(self, &codebytes, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004903 return -1;
4904 code = calc_binint(codebytes, nbytes);
4905 if (code <= 0) { /* note that 0 is forbidden */
4906 /* Corrupt or hostile pickle. */
4907 PyErr_SetString(UnpicklingError, "EXT specifies code <= 0");
4908 return -1;
4909 }
4910
4911 /* Look for the code in the cache. */
4912 py_code = PyLong_FromLong(code);
4913 if (py_code == NULL)
4914 return -1;
4915 obj = PyDict_GetItem(extension_cache, py_code);
4916 if (obj != NULL) {
4917 /* Bingo. */
4918 Py_DECREF(py_code);
4919 PDATA_APPEND(self->stack, obj, -1);
4920 return 0;
4921 }
4922
4923 /* Look up the (module_name, class_name) pair. */
4924 pair = PyDict_GetItem(inverted_registry, py_code);
4925 if (pair == NULL) {
4926 Py_DECREF(py_code);
4927 PyErr_Format(PyExc_ValueError, "unregistered extension "
4928 "code %ld", code);
4929 return -1;
4930 }
4931 /* Since the extension registry is manipulable via Python code,
4932 * confirm that pair is really a 2-tuple of strings.
4933 */
4934 if (!PyTuple_Check(pair) || PyTuple_Size(pair) != 2 ||
4935 !PyUnicode_Check(module_name = PyTuple_GET_ITEM(pair, 0)) ||
4936 !PyUnicode_Check(class_name = PyTuple_GET_ITEM(pair, 1))) {
4937 Py_DECREF(py_code);
4938 PyErr_Format(PyExc_ValueError, "_inverted_registry[%ld] "
4939 "isn't a 2-tuple of strings", code);
4940 return -1;
4941 }
4942 /* Load the object. */
4943 obj = find_class(self, module_name, class_name);
4944 if (obj == NULL) {
4945 Py_DECREF(py_code);
4946 return -1;
4947 }
4948 /* Cache code -> obj. */
4949 code = PyDict_SetItem(extension_cache, py_code, obj);
4950 Py_DECREF(py_code);
4951 if (code < 0) {
4952 Py_DECREF(obj);
4953 return -1;
4954 }
4955 PDATA_PUSH(self->stack, obj, -1);
4956 return 0;
4957}
4958
4959static int
4960load_put(UnpicklerObject *self)
4961{
4962 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004963 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004964 Py_ssize_t len;
4965 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004966
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004967 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004968 return -1;
4969 if (len < 2)
4970 return bad_readline();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004971 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004972 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004973 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004974
4975 key = PyLong_FromString(s, NULL, 10);
4976 if (key == NULL)
4977 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004978 idx = PyLong_AsSsize_t(key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004979 Py_DECREF(key);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02004980 if (idx < 0) {
4981 if (!PyErr_Occurred())
4982 PyErr_SetString(PyExc_ValueError,
4983 "negative PUT argument");
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004984 return -1;
Antoine Pitrou55549ec2011-08-30 00:27:10 +02004985 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004986
4987 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004988}
4989
4990static int
4991load_binput(UnpicklerObject *self)
4992{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004993 PyObject *value;
4994 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004995 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004996
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004997 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004998 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004999
5000 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005001 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005002 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005003
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005004 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005005
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005006 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005007}
5008
5009static int
5010load_long_binput(UnpicklerObject *self)
5011{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005012 PyObject *value;
5013 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005014 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005015
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005016 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005017 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005018
5019 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005020 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005021 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005022
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005023 idx = calc_binsize(s, 4);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02005024 if (idx < 0) {
5025 PyErr_SetString(PyExc_ValueError,
5026 "negative LONG_BINPUT argument");
5027 return -1;
5028 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005029
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005030 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005031}
5032
5033static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005034do_append(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005035{
5036 PyObject *value;
5037 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005038 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005039
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005040 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005041 if (x > len || x <= 0)
5042 return stack_underflow();
5043 if (len == x) /* nothing to do */
5044 return 0;
5045
5046 list = self->stack->data[x - 1];
5047
5048 if (PyList_Check(list)) {
5049 PyObject *slice;
5050 Py_ssize_t list_len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005051 int ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005052
5053 slice = Pdata_poplist(self->stack, x);
5054 if (!slice)
5055 return -1;
5056 list_len = PyList_GET_SIZE(list);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005057 ret = PyList_SetSlice(list, list_len, list_len, slice);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005058 Py_DECREF(slice);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005059 return ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005060 }
5061 else {
5062 PyObject *append_func;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005063 _Py_IDENTIFIER(append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005064
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005065 append_func = _PyObject_GetAttrId(list, &PyId_append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005066 if (append_func == NULL)
5067 return -1;
5068 for (i = x; i < len; i++) {
5069 PyObject *result;
5070
5071 value = self->stack->data[i];
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005072 result = _Unpickler_FastCall(self, append_func, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005073 if (result == NULL) {
5074 Pdata_clear(self->stack, i + 1);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005075 Py_SIZE(self->stack) = x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005076 return -1;
5077 }
5078 Py_DECREF(result);
5079 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005080 Py_SIZE(self->stack) = x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005081 }
5082
5083 return 0;
5084}
5085
5086static int
5087load_append(UnpicklerObject *self)
5088{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005089 return do_append(self, Py_SIZE(self->stack) - 1);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005090}
5091
5092static int
5093load_appends(UnpicklerObject *self)
5094{
5095 return do_append(self, marker(self));
5096}
5097
5098static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005099do_setitems(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005100{
5101 PyObject *value, *key;
5102 PyObject *dict;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005103 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005104 int status = 0;
5105
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005106 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005107 if (x > len || x <= 0)
5108 return stack_underflow();
5109 if (len == x) /* nothing to do */
5110 return 0;
Victor Stinner121aab42011-09-29 23:40:53 +02005111 if ((len - x) % 2 != 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005112 /* Currupt or hostile pickle -- we never write one like this. */
5113 PyErr_SetString(UnpicklingError, "odd number of items for SETITEMS");
5114 return -1;
5115 }
5116
5117 /* Here, dict does not actually need to be a PyDict; it could be anything
5118 that supports the __setitem__ attribute. */
5119 dict = self->stack->data[x - 1];
5120
5121 for (i = x + 1; i < len; i += 2) {
5122 key = self->stack->data[i - 1];
5123 value = self->stack->data[i];
5124 if (PyObject_SetItem(dict, key, value) < 0) {
5125 status = -1;
5126 break;
5127 }
5128 }
5129
5130 Pdata_clear(self->stack, x);
5131 return status;
5132}
5133
5134static int
5135load_setitem(UnpicklerObject *self)
5136{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005137 return do_setitems(self, Py_SIZE(self->stack) - 2);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005138}
5139
5140static int
5141load_setitems(UnpicklerObject *self)
5142{
5143 return do_setitems(self, marker(self));
5144}
5145
5146static int
5147load_build(UnpicklerObject *self)
5148{
5149 PyObject *state, *inst, *slotstate;
5150 PyObject *setstate;
5151 int status = 0;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005152 _Py_IDENTIFIER(__setstate__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005153
5154 /* Stack is ... instance, state. We want to leave instance at
5155 * the stack top, possibly mutated via instance.__setstate__(state).
5156 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005157 if (Py_SIZE(self->stack) < 2)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005158 return stack_underflow();
5159
5160 PDATA_POP(self->stack, state);
5161 if (state == NULL)
5162 return -1;
5163
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005164 inst = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005165
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005166 setstate = _PyObject_GetAttrId(inst, &PyId___setstate__);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005167 if (setstate == NULL) {
5168 if (PyErr_ExceptionMatches(PyExc_AttributeError))
5169 PyErr_Clear();
Antoine Pitroud79dc622008-09-05 00:03:33 +00005170 else {
5171 Py_DECREF(state);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005172 return -1;
Antoine Pitroud79dc622008-09-05 00:03:33 +00005173 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005174 }
5175 else {
5176 PyObject *result;
5177
5178 /* The explicit __setstate__ is responsible for everything. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005179 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Antoine Pitroud79dc622008-09-05 00:03:33 +00005180 reference to state first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005181 result = _Unpickler_FastCall(self, setstate, state);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005182 Py_DECREF(setstate);
5183 if (result == NULL)
5184 return -1;
5185 Py_DECREF(result);
5186 return 0;
5187 }
5188
5189 /* A default __setstate__. First see whether state embeds a
5190 * slot state dict too (a proto 2 addition).
5191 */
5192 if (PyTuple_Check(state) && Py_SIZE(state) == 2) {
5193 PyObject *tmp = state;
5194
5195 state = PyTuple_GET_ITEM(tmp, 0);
5196 slotstate = PyTuple_GET_ITEM(tmp, 1);
5197 Py_INCREF(state);
5198 Py_INCREF(slotstate);
5199 Py_DECREF(tmp);
5200 }
5201 else
5202 slotstate = NULL;
5203
5204 /* Set inst.__dict__ from the state dict (if any). */
5205 if (state != Py_None) {
5206 PyObject *dict;
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005207 PyObject *d_key, *d_value;
5208 Py_ssize_t i;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005209 _Py_IDENTIFIER(__dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005210
5211 if (!PyDict_Check(state)) {
5212 PyErr_SetString(UnpicklingError, "state is not a dictionary");
5213 goto error;
5214 }
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005215 dict = _PyObject_GetAttrId(inst, &PyId___dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005216 if (dict == NULL)
5217 goto error;
5218
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005219 i = 0;
5220 while (PyDict_Next(state, &i, &d_key, &d_value)) {
5221 /* normally the keys for instance attributes are
5222 interned. we should try to do that here. */
5223 Py_INCREF(d_key);
5224 if (PyUnicode_CheckExact(d_key))
5225 PyUnicode_InternInPlace(&d_key);
5226 if (PyObject_SetItem(dict, d_key, d_value) < 0) {
5227 Py_DECREF(d_key);
5228 goto error;
5229 }
5230 Py_DECREF(d_key);
5231 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005232 Py_DECREF(dict);
5233 }
5234
5235 /* Also set instance attributes from the slotstate dict (if any). */
5236 if (slotstate != NULL) {
5237 PyObject *d_key, *d_value;
5238 Py_ssize_t i;
5239
5240 if (!PyDict_Check(slotstate)) {
5241 PyErr_SetString(UnpicklingError,
5242 "slot state is not a dictionary");
5243 goto error;
5244 }
5245 i = 0;
5246 while (PyDict_Next(slotstate, &i, &d_key, &d_value)) {
5247 if (PyObject_SetAttr(inst, d_key, d_value) < 0)
5248 goto error;
5249 }
5250 }
5251
5252 if (0) {
5253 error:
5254 status = -1;
5255 }
5256
5257 Py_DECREF(state);
5258 Py_XDECREF(slotstate);
5259 return status;
5260}
5261
5262static int
5263load_mark(UnpicklerObject *self)
5264{
5265
5266 /* Note that we split the (pickle.py) stack into two stacks, an
5267 * object stack and a mark stack. Here we push a mark onto the
5268 * mark stack.
5269 */
5270
5271 if ((self->num_marks + 1) >= self->marks_size) {
5272 size_t alloc;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005273 Py_ssize_t *marks;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005274
5275 /* Use the size_t type to check for overflow. */
5276 alloc = ((size_t)self->num_marks << 1) + 20;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005277 if (alloc > (PY_SSIZE_T_MAX / sizeof(Py_ssize_t)) ||
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00005278 alloc <= ((size_t)self->num_marks + 1)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005279 PyErr_NoMemory();
5280 return -1;
5281 }
5282
5283 if (self->marks == NULL)
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005284 marks = (Py_ssize_t *) PyMem_Malloc(alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005285 else
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005286 marks = (Py_ssize_t *) PyMem_Realloc(self->marks,
5287 alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005288 if (marks == NULL) {
5289 PyErr_NoMemory();
5290 return -1;
5291 }
5292 self->marks = marks;
5293 self->marks_size = (Py_ssize_t)alloc;
5294 }
5295
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005296 self->marks[self->num_marks++] = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005297
5298 return 0;
5299}
5300
5301static int
5302load_reduce(UnpicklerObject *self)
5303{
5304 PyObject *callable = NULL;
5305 PyObject *argtup = NULL;
5306 PyObject *obj = NULL;
5307
5308 PDATA_POP(self->stack, argtup);
5309 if (argtup == NULL)
5310 return -1;
5311 PDATA_POP(self->stack, callable);
5312 if (callable) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005313 obj = PyObject_CallObject(callable, argtup);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005314 Py_DECREF(callable);
5315 }
5316 Py_DECREF(argtup);
5317
5318 if (obj == NULL)
5319 return -1;
5320
5321 PDATA_PUSH(self->stack, obj, -1);
5322 return 0;
5323}
5324
5325/* Just raises an error if we don't know the protocol specified. PROTO
5326 * is the first opcode for protocols >= 2.
5327 */
5328static int
5329load_proto(UnpicklerObject *self)
5330{
5331 char *s;
5332 int i;
5333
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005334 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005335 return -1;
5336
5337 i = (unsigned char)s[0];
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005338 if (i <= HIGHEST_PROTOCOL) {
5339 self->proto = i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005340 return 0;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005341 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005342
5343 PyErr_Format(PyExc_ValueError, "unsupported pickle protocol: %d", i);
5344 return -1;
5345}
5346
5347static PyObject *
5348load(UnpicklerObject *self)
5349{
5350 PyObject *err;
5351 PyObject *value = NULL;
5352 char *s;
5353
5354 self->num_marks = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005355 if (Py_SIZE(self->stack))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005356 Pdata_clear(self->stack, 0);
5357
5358 /* Convenient macros for the dispatch while-switch loop just below. */
5359#define OP(opcode, load_func) \
5360 case opcode: if (load_func(self) < 0) break; continue;
5361
5362#define OP_ARG(opcode, load_func, arg) \
5363 case opcode: if (load_func(self, (arg)) < 0) break; continue;
5364
5365 while (1) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005366 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005367 break;
5368
5369 switch ((enum opcode)s[0]) {
5370 OP(NONE, load_none)
5371 OP(BININT, load_binint)
5372 OP(BININT1, load_binint1)
5373 OP(BININT2, load_binint2)
5374 OP(INT, load_int)
5375 OP(LONG, load_long)
5376 OP_ARG(LONG1, load_counted_long, 1)
5377 OP_ARG(LONG4, load_counted_long, 4)
5378 OP(FLOAT, load_float)
5379 OP(BINFLOAT, load_binfloat)
5380 OP(BINBYTES, load_binbytes)
5381 OP(SHORT_BINBYTES, load_short_binbytes)
5382 OP(BINSTRING, load_binstring)
5383 OP(SHORT_BINSTRING, load_short_binstring)
5384 OP(STRING, load_string)
5385 OP(UNICODE, load_unicode)
5386 OP(BINUNICODE, load_binunicode)
5387 OP_ARG(EMPTY_TUPLE, load_counted_tuple, 0)
5388 OP_ARG(TUPLE1, load_counted_tuple, 1)
5389 OP_ARG(TUPLE2, load_counted_tuple, 2)
5390 OP_ARG(TUPLE3, load_counted_tuple, 3)
5391 OP(TUPLE, load_tuple)
5392 OP(EMPTY_LIST, load_empty_list)
5393 OP(LIST, load_list)
5394 OP(EMPTY_DICT, load_empty_dict)
5395 OP(DICT, load_dict)
5396 OP(OBJ, load_obj)
5397 OP(INST, load_inst)
5398 OP(NEWOBJ, load_newobj)
5399 OP(GLOBAL, load_global)
5400 OP(APPEND, load_append)
5401 OP(APPENDS, load_appends)
5402 OP(BUILD, load_build)
5403 OP(DUP, load_dup)
5404 OP(BINGET, load_binget)
5405 OP(LONG_BINGET, load_long_binget)
5406 OP(GET, load_get)
5407 OP(MARK, load_mark)
5408 OP(BINPUT, load_binput)
5409 OP(LONG_BINPUT, load_long_binput)
5410 OP(PUT, load_put)
5411 OP(POP, load_pop)
5412 OP(POP_MARK, load_pop_mark)
5413 OP(SETITEM, load_setitem)
5414 OP(SETITEMS, load_setitems)
5415 OP(PERSID, load_persid)
5416 OP(BINPERSID, load_binpersid)
5417 OP(REDUCE, load_reduce)
5418 OP(PROTO, load_proto)
5419 OP_ARG(EXT1, load_extension, 1)
5420 OP_ARG(EXT2, load_extension, 2)
5421 OP_ARG(EXT4, load_extension, 4)
5422 OP_ARG(NEWTRUE, load_bool, Py_True)
5423 OP_ARG(NEWFALSE, load_bool, Py_False)
5424
5425 case STOP:
5426 break;
5427
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005428 default:
Benjamin Petersonadde86d2011-09-23 13:41:41 -04005429 if (s[0] == '\0')
5430 PyErr_SetNone(PyExc_EOFError);
5431 else
5432 PyErr_Format(UnpicklingError,
5433 "invalid load key, '%c'.", s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005434 return NULL;
5435 }
5436
5437 break; /* and we are done! */
5438 }
5439
Antoine Pitrou04248a82010-10-12 20:51:21 +00005440 if (_Unpickler_SkipConsumed(self) < 0)
5441 return NULL;
5442
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005443 /* XXX: It is not clear what this is actually for. */
5444 if ((err = PyErr_Occurred())) {
5445 if (err == PyExc_EOFError) {
5446 PyErr_SetNone(PyExc_EOFError);
5447 }
5448 return NULL;
5449 }
5450
5451 PDATA_POP(self->stack, value);
5452 return value;
5453}
5454
5455PyDoc_STRVAR(Unpickler_load_doc,
5456"load() -> object. Load a pickle."
5457"\n"
5458"Read a pickled object representation from the open file object given in\n"
5459"the constructor, and return the reconstituted object hierarchy specified\n"
5460"therein.\n");
5461
5462static PyObject *
5463Unpickler_load(UnpicklerObject *self)
5464{
5465 /* Check whether the Unpickler was initialized correctly. This prevents
5466 segfaulting if a subclass overridden __init__ with a function that does
5467 not call Unpickler.__init__(). Here, we simply ensure that self->read
5468 is not NULL. */
5469 if (self->read == NULL) {
Victor Stinner121aab42011-09-29 23:40:53 +02005470 PyErr_Format(UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005471 "Unpickler.__init__() was not called by %s.__init__()",
5472 Py_TYPE(self)->tp_name);
5473 return NULL;
5474 }
5475
5476 return load(self);
5477}
5478
5479/* The name of find_class() is misleading. In newer pickle protocols, this
5480 function is used for loading any global (i.e., functions), not just
5481 classes. The name is kept only for backward compatibility. */
5482
5483PyDoc_STRVAR(Unpickler_find_class_doc,
5484"find_class(module_name, global_name) -> object.\n"
5485"\n"
5486"Return an object from a specified module, importing the module if\n"
5487"necessary. Subclasses may override this method (e.g. to restrict\n"
5488"unpickling of arbitrary classes and functions).\n"
5489"\n"
5490"This method is called whenever a class or a function object is\n"
5491"needed. Both arguments passed are str objects.\n");
5492
5493static PyObject *
5494Unpickler_find_class(UnpicklerObject *self, PyObject *args)
5495{
5496 PyObject *global;
5497 PyObject *modules_dict;
5498 PyObject *module;
5499 PyObject *module_name, *global_name;
5500
5501 if (!PyArg_UnpackTuple(args, "find_class", 2, 2,
5502 &module_name, &global_name))
5503 return NULL;
5504
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005505 /* Try to map the old names used in Python 2.x to the new ones used in
5506 Python 3.x. We do this only with old pickle protocols and when the
5507 user has not disabled the feature. */
5508 if (self->proto < 3 && self->fix_imports) {
5509 PyObject *key;
5510 PyObject *item;
5511
5512 /* Check if the global (i.e., a function or a class) was renamed
5513 or moved to another module. */
5514 key = PyTuple_Pack(2, module_name, global_name);
5515 if (key == NULL)
5516 return NULL;
5517 item = PyDict_GetItemWithError(name_mapping_2to3, key);
5518 Py_DECREF(key);
5519 if (item) {
5520 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
5521 PyErr_Format(PyExc_RuntimeError,
5522 "_compat_pickle.NAME_MAPPING values should be "
5523 "2-tuples, not %.200s", Py_TYPE(item)->tp_name);
5524 return NULL;
5525 }
5526 module_name = PyTuple_GET_ITEM(item, 0);
5527 global_name = PyTuple_GET_ITEM(item, 1);
5528 if (!PyUnicode_Check(module_name) ||
5529 !PyUnicode_Check(global_name)) {
5530 PyErr_Format(PyExc_RuntimeError,
5531 "_compat_pickle.NAME_MAPPING values should be "
5532 "pairs of str, not (%.200s, %.200s)",
5533 Py_TYPE(module_name)->tp_name,
5534 Py_TYPE(global_name)->tp_name);
5535 return NULL;
5536 }
5537 }
5538 else if (PyErr_Occurred()) {
5539 return NULL;
5540 }
5541
5542 /* Check if the module was renamed. */
5543 item = PyDict_GetItemWithError(import_mapping_2to3, module_name);
5544 if (item) {
5545 if (!PyUnicode_Check(item)) {
5546 PyErr_Format(PyExc_RuntimeError,
5547 "_compat_pickle.IMPORT_MAPPING values should be "
5548 "strings, not %.200s", Py_TYPE(item)->tp_name);
5549 return NULL;
5550 }
5551 module_name = item;
5552 }
5553 else if (PyErr_Occurred()) {
5554 return NULL;
5555 }
5556 }
5557
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005558 modules_dict = PySys_GetObject("modules");
5559 if (modules_dict == NULL)
5560 return NULL;
5561
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005562 module = PyDict_GetItemWithError(modules_dict, module_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005563 if (module == NULL) {
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005564 if (PyErr_Occurred())
5565 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005566 module = PyImport_Import(module_name);
5567 if (module == NULL)
5568 return NULL;
5569 global = PyObject_GetAttr(module, global_name);
5570 Py_DECREF(module);
5571 }
Victor Stinner121aab42011-09-29 23:40:53 +02005572 else {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005573 global = PyObject_GetAttr(module, global_name);
5574 }
5575 return global;
5576}
5577
5578static struct PyMethodDef Unpickler_methods[] = {
5579 {"load", (PyCFunction)Unpickler_load, METH_NOARGS,
5580 Unpickler_load_doc},
5581 {"find_class", (PyCFunction)Unpickler_find_class, METH_VARARGS,
5582 Unpickler_find_class_doc},
5583 {NULL, NULL} /* sentinel */
5584};
5585
5586static void
5587Unpickler_dealloc(UnpicklerObject *self)
5588{
5589 PyObject_GC_UnTrack((PyObject *)self);
5590 Py_XDECREF(self->readline);
5591 Py_XDECREF(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005592 Py_XDECREF(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005593 Py_XDECREF(self->stack);
5594 Py_XDECREF(self->pers_func);
5595 Py_XDECREF(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005596 if (self->buffer.buf != NULL) {
5597 PyBuffer_Release(&self->buffer);
5598 self->buffer.buf = NULL;
5599 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005600
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005601 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005602 PyMem_Free(self->marks);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005603 PyMem_Free(self->input_line);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005604 free(self->encoding);
5605 free(self->errors);
5606
5607 Py_TYPE(self)->tp_free((PyObject *)self);
5608}
5609
5610static int
5611Unpickler_traverse(UnpicklerObject *self, visitproc visit, void *arg)
5612{
5613 Py_VISIT(self->readline);
5614 Py_VISIT(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005615 Py_VISIT(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005616 Py_VISIT(self->stack);
5617 Py_VISIT(self->pers_func);
5618 Py_VISIT(self->arg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005619 return 0;
5620}
5621
5622static int
5623Unpickler_clear(UnpicklerObject *self)
5624{
5625 Py_CLEAR(self->readline);
5626 Py_CLEAR(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005627 Py_CLEAR(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005628 Py_CLEAR(self->stack);
5629 Py_CLEAR(self->pers_func);
5630 Py_CLEAR(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005631 if (self->buffer.buf != NULL) {
5632 PyBuffer_Release(&self->buffer);
5633 self->buffer.buf = NULL;
5634 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005635
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005636 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005637 PyMem_Free(self->marks);
5638 self->marks = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005639 PyMem_Free(self->input_line);
5640 self->input_line = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005641 free(self->encoding);
5642 self->encoding = NULL;
5643 free(self->errors);
5644 self->errors = NULL;
5645
5646 return 0;
5647}
5648
5649PyDoc_STRVAR(Unpickler_doc,
5650"Unpickler(file, *, encoding='ASCII', errors='strict')"
5651"\n"
5652"This takes a binary file for reading a pickle data stream.\n"
5653"\n"
5654"The protocol version of the pickle is detected automatically, so no\n"
5655"proto argument is needed.\n"
5656"\n"
5657"The file-like object must have two methods, a read() method\n"
5658"that takes an integer argument, and a readline() method that\n"
5659"requires no arguments. Both methods should return bytes.\n"
5660"Thus file-like object can be a binary file object opened for\n"
5661"reading, a BytesIO object, or any other custom object that\n"
5662"meets this interface.\n"
5663"\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005664"Optional keyword arguments are *fix_imports*, *encoding* and *errors*,\n"
5665"which are used to control compatiblity support for pickle stream\n"
5666"generated by Python 2.x. If *fix_imports* is True, pickle will try to\n"
5667"map the old Python 2.x names to the new names used in Python 3.x. The\n"
5668"*encoding* and *errors* tell pickle how to decode 8-bit string\n"
5669"instances pickled by Python 2.x; these default to 'ASCII' and\n"
5670"'strict', respectively.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005671
5672static int
5673Unpickler_init(UnpicklerObject *self, PyObject *args, PyObject *kwds)
5674{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005675 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005676 PyObject *file;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005677 PyObject *fix_imports = Py_True;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005678 char *encoding = NULL;
5679 char *errors = NULL;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005680 _Py_IDENTIFIER(persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005681
5682 /* XXX: That is an horrible error message. But, I don't know how to do
5683 better... */
5684 if (Py_SIZE(args) != 1) {
5685 PyErr_Format(PyExc_TypeError,
5686 "%s takes exactly one positional argument (%zd given)",
5687 Py_TYPE(self)->tp_name, Py_SIZE(args));
5688 return -1;
5689 }
5690
5691 /* Arguments parsing needs to be done in the __init__() method to allow
5692 subclasses to define their own __init__() method, which may (or may
5693 not) support Unpickler arguments. However, this means we need to be
5694 extra careful in the other Unpickler methods, since a subclass could
5695 forget to call Unpickler.__init__() thus breaking our internal
5696 invariants. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005697 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:Unpickler", kwlist,
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005698 &file, &fix_imports, &encoding, &errors))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005699 return -1;
5700
5701 /* In case of multiple __init__() calls, clear previous content. */
5702 if (self->read != NULL)
5703 (void)Unpickler_clear(self);
5704
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005705 if (_Unpickler_SetInputStream(self, file) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005706 return -1;
5707
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005708 if (_Unpickler_SetInputEncoding(self, encoding, errors) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005709 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005710
5711 self->fix_imports = PyObject_IsTrue(fix_imports);
5712 if (self->fix_imports == -1)
5713 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005714
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005715 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_load)) {
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005716 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
5717 &PyId_persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005718 if (self->pers_func == NULL)
5719 return -1;
5720 }
5721 else {
5722 self->pers_func = NULL;
5723 }
5724
5725 self->stack = (Pdata *)Pdata_New();
5726 if (self->stack == NULL)
5727 return -1;
5728
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005729 self->memo_size = 32;
5730 self->memo = _Unpickler_NewMemo(self->memo_size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005731 if (self->memo == NULL)
5732 return -1;
5733
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005734 self->arg = NULL;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005735 self->proto = 0;
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005736
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005737 return 0;
5738}
5739
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005740/* Define a proxy object for the Unpickler's internal memo object. This is to
5741 * avoid breaking code like:
5742 * unpickler.memo.clear()
5743 * and
5744 * unpickler.memo = saved_memo
5745 * Is this a good idea? Not really, but we don't want to break code that uses
5746 * it. Note that we don't implement the entire mapping API here. This is
5747 * intentional, as these should be treated as black-box implementation details.
5748 *
5749 * We do, however, have to implement pickling/unpickling support because of
Victor Stinner121aab42011-09-29 23:40:53 +02005750 * real-world code like cvs2svn.
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005751 */
5752
5753typedef struct {
5754 PyObject_HEAD
5755 UnpicklerObject *unpickler;
5756} UnpicklerMemoProxyObject;
5757
5758PyDoc_STRVAR(ump_clear_doc,
5759"memo.clear() -> None. Remove all items from memo.");
5760
5761static PyObject *
5762ump_clear(UnpicklerMemoProxyObject *self)
5763{
5764 _Unpickler_MemoCleanup(self->unpickler);
5765 self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size);
5766 if (self->unpickler->memo == NULL)
5767 return NULL;
5768 Py_RETURN_NONE;
5769}
5770
5771PyDoc_STRVAR(ump_copy_doc,
5772"memo.copy() -> new_memo. Copy the memo to a new object.");
5773
5774static PyObject *
5775ump_copy(UnpicklerMemoProxyObject *self)
5776{
5777 Py_ssize_t i;
5778 PyObject *new_memo = PyDict_New();
5779 if (new_memo == NULL)
5780 return NULL;
5781
5782 for (i = 0; i < self->unpickler->memo_size; i++) {
5783 int status;
5784 PyObject *key, *value;
5785
5786 value = self->unpickler->memo[i];
5787 if (value == NULL)
5788 continue;
5789
5790 key = PyLong_FromSsize_t(i);
5791 if (key == NULL)
5792 goto error;
5793 status = PyDict_SetItem(new_memo, key, value);
5794 Py_DECREF(key);
5795 if (status < 0)
5796 goto error;
5797 }
5798 return new_memo;
5799
5800error:
5801 Py_DECREF(new_memo);
5802 return NULL;
5803}
5804
5805PyDoc_STRVAR(ump_reduce_doc,
5806"memo.__reduce__(). Pickling support.");
5807
5808static PyObject *
5809ump_reduce(UnpicklerMemoProxyObject *self, PyObject *args)
5810{
5811 PyObject *reduce_value;
5812 PyObject *constructor_args;
5813 PyObject *contents = ump_copy(self);
5814 if (contents == NULL)
5815 return NULL;
5816
5817 reduce_value = PyTuple_New(2);
5818 if (reduce_value == NULL) {
5819 Py_DECREF(contents);
5820 return NULL;
5821 }
5822 constructor_args = PyTuple_New(1);
5823 if (constructor_args == NULL) {
5824 Py_DECREF(contents);
5825 Py_DECREF(reduce_value);
5826 return NULL;
5827 }
5828 PyTuple_SET_ITEM(constructor_args, 0, contents);
5829 Py_INCREF((PyObject *)&PyDict_Type);
5830 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
5831 PyTuple_SET_ITEM(reduce_value, 1, constructor_args);
5832 return reduce_value;
5833}
5834
5835static PyMethodDef unpicklerproxy_methods[] = {
5836 {"clear", (PyCFunction)ump_clear, METH_NOARGS, ump_clear_doc},
5837 {"copy", (PyCFunction)ump_copy, METH_NOARGS, ump_copy_doc},
5838 {"__reduce__", (PyCFunction)ump_reduce, METH_VARARGS, ump_reduce_doc},
5839 {NULL, NULL} /* sentinel */
5840};
5841
5842static void
5843UnpicklerMemoProxy_dealloc(UnpicklerMemoProxyObject *self)
5844{
5845 PyObject_GC_UnTrack(self);
5846 Py_XDECREF(self->unpickler);
5847 PyObject_GC_Del((PyObject *)self);
5848}
5849
5850static int
5851UnpicklerMemoProxy_traverse(UnpicklerMemoProxyObject *self,
5852 visitproc visit, void *arg)
5853{
5854 Py_VISIT(self->unpickler);
5855 return 0;
5856}
5857
5858static int
5859UnpicklerMemoProxy_clear(UnpicklerMemoProxyObject *self)
5860{
5861 Py_CLEAR(self->unpickler);
5862 return 0;
5863}
5864
5865static PyTypeObject UnpicklerMemoProxyType = {
5866 PyVarObject_HEAD_INIT(NULL, 0)
5867 "_pickle.UnpicklerMemoProxy", /*tp_name*/
5868 sizeof(UnpicklerMemoProxyObject), /*tp_basicsize*/
5869 0,
5870 (destructor)UnpicklerMemoProxy_dealloc, /* tp_dealloc */
5871 0, /* tp_print */
5872 0, /* tp_getattr */
5873 0, /* tp_setattr */
5874 0, /* tp_compare */
5875 0, /* tp_repr */
5876 0, /* tp_as_number */
5877 0, /* tp_as_sequence */
5878 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00005879 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005880 0, /* tp_call */
5881 0, /* tp_str */
5882 PyObject_GenericGetAttr, /* tp_getattro */
5883 PyObject_GenericSetAttr, /* tp_setattro */
5884 0, /* tp_as_buffer */
5885 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5886 0, /* tp_doc */
5887 (traverseproc)UnpicklerMemoProxy_traverse, /* tp_traverse */
5888 (inquiry)UnpicklerMemoProxy_clear, /* tp_clear */
5889 0, /* tp_richcompare */
5890 0, /* tp_weaklistoffset */
5891 0, /* tp_iter */
5892 0, /* tp_iternext */
5893 unpicklerproxy_methods, /* tp_methods */
5894};
5895
5896static PyObject *
5897UnpicklerMemoProxy_New(UnpicklerObject *unpickler)
5898{
5899 UnpicklerMemoProxyObject *self;
5900
5901 self = PyObject_GC_New(UnpicklerMemoProxyObject,
5902 &UnpicklerMemoProxyType);
5903 if (self == NULL)
5904 return NULL;
5905 Py_INCREF(unpickler);
5906 self->unpickler = unpickler;
5907 PyObject_GC_Track(self);
5908 return (PyObject *)self;
5909}
5910
5911/*****************************************************************************/
5912
5913
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005914static PyObject *
5915Unpickler_get_memo(UnpicklerObject *self)
5916{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005917 return UnpicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005918}
5919
5920static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005921Unpickler_set_memo(UnpicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005922{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005923 PyObject **new_memo;
5924 Py_ssize_t new_memo_size = 0;
5925 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005926
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005927 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005928 PyErr_SetString(PyExc_TypeError,
5929 "attribute deletion is not supported");
5930 return -1;
5931 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005932
5933 if (Py_TYPE(obj) == &UnpicklerMemoProxyType) {
5934 UnpicklerObject *unpickler =
5935 ((UnpicklerMemoProxyObject *)obj)->unpickler;
5936
5937 new_memo_size = unpickler->memo_size;
5938 new_memo = _Unpickler_NewMemo(new_memo_size);
5939 if (new_memo == NULL)
5940 return -1;
5941
5942 for (i = 0; i < new_memo_size; i++) {
5943 Py_XINCREF(unpickler->memo[i]);
5944 new_memo[i] = unpickler->memo[i];
5945 }
5946 }
5947 else if (PyDict_Check(obj)) {
5948 Py_ssize_t i = 0;
5949 PyObject *key, *value;
5950
5951 new_memo_size = PyDict_Size(obj);
5952 new_memo = _Unpickler_NewMemo(new_memo_size);
5953 if (new_memo == NULL)
5954 return -1;
5955
5956 while (PyDict_Next(obj, &i, &key, &value)) {
5957 Py_ssize_t idx;
5958 if (!PyLong_Check(key)) {
5959 PyErr_SetString(PyExc_TypeError,
5960 "memo key must be integers");
5961 goto error;
5962 }
5963 idx = PyLong_AsSsize_t(key);
5964 if (idx == -1 && PyErr_Occurred())
5965 goto error;
5966 if (_Unpickler_MemoPut(self, idx, value) < 0)
5967 goto error;
5968 }
5969 }
5970 else {
5971 PyErr_Format(PyExc_TypeError,
5972 "'memo' attribute must be an UnpicklerMemoProxy object"
5973 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005974 return -1;
5975 }
5976
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005977 _Unpickler_MemoCleanup(self);
5978 self->memo_size = new_memo_size;
5979 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005980
5981 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005982
5983 error:
5984 if (new_memo_size) {
5985 i = new_memo_size;
5986 while (--i >= 0) {
5987 Py_XDECREF(new_memo[i]);
5988 }
5989 PyMem_FREE(new_memo);
5990 }
5991 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005992}
5993
5994static PyObject *
5995Unpickler_get_persload(UnpicklerObject *self)
5996{
5997 if (self->pers_func == NULL)
5998 PyErr_SetString(PyExc_AttributeError, "persistent_load");
5999 else
6000 Py_INCREF(self->pers_func);
6001 return self->pers_func;
6002}
6003
6004static int
6005Unpickler_set_persload(UnpicklerObject *self, PyObject *value)
6006{
6007 PyObject *tmp;
6008
6009 if (value == NULL) {
6010 PyErr_SetString(PyExc_TypeError,
6011 "attribute deletion is not supported");
6012 return -1;
6013 }
6014 if (!PyCallable_Check(value)) {
6015 PyErr_SetString(PyExc_TypeError,
6016 "persistent_load must be a callable taking "
6017 "one argument");
6018 return -1;
6019 }
6020
6021 tmp = self->pers_func;
6022 Py_INCREF(value);
6023 self->pers_func = value;
6024 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
6025
6026 return 0;
6027}
6028
6029static PyGetSetDef Unpickler_getsets[] = {
6030 {"memo", (getter)Unpickler_get_memo, (setter)Unpickler_set_memo},
6031 {"persistent_load", (getter)Unpickler_get_persload,
6032 (setter)Unpickler_set_persload},
6033 {NULL}
6034};
6035
6036static PyTypeObject Unpickler_Type = {
6037 PyVarObject_HEAD_INIT(NULL, 0)
6038 "_pickle.Unpickler", /*tp_name*/
6039 sizeof(UnpicklerObject), /*tp_basicsize*/
6040 0, /*tp_itemsize*/
6041 (destructor)Unpickler_dealloc, /*tp_dealloc*/
6042 0, /*tp_print*/
6043 0, /*tp_getattr*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006044 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00006045 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006046 0, /*tp_repr*/
6047 0, /*tp_as_number*/
6048 0, /*tp_as_sequence*/
6049 0, /*tp_as_mapping*/
6050 0, /*tp_hash*/
6051 0, /*tp_call*/
6052 0, /*tp_str*/
6053 0, /*tp_getattro*/
6054 0, /*tp_setattro*/
6055 0, /*tp_as_buffer*/
6056 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
6057 Unpickler_doc, /*tp_doc*/
6058 (traverseproc)Unpickler_traverse, /*tp_traverse*/
6059 (inquiry)Unpickler_clear, /*tp_clear*/
6060 0, /*tp_richcompare*/
6061 0, /*tp_weaklistoffset*/
6062 0, /*tp_iter*/
6063 0, /*tp_iternext*/
6064 Unpickler_methods, /*tp_methods*/
6065 0, /*tp_members*/
6066 Unpickler_getsets, /*tp_getset*/
6067 0, /*tp_base*/
6068 0, /*tp_dict*/
6069 0, /*tp_descr_get*/
6070 0, /*tp_descr_set*/
6071 0, /*tp_dictoffset*/
6072 (initproc)Unpickler_init, /*tp_init*/
6073 PyType_GenericAlloc, /*tp_alloc*/
6074 PyType_GenericNew, /*tp_new*/
6075 PyObject_GC_Del, /*tp_free*/
6076 0, /*tp_is_gc*/
6077};
6078
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006079PyDoc_STRVAR(pickle_dump_doc,
6080"dump(obj, file, protocol=None, *, fix_imports=True) -> None\n"
6081"\n"
6082"Write a pickled representation of obj to the open file object file. This\n"
6083"is equivalent to ``Pickler(file, protocol).dump(obj)``, but may be more\n"
6084"efficient.\n"
6085"\n"
6086"The optional protocol argument tells the pickler to use the given protocol;\n"
6087"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
6088"backward-incompatible protocol designed for Python 3.0.\n"
6089"\n"
6090"Specifying a negative protocol version selects the highest protocol version\n"
6091"supported. The higher the protocol used, the more recent the version of\n"
6092"Python needed to read the pickle produced.\n"
6093"\n"
6094"The file argument must have a write() method that accepts a single bytes\n"
6095"argument. It can thus be a file object opened for binary writing, a\n"
6096"io.BytesIO instance, or any other custom object that meets this interface.\n"
6097"\n"
6098"If fix_imports is True and protocol is less than 3, pickle will try to\n"
6099"map the new Python 3.x names to the old module names used in Python 2.x,\n"
6100"so that the pickle data stream is readable with Python 2.x.\n");
6101
6102static PyObject *
6103pickle_dump(PyObject *self, PyObject *args, PyObject *kwds)
6104{
6105 static char *kwlist[] = {"obj", "file", "protocol", "fix_imports", 0};
6106 PyObject *obj;
6107 PyObject *file;
6108 PyObject *proto = NULL;
6109 PyObject *fix_imports = Py_True;
6110 PicklerObject *pickler;
6111
6112 /* fix_imports is a keyword-only argument. */
6113 if (Py_SIZE(args) > 3) {
6114 PyErr_Format(PyExc_TypeError,
6115 "pickle.dump() takes at most 3 positional "
6116 "argument (%zd given)", Py_SIZE(args));
6117 return NULL;
6118 }
6119
6120 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|OO:dump", kwlist,
6121 &obj, &file, &proto, &fix_imports))
6122 return NULL;
6123
6124 pickler = _Pickler_New();
6125 if (pickler == NULL)
6126 return NULL;
6127
6128 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6129 goto error;
6130
6131 if (_Pickler_SetOutputStream(pickler, file) < 0)
6132 goto error;
6133
6134 if (dump(pickler, obj) < 0)
6135 goto error;
6136
6137 if (_Pickler_FlushToFile(pickler) < 0)
6138 goto error;
6139
6140 Py_DECREF(pickler);
6141 Py_RETURN_NONE;
6142
6143 error:
6144 Py_XDECREF(pickler);
6145 return NULL;
6146}
6147
6148PyDoc_STRVAR(pickle_dumps_doc,
6149"dumps(obj, protocol=None, *, fix_imports=True) -> bytes\n"
6150"\n"
6151"Return the pickled representation of the object as a bytes\n"
6152"object, instead of writing it to a file.\n"
6153"\n"
6154"The optional protocol argument tells the pickler to use the given protocol;\n"
6155"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
6156"backward-incompatible protocol designed for Python 3.0.\n"
6157"\n"
6158"Specifying a negative protocol version selects the highest protocol version\n"
6159"supported. The higher the protocol used, the more recent the version of\n"
6160"Python needed to read the pickle produced.\n"
6161"\n"
6162"If fix_imports is True and *protocol* is less than 3, pickle will try to\n"
6163"map the new Python 3.x names to the old module names used in Python 2.x,\n"
6164"so that the pickle data stream is readable with Python 2.x.\n");
6165
6166static PyObject *
6167pickle_dumps(PyObject *self, PyObject *args, PyObject *kwds)
6168{
6169 static char *kwlist[] = {"obj", "protocol", "fix_imports", 0};
6170 PyObject *obj;
6171 PyObject *proto = NULL;
6172 PyObject *result;
6173 PyObject *fix_imports = Py_True;
6174 PicklerObject *pickler;
6175
6176 /* fix_imports is a keyword-only argument. */
6177 if (Py_SIZE(args) > 2) {
6178 PyErr_Format(PyExc_TypeError,
6179 "pickle.dumps() takes at most 2 positional "
6180 "argument (%zd given)", Py_SIZE(args));
6181 return NULL;
6182 }
6183
6184 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:dumps", kwlist,
6185 &obj, &proto, &fix_imports))
6186 return NULL;
6187
6188 pickler = _Pickler_New();
6189 if (pickler == NULL)
6190 return NULL;
6191
6192 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6193 goto error;
6194
6195 if (dump(pickler, obj) < 0)
6196 goto error;
6197
6198 result = _Pickler_GetString(pickler);
6199 Py_DECREF(pickler);
6200 return result;
6201
6202 error:
6203 Py_XDECREF(pickler);
6204 return NULL;
6205}
6206
6207PyDoc_STRVAR(pickle_load_doc,
6208"load(file, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6209"\n"
6210"Read a pickled object representation from the open file object file and\n"
6211"return the reconstituted object hierarchy specified therein. This is\n"
6212"equivalent to ``Unpickler(file).load()``, but may be more efficient.\n"
6213"\n"
6214"The protocol version of the pickle is detected automatically, so no protocol\n"
6215"argument is needed. Bytes past the pickled object's representation are\n"
6216"ignored.\n"
6217"\n"
6218"The argument file must have two methods, a read() method that takes an\n"
6219"integer argument, and a readline() method that requires no arguments. Both\n"
6220"methods should return bytes. Thus *file* can be a binary file object opened\n"
6221"for reading, a BytesIO object, or any other custom object that meets this\n"
6222"interface.\n"
6223"\n"
6224"Optional keyword arguments are fix_imports, encoding and errors,\n"
6225"which are used to control compatiblity support for pickle stream generated\n"
6226"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6227"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6228"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6229"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6230
6231static PyObject *
6232pickle_load(PyObject *self, PyObject *args, PyObject *kwds)
6233{
6234 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
6235 PyObject *file;
6236 PyObject *fix_imports = Py_True;
6237 PyObject *result;
6238 char *encoding = NULL;
6239 char *errors = NULL;
6240 UnpicklerObject *unpickler;
6241
6242 /* fix_imports, encoding and errors are a keyword-only argument. */
6243 if (Py_SIZE(args) != 1) {
6244 PyErr_Format(PyExc_TypeError,
6245 "pickle.load() takes exactly one positional "
6246 "argument (%zd given)", Py_SIZE(args));
6247 return NULL;
6248 }
6249
6250 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:load", kwlist,
6251 &file, &fix_imports, &encoding, &errors))
6252 return NULL;
6253
6254 unpickler = _Unpickler_New();
6255 if (unpickler == NULL)
6256 return NULL;
6257
6258 if (_Unpickler_SetInputStream(unpickler, file) < 0)
6259 goto error;
6260
6261 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6262 goto error;
6263
6264 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6265 if (unpickler->fix_imports == -1)
6266 goto error;
6267
6268 result = load(unpickler);
6269 Py_DECREF(unpickler);
6270 return result;
6271
6272 error:
6273 Py_XDECREF(unpickler);
6274 return NULL;
6275}
6276
6277PyDoc_STRVAR(pickle_loads_doc,
6278"loads(input, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6279"\n"
6280"Read a pickled object hierarchy from a bytes object and return the\n"
6281"reconstituted object hierarchy specified therein\n"
6282"\n"
6283"The protocol version of the pickle is detected automatically, so no protocol\n"
6284"argument is needed. Bytes past the pickled object's representation are\n"
6285"ignored.\n"
6286"\n"
6287"Optional keyword arguments are fix_imports, encoding and errors, which\n"
6288"are used to control compatiblity support for pickle stream generated\n"
6289"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6290"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6291"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6292"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6293
6294static PyObject *
6295pickle_loads(PyObject *self, PyObject *args, PyObject *kwds)
6296{
6297 static char *kwlist[] = {"input", "fix_imports", "encoding", "errors", 0};
6298 PyObject *input;
6299 PyObject *fix_imports = Py_True;
6300 PyObject *result;
6301 char *encoding = NULL;
6302 char *errors = NULL;
6303 UnpicklerObject *unpickler;
6304
6305 /* fix_imports, encoding and errors are a keyword-only argument. */
6306 if (Py_SIZE(args) != 1) {
6307 PyErr_Format(PyExc_TypeError,
6308 "pickle.loads() takes exactly one positional "
6309 "argument (%zd given)", Py_SIZE(args));
6310 return NULL;
6311 }
6312
6313 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:loads", kwlist,
6314 &input, &fix_imports, &encoding, &errors))
6315 return NULL;
6316
6317 unpickler = _Unpickler_New();
6318 if (unpickler == NULL)
6319 return NULL;
6320
6321 if (_Unpickler_SetStringInput(unpickler, input) < 0)
6322 goto error;
6323
6324 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6325 goto error;
6326
6327 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6328 if (unpickler->fix_imports == -1)
6329 goto error;
6330
6331 result = load(unpickler);
6332 Py_DECREF(unpickler);
6333 return result;
6334
6335 error:
6336 Py_XDECREF(unpickler);
6337 return NULL;
6338}
6339
6340
6341static struct PyMethodDef pickle_methods[] = {
6342 {"dump", (PyCFunction)pickle_dump, METH_VARARGS|METH_KEYWORDS,
6343 pickle_dump_doc},
6344 {"dumps", (PyCFunction)pickle_dumps, METH_VARARGS|METH_KEYWORDS,
6345 pickle_dumps_doc},
6346 {"load", (PyCFunction)pickle_load, METH_VARARGS|METH_KEYWORDS,
6347 pickle_load_doc},
6348 {"loads", (PyCFunction)pickle_loads, METH_VARARGS|METH_KEYWORDS,
6349 pickle_loads_doc},
6350 {NULL, NULL} /* sentinel */
6351};
6352
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006353static int
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006354initmodule(void)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006355{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006356 PyObject *copyreg = NULL;
6357 PyObject *compat_pickle = NULL;
6358
6359 /* XXX: We should ensure that the types of the dictionaries imported are
6360 exactly PyDict objects. Otherwise, it is possible to crash the pickle
6361 since we use the PyDict API directly to access these dictionaries. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006362
6363 copyreg = PyImport_ImportModule("copyreg");
6364 if (!copyreg)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006365 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006366 dispatch_table = PyObject_GetAttrString(copyreg, "dispatch_table");
6367 if (!dispatch_table)
6368 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006369 extension_registry = \
6370 PyObject_GetAttrString(copyreg, "_extension_registry");
6371 if (!extension_registry)
6372 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006373 inverted_registry = PyObject_GetAttrString(copyreg, "_inverted_registry");
6374 if (!inverted_registry)
6375 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006376 extension_cache = PyObject_GetAttrString(copyreg, "_extension_cache");
6377 if (!extension_cache)
6378 goto error;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006379 Py_CLEAR(copyreg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006380
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006381 /* Load the 2.x -> 3.x stdlib module mapping tables */
6382 compat_pickle = PyImport_ImportModule("_compat_pickle");
6383 if (!compat_pickle)
6384 goto error;
6385 name_mapping_2to3 = PyObject_GetAttrString(compat_pickle, "NAME_MAPPING");
6386 if (!name_mapping_2to3)
6387 goto error;
6388 if (!PyDict_CheckExact(name_mapping_2to3)) {
6389 PyErr_Format(PyExc_RuntimeError,
6390 "_compat_pickle.NAME_MAPPING should be a dict, not %.200s",
6391 Py_TYPE(name_mapping_2to3)->tp_name);
6392 goto error;
6393 }
6394 import_mapping_2to3 = PyObject_GetAttrString(compat_pickle,
6395 "IMPORT_MAPPING");
6396 if (!import_mapping_2to3)
6397 goto error;
6398 if (!PyDict_CheckExact(import_mapping_2to3)) {
6399 PyErr_Format(PyExc_RuntimeError,
6400 "_compat_pickle.IMPORT_MAPPING should be a dict, "
6401 "not %.200s", Py_TYPE(import_mapping_2to3)->tp_name);
6402 goto error;
6403 }
6404 /* ... and the 3.x -> 2.x mapping tables */
6405 name_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6406 "REVERSE_NAME_MAPPING");
6407 if (!name_mapping_3to2)
6408 goto error;
6409 if (!PyDict_CheckExact(name_mapping_3to2)) {
6410 PyErr_Format(PyExc_RuntimeError,
Ezio Melotti13925002011-03-16 11:05:33 +02006411 "_compat_pickle.REVERSE_NAME_MAPPING should be a dict, "
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006412 "not %.200s", Py_TYPE(name_mapping_3to2)->tp_name);
6413 goto error;
6414 }
6415 import_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6416 "REVERSE_IMPORT_MAPPING");
6417 if (!import_mapping_3to2)
6418 goto error;
6419 if (!PyDict_CheckExact(import_mapping_3to2)) {
6420 PyErr_Format(PyExc_RuntimeError,
6421 "_compat_pickle.REVERSE_IMPORT_MAPPING should be a dict, "
6422 "not %.200s", Py_TYPE(import_mapping_3to2)->tp_name);
6423 goto error;
6424 }
6425 Py_CLEAR(compat_pickle);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006426
6427 empty_tuple = PyTuple_New(0);
6428 if (empty_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006429 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006430 two_tuple = PyTuple_New(2);
6431 if (two_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006432 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006433 /* We use this temp container with no regard to refcounts, or to
6434 * keeping containees alive. Exempt from GC, because we don't
6435 * want anything looking at two_tuple() by magic.
6436 */
6437 PyObject_GC_UnTrack(two_tuple);
6438
6439 return 0;
6440
6441 error:
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006442 Py_CLEAR(copyreg);
6443 Py_CLEAR(dispatch_table);
6444 Py_CLEAR(extension_registry);
6445 Py_CLEAR(inverted_registry);
6446 Py_CLEAR(extension_cache);
6447 Py_CLEAR(compat_pickle);
6448 Py_CLEAR(name_mapping_2to3);
6449 Py_CLEAR(import_mapping_2to3);
6450 Py_CLEAR(name_mapping_3to2);
6451 Py_CLEAR(import_mapping_3to2);
6452 Py_CLEAR(empty_tuple);
6453 Py_CLEAR(two_tuple);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006454 return -1;
6455}
6456
6457static struct PyModuleDef _picklemodule = {
6458 PyModuleDef_HEAD_INIT,
6459 "_pickle",
6460 pickle_module_doc,
6461 -1,
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006462 pickle_methods,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006463 NULL,
6464 NULL,
6465 NULL,
6466 NULL
6467};
6468
6469PyMODINIT_FUNC
6470PyInit__pickle(void)
6471{
6472 PyObject *m;
6473
6474 if (PyType_Ready(&Unpickler_Type) < 0)
6475 return NULL;
6476 if (PyType_Ready(&Pickler_Type) < 0)
6477 return NULL;
6478 if (PyType_Ready(&Pdata_Type) < 0)
6479 return NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006480 if (PyType_Ready(&PicklerMemoProxyType) < 0)
6481 return NULL;
6482 if (PyType_Ready(&UnpicklerMemoProxyType) < 0)
6483 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006484
6485 /* Create the module and add the functions. */
6486 m = PyModule_Create(&_picklemodule);
6487 if (m == NULL)
6488 return NULL;
6489
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006490 Py_INCREF(&Pickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006491 if (PyModule_AddObject(m, "Pickler", (PyObject *)&Pickler_Type) < 0)
6492 return NULL;
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006493 Py_INCREF(&Unpickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006494 if (PyModule_AddObject(m, "Unpickler", (PyObject *)&Unpickler_Type) < 0)
6495 return NULL;
6496
6497 /* Initialize the exceptions. */
6498 PickleError = PyErr_NewException("_pickle.PickleError", NULL, NULL);
6499 if (PickleError == NULL)
6500 return NULL;
6501 PicklingError = \
6502 PyErr_NewException("_pickle.PicklingError", PickleError, NULL);
6503 if (PicklingError == NULL)
6504 return NULL;
6505 UnpicklingError = \
6506 PyErr_NewException("_pickle.UnpicklingError", PickleError, NULL);
6507 if (UnpicklingError == NULL)
6508 return NULL;
6509
6510 if (PyModule_AddObject(m, "PickleError", PickleError) < 0)
6511 return NULL;
6512 if (PyModule_AddObject(m, "PicklingError", PicklingError) < 0)
6513 return NULL;
6514 if (PyModule_AddObject(m, "UnpicklingError", UnpicklingError) < 0)
6515 return NULL;
6516
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006517 if (initmodule() < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006518 return NULL;
6519
6520 return m;
6521}