blob: 2dc3a4116dc6c41a7eb7007a86e56647d4049b87 [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;
1592 if (nbytes > INT_MAX) {
1593 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,
1791 "cannot serialize a bytes object larger than 4GB");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001792 return -1; /* string too large */
1793 }
1794
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001795 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001796 return -1;
1797
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001798 if (_Pickler_Write(self, PyBytes_AS_STRING(obj), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001799 return -1;
1800
1801 if (memo_put(self, obj) < 0)
1802 return -1;
1803
1804 return 0;
1805 }
1806}
1807
1808/* A copy of PyUnicode_EncodeRawUnicodeEscape() that also translates
1809 backslash and newline characters to \uXXXX escapes. */
1810static PyObject *
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001811raw_unicode_escape(PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001812{
1813 PyObject *repr, *result;
1814 char *p;
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001815 Py_ssize_t i, size, expandsize;
1816 void *data;
1817 unsigned int kind;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001818
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001819 if (PyUnicode_READY(obj))
1820 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001821
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001822 size = PyUnicode_GET_LENGTH(obj);
1823 data = PyUnicode_DATA(obj);
1824 kind = PyUnicode_KIND(obj);
1825 if (kind == PyUnicode_4BYTE_KIND)
1826 expandsize = 10;
1827 else
1828 expandsize = 6;
Victor Stinner121aab42011-09-29 23:40:53 +02001829
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001830 if (size > PY_SSIZE_T_MAX / expandsize)
1831 return PyErr_NoMemory();
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001832 repr = PyByteArray_FromStringAndSize(NULL, expandsize * size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001833 if (repr == NULL)
1834 return NULL;
1835 if (size == 0)
1836 goto done;
1837
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001838 p = PyByteArray_AS_STRING(repr);
1839 for (i=0; i < size; i++) {
1840 Py_UCS4 ch = PyUnicode_READ(kind, data, i);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001841 /* Map 32-bit characters to '\Uxxxxxxxx' */
1842 if (ch >= 0x10000) {
1843 *p++ = '\\';
1844 *p++ = 'U';
Victor Stinnerf5cff562011-10-14 02:13:11 +02001845 *p++ = Py_hexdigits[(ch >> 28) & 0xf];
1846 *p++ = Py_hexdigits[(ch >> 24) & 0xf];
1847 *p++ = Py_hexdigits[(ch >> 20) & 0xf];
1848 *p++ = Py_hexdigits[(ch >> 16) & 0xf];
1849 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
1850 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
1851 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
1852 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001853 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001854 /* Map 16-bit characters to '\uxxxx' */
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001855 else if (ch >= 256 || ch == '\\' || ch == '\n') {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001856 *p++ = '\\';
1857 *p++ = 'u';
Victor Stinnerf5cff562011-10-14 02:13:11 +02001858 *p++ = Py_hexdigits[(ch >> 12) & 0xf];
1859 *p++ = Py_hexdigits[(ch >> 8) & 0xf];
1860 *p++ = Py_hexdigits[(ch >> 4) & 0xf];
1861 *p++ = Py_hexdigits[ch & 15];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001862 }
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001863 /* Copy everything else as-is */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001864 else
1865 *p++ = (char) ch;
1866 }
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001867 size = p - PyByteArray_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001868
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001869done:
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001870 result = PyBytes_FromStringAndSize(PyByteArray_AS_STRING(repr), size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001871 Py_DECREF(repr);
1872 return result;
1873}
1874
1875static int
1876save_unicode(PicklerObject *self, PyObject *obj)
1877{
1878 Py_ssize_t size;
1879 PyObject *encoded = NULL;
1880
1881 if (self->bin) {
1882 char pdata[5];
1883
Martin v. Löwisd63a3b82011-09-28 07:41:54 +02001884 encoded = PyUnicode_AsEncodedString(obj, "utf-8", "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001885 if (encoded == NULL)
1886 goto error;
1887
1888 size = PyBytes_GET_SIZE(encoded);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001889 if (size > 0xffffffffL) {
1890 PyErr_SetString(PyExc_OverflowError,
1891 "cannot serialize a string larger than 4GB");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001892 goto error; /* string too large */
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001893 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001894
1895 pdata[0] = BINUNICODE;
1896 pdata[1] = (unsigned char)(size & 0xff);
1897 pdata[2] = (unsigned char)((size >> 8) & 0xff);
1898 pdata[3] = (unsigned char)((size >> 16) & 0xff);
1899 pdata[4] = (unsigned char)((size >> 24) & 0xff);
1900
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001901 if (_Pickler_Write(self, pdata, 5) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001902 goto error;
1903
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001904 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001905 goto error;
1906 }
1907 else {
1908 const char unicode_op = UNICODE;
1909
Victor Stinnerc806fdc2011-09-29 23:50:23 +02001910 encoded = raw_unicode_escape(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001911 if (encoded == NULL)
1912 goto error;
1913
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001914 if (_Pickler_Write(self, &unicode_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001915 goto error;
1916
1917 size = PyBytes_GET_SIZE(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001918 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001919 goto error;
1920
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001921 if (_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001922 goto error;
1923 }
1924 if (memo_put(self, obj) < 0)
1925 goto error;
1926
1927 Py_DECREF(encoded);
1928 return 0;
1929
1930 error:
1931 Py_XDECREF(encoded);
1932 return -1;
1933}
1934
1935/* A helper for save_tuple. Push the len elements in tuple t on the stack. */
1936static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001937store_tuple_elements(PicklerObject *self, PyObject *t, Py_ssize_t len)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001938{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001939 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001940
1941 assert(PyTuple_Size(t) == len);
1942
1943 for (i = 0; i < len; i++) {
1944 PyObject *element = PyTuple_GET_ITEM(t, i);
1945
1946 if (element == NULL)
1947 return -1;
1948 if (save(self, element, 0) < 0)
1949 return -1;
1950 }
1951
1952 return 0;
1953}
1954
1955/* Tuples are ubiquitous in the pickle protocols, so many techniques are
1956 * used across protocols to minimize the space needed to pickle them.
1957 * Tuples are also the only builtin immutable type that can be recursive
1958 * (a tuple can be reached from itself), and that requires some subtle
1959 * magic so that it works in all cases. IOW, this is a long routine.
1960 */
1961static int
1962save_tuple(PicklerObject *self, PyObject *obj)
1963{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02001964 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001965
1966 const char mark_op = MARK;
1967 const char tuple_op = TUPLE;
1968 const char pop_op = POP;
1969 const char pop_mark_op = POP_MARK;
1970 const char len2opcode[] = {EMPTY_TUPLE, TUPLE1, TUPLE2, TUPLE3};
1971
1972 if ((len = PyTuple_Size(obj)) < 0)
1973 return -1;
1974
1975 if (len == 0) {
1976 char pdata[2];
1977
1978 if (self->proto) {
1979 pdata[0] = EMPTY_TUPLE;
1980 len = 1;
1981 }
1982 else {
1983 pdata[0] = MARK;
1984 pdata[1] = TUPLE;
1985 len = 2;
1986 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001987 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001988 return -1;
1989 return 0;
1990 }
1991
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001992 /* The tuple isn't in the memo now. If it shows up there after
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001993 * saving the tuple elements, the tuple must be recursive, in
1994 * which case we'll pop everything we put on the stack, and fetch
1995 * its value from the memo.
1996 */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001997 if (len <= 3 && self->proto >= 2) {
1998 /* Use TUPLE{1,2,3} opcodes. */
1999 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002000 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002001
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002002 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002003 /* pop the len elements */
2004 for (i = 0; i < len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002005 if (_Pickler_Write(self, &pop_op, 1) < 0)
2006 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002007 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002008 if (memo_get(self, obj) < 0)
2009 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002010
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002011 return 0;
2012 }
2013 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002014 if (_Pickler_Write(self, len2opcode + len, 1) < 0)
2015 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002016 }
2017 goto memoize;
2018 }
2019
2020 /* proto < 2 and len > 0, or proto >= 2 and len > 3.
2021 * Generate MARK e1 e2 ... TUPLE
2022 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002023 if (_Pickler_Write(self, &mark_op, 1) < 0)
2024 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002025
2026 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002027 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002028
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002029 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002030 /* pop the stack stuff we pushed */
2031 if (self->bin) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002032 if (_Pickler_Write(self, &pop_mark_op, 1) < 0)
2033 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002034 }
2035 else {
2036 /* Note that we pop one more than len, to remove
2037 * the MARK too.
2038 */
2039 for (i = 0; i <= len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002040 if (_Pickler_Write(self, &pop_op, 1) < 0)
2041 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002042 }
2043 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002044 if (memo_get(self, obj) < 0)
2045 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002046
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002047 return 0;
2048 }
2049 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002050 if (_Pickler_Write(self, &tuple_op, 1) < 0)
2051 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002052 }
2053
2054 memoize:
2055 if (memo_put(self, obj) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002056 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002057
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002058 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002059}
2060
2061/* iter is an iterator giving items, and we batch up chunks of
2062 * MARK item item ... item APPENDS
2063 * opcode sequences. Calling code should have arranged to first create an
2064 * empty list, or list-like object, for the APPENDS to operate on.
2065 * Returns 0 on success, <0 on error.
2066 */
2067static int
2068batch_list(PicklerObject *self, PyObject *iter)
2069{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002070 PyObject *obj = NULL;
2071 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002072 int i, n;
2073
2074 const char mark_op = MARK;
2075 const char append_op = APPEND;
2076 const char appends_op = APPENDS;
2077
2078 assert(iter != NULL);
2079
2080 /* XXX: I think this function could be made faster by avoiding the
2081 iterator interface and fetching objects directly from list using
2082 PyList_GET_ITEM.
2083 */
2084
2085 if (self->proto == 0) {
2086 /* APPENDS isn't available; do one at a time. */
2087 for (;;) {
2088 obj = PyIter_Next(iter);
2089 if (obj == NULL) {
2090 if (PyErr_Occurred())
2091 return -1;
2092 break;
2093 }
2094 i = save(self, obj, 0);
2095 Py_DECREF(obj);
2096 if (i < 0)
2097 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002098 if (_Pickler_Write(self, &append_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002099 return -1;
2100 }
2101 return 0;
2102 }
2103
2104 /* proto > 0: write in batches of BATCHSIZE. */
2105 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002106 /* Get first item */
2107 firstitem = PyIter_Next(iter);
2108 if (firstitem == NULL) {
2109 if (PyErr_Occurred())
2110 goto error;
2111
2112 /* nothing more to add */
2113 break;
2114 }
2115
2116 /* Try to get a second item */
2117 obj = PyIter_Next(iter);
2118 if (obj == NULL) {
2119 if (PyErr_Occurred())
2120 goto error;
2121
2122 /* Only one item to write */
2123 if (save(self, firstitem, 0) < 0)
2124 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002125 if (_Pickler_Write(self, &append_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002126 goto error;
2127 Py_CLEAR(firstitem);
2128 break;
2129 }
2130
2131 /* More than one item to write */
2132
2133 /* Pump out MARK, items, APPENDS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002134 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002135 goto error;
2136
2137 if (save(self, firstitem, 0) < 0)
2138 goto error;
2139 Py_CLEAR(firstitem);
2140 n = 1;
2141
2142 /* Fetch and save up to BATCHSIZE items */
2143 while (obj) {
2144 if (save(self, obj, 0) < 0)
2145 goto error;
2146 Py_CLEAR(obj);
2147 n += 1;
2148
2149 if (n == BATCHSIZE)
2150 break;
2151
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002152 obj = PyIter_Next(iter);
2153 if (obj == NULL) {
2154 if (PyErr_Occurred())
2155 goto error;
2156 break;
2157 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002158 }
2159
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002160 if (_Pickler_Write(self, &appends_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002161 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002162
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002163 } while (n == BATCHSIZE);
2164 return 0;
2165
2166 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002167 Py_XDECREF(firstitem);
2168 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002169 return -1;
2170}
2171
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002172/* This is a variant of batch_list() above, specialized for lists (with no
2173 * support for list subclasses). Like batch_list(), we batch up chunks of
2174 * MARK item item ... item APPENDS
2175 * opcode sequences. Calling code should have arranged to first create an
2176 * empty list, or list-like object, for the APPENDS to operate on.
2177 * Returns 0 on success, -1 on error.
2178 *
2179 * This version is considerably faster than batch_list(), if less general.
2180 *
2181 * Note that this only works for protocols > 0.
2182 */
2183static int
2184batch_list_exact(PicklerObject *self, PyObject *obj)
2185{
2186 PyObject *item = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002187 Py_ssize_t this_batch, total;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002188
2189 const char append_op = APPEND;
2190 const char appends_op = APPENDS;
2191 const char mark_op = MARK;
2192
2193 assert(obj != NULL);
2194 assert(self->proto > 0);
2195 assert(PyList_CheckExact(obj));
2196
2197 if (PyList_GET_SIZE(obj) == 1) {
2198 item = PyList_GET_ITEM(obj, 0);
2199 if (save(self, item, 0) < 0)
2200 return -1;
2201 if (_Pickler_Write(self, &append_op, 1) < 0)
2202 return -1;
2203 return 0;
2204 }
2205
2206 /* Write in batches of BATCHSIZE. */
2207 total = 0;
2208 do {
2209 this_batch = 0;
2210 if (_Pickler_Write(self, &mark_op, 1) < 0)
2211 return -1;
2212 while (total < PyList_GET_SIZE(obj)) {
2213 item = PyList_GET_ITEM(obj, total);
2214 if (save(self, item, 0) < 0)
2215 return -1;
2216 total++;
2217 if (++this_batch == BATCHSIZE)
2218 break;
2219 }
2220 if (_Pickler_Write(self, &appends_op, 1) < 0)
2221 return -1;
2222
2223 } while (total < PyList_GET_SIZE(obj));
2224
2225 return 0;
2226}
2227
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002228static int
2229save_list(PicklerObject *self, PyObject *obj)
2230{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002231 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002232 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002233 int status = 0;
2234
2235 if (self->fast && !fast_save_enter(self, obj))
2236 goto error;
2237
2238 /* Create an empty list. */
2239 if (self->bin) {
2240 header[0] = EMPTY_LIST;
2241 len = 1;
2242 }
2243 else {
2244 header[0] = MARK;
2245 header[1] = LIST;
2246 len = 2;
2247 }
2248
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002249 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002250 goto error;
2251
2252 /* Get list length, and bow out early if empty. */
2253 if ((len = PyList_Size(obj)) < 0)
2254 goto error;
2255
2256 if (memo_put(self, obj) < 0)
2257 goto error;
2258
2259 if (len != 0) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002260 /* Materialize the list elements. */
2261 if (PyList_CheckExact(obj) && self->proto > 0) {
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002262 if (Py_EnterRecursiveCall(" while pickling an object"))
2263 goto error;
2264 status = batch_list_exact(self, obj);
2265 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002266 } else {
2267 PyObject *iter = PyObject_GetIter(obj);
2268 if (iter == NULL)
2269 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002270
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002271 if (Py_EnterRecursiveCall(" while pickling an object")) {
2272 Py_DECREF(iter);
2273 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002274 }
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002275 status = batch_list(self, iter);
2276 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002277 Py_DECREF(iter);
2278 }
2279 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002280 if (0) {
2281 error:
2282 status = -1;
2283 }
2284
2285 if (self->fast && !fast_save_leave(self, obj))
2286 status = -1;
2287
2288 return status;
2289}
2290
2291/* iter is an iterator giving (key, value) pairs, and we batch up chunks of
2292 * MARK key value ... key value SETITEMS
2293 * opcode sequences. Calling code should have arranged to first create an
2294 * empty dict, or dict-like object, for the SETITEMS to operate on.
2295 * Returns 0 on success, <0 on error.
2296 *
2297 * This is very much like batch_list(). The difference between saving
2298 * elements directly, and picking apart two-tuples, is so long-winded at
2299 * the C level, though, that attempts to combine these routines were too
2300 * ugly to bear.
2301 */
2302static int
2303batch_dict(PicklerObject *self, PyObject *iter)
2304{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002305 PyObject *obj = NULL;
2306 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002307 int i, n;
2308
2309 const char mark_op = MARK;
2310 const char setitem_op = SETITEM;
2311 const char setitems_op = SETITEMS;
2312
2313 assert(iter != NULL);
2314
2315 if (self->proto == 0) {
2316 /* SETITEMS isn't available; do one at a time. */
2317 for (;;) {
2318 obj = PyIter_Next(iter);
2319 if (obj == NULL) {
2320 if (PyErr_Occurred())
2321 return -1;
2322 break;
2323 }
2324 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2325 PyErr_SetString(PyExc_TypeError, "dict items "
2326 "iterator must return 2-tuples");
2327 return -1;
2328 }
2329 i = save(self, PyTuple_GET_ITEM(obj, 0), 0);
2330 if (i >= 0)
2331 i = save(self, PyTuple_GET_ITEM(obj, 1), 0);
2332 Py_DECREF(obj);
2333 if (i < 0)
2334 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002335 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002336 return -1;
2337 }
2338 return 0;
2339 }
2340
2341 /* proto > 0: write in batches of BATCHSIZE. */
2342 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002343 /* Get first item */
2344 firstitem = PyIter_Next(iter);
2345 if (firstitem == NULL) {
2346 if (PyErr_Occurred())
2347 goto error;
2348
2349 /* nothing more to add */
2350 break;
2351 }
2352 if (!PyTuple_Check(firstitem) || PyTuple_Size(firstitem) != 2) {
2353 PyErr_SetString(PyExc_TypeError, "dict items "
2354 "iterator must return 2-tuples");
2355 goto error;
2356 }
2357
2358 /* Try to get a second item */
2359 obj = PyIter_Next(iter);
2360 if (obj == NULL) {
2361 if (PyErr_Occurred())
2362 goto error;
2363
2364 /* Only one item to write */
2365 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2366 goto error;
2367 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2368 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002369 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002370 goto error;
2371 Py_CLEAR(firstitem);
2372 break;
2373 }
2374
2375 /* More than one item to write */
2376
2377 /* Pump out MARK, items, SETITEMS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002378 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002379 goto error;
2380
2381 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2382 goto error;
2383 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2384 goto error;
2385 Py_CLEAR(firstitem);
2386 n = 1;
2387
2388 /* Fetch and save up to BATCHSIZE items */
2389 while (obj) {
2390 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2391 PyErr_SetString(PyExc_TypeError, "dict items "
2392 "iterator must return 2-tuples");
2393 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002394 }
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002395 if (save(self, PyTuple_GET_ITEM(obj, 0), 0) < 0 ||
2396 save(self, PyTuple_GET_ITEM(obj, 1), 0) < 0)
2397 goto error;
2398 Py_CLEAR(obj);
2399 n += 1;
2400
2401 if (n == BATCHSIZE)
2402 break;
2403
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002404 obj = PyIter_Next(iter);
2405 if (obj == NULL) {
2406 if (PyErr_Occurred())
2407 goto error;
2408 break;
2409 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002410 }
2411
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002412 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002413 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002414
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002415 } while (n == BATCHSIZE);
2416 return 0;
2417
2418 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002419 Py_XDECREF(firstitem);
2420 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002421 return -1;
2422}
2423
Collin Winter5c9b02d2009-05-25 05:43:30 +00002424/* This is a variant of batch_dict() above that specializes for dicts, with no
2425 * support for dict subclasses. Like batch_dict(), we batch up chunks of
2426 * MARK key value ... key value SETITEMS
2427 * opcode sequences. Calling code should have arranged to first create an
2428 * empty dict, or dict-like object, for the SETITEMS to operate on.
2429 * Returns 0 on success, -1 on error.
2430 *
2431 * Note that this currently doesn't work for protocol 0.
2432 */
2433static int
2434batch_dict_exact(PicklerObject *self, PyObject *obj)
2435{
2436 PyObject *key = NULL, *value = NULL;
2437 int i;
2438 Py_ssize_t dict_size, ppos = 0;
2439
Alexandre Vassalottif70b1292009-05-25 18:00:52 +00002440 const char mark_op = MARK;
2441 const char setitem_op = SETITEM;
2442 const char setitems_op = SETITEMS;
Collin Winter5c9b02d2009-05-25 05:43:30 +00002443
2444 assert(obj != NULL);
2445 assert(self->proto > 0);
2446
2447 dict_size = PyDict_Size(obj);
2448
2449 /* Special-case len(d) == 1 to save space. */
2450 if (dict_size == 1) {
2451 PyDict_Next(obj, &ppos, &key, &value);
2452 if (save(self, key, 0) < 0)
2453 return -1;
2454 if (save(self, value, 0) < 0)
2455 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002456 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002457 return -1;
2458 return 0;
2459 }
2460
2461 /* Write in batches of BATCHSIZE. */
2462 do {
2463 i = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002464 if (_Pickler_Write(self, &mark_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002465 return -1;
2466 while (PyDict_Next(obj, &ppos, &key, &value)) {
2467 if (save(self, key, 0) < 0)
2468 return -1;
2469 if (save(self, value, 0) < 0)
2470 return -1;
2471 if (++i == BATCHSIZE)
2472 break;
2473 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002474 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002475 return -1;
2476 if (PyDict_Size(obj) != dict_size) {
2477 PyErr_Format(
2478 PyExc_RuntimeError,
2479 "dictionary changed size during iteration");
2480 return -1;
2481 }
2482
2483 } while (i == BATCHSIZE);
2484 return 0;
2485}
2486
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002487static int
2488save_dict(PicklerObject *self, PyObject *obj)
2489{
2490 PyObject *items, *iter;
2491 char header[3];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002492 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002493 int status = 0;
2494
2495 if (self->fast && !fast_save_enter(self, obj))
2496 goto error;
2497
2498 /* Create an empty dict. */
2499 if (self->bin) {
2500 header[0] = EMPTY_DICT;
2501 len = 1;
2502 }
2503 else {
2504 header[0] = MARK;
2505 header[1] = DICT;
2506 len = 2;
2507 }
2508
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002509 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002510 goto error;
2511
2512 /* Get dict size, and bow out early if empty. */
2513 if ((len = PyDict_Size(obj)) < 0)
2514 goto error;
2515
2516 if (memo_put(self, obj) < 0)
2517 goto error;
2518
2519 if (len != 0) {
2520 /* Save the dict items. */
Collin Winter5c9b02d2009-05-25 05:43:30 +00002521 if (PyDict_CheckExact(obj) && self->proto > 0) {
2522 /* We can take certain shortcuts if we know this is a dict and
2523 not a dict subclass. */
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002524 if (Py_EnterRecursiveCall(" while pickling an object"))
2525 goto error;
2526 status = batch_dict_exact(self, obj);
2527 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002528 } else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02002529 _Py_IDENTIFIER(items);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02002530
2531 items = _PyObject_CallMethodId(obj, &PyId_items, "()");
Collin Winter5c9b02d2009-05-25 05:43:30 +00002532 if (items == NULL)
2533 goto error;
2534 iter = PyObject_GetIter(items);
2535 Py_DECREF(items);
2536 if (iter == NULL)
2537 goto error;
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002538 if (Py_EnterRecursiveCall(" while pickling an object")) {
2539 Py_DECREF(iter);
2540 goto error;
2541 }
Collin Winter5c9b02d2009-05-25 05:43:30 +00002542 status = batch_dict(self, iter);
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002543 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002544 Py_DECREF(iter);
2545 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002546 }
2547
2548 if (0) {
2549 error:
2550 status = -1;
2551 }
2552
2553 if (self->fast && !fast_save_leave(self, obj))
2554 status = -1;
2555
2556 return status;
2557}
2558
2559static int
2560save_global(PicklerObject *self, PyObject *obj, PyObject *name)
2561{
2562 static PyObject *name_str = NULL;
2563 PyObject *global_name = NULL;
2564 PyObject *module_name = NULL;
2565 PyObject *module = NULL;
2566 PyObject *cls;
2567 int status = 0;
2568
2569 const char global_op = GLOBAL;
2570
2571 if (name_str == NULL) {
2572 name_str = PyUnicode_InternFromString("__name__");
2573 if (name_str == NULL)
2574 goto error;
2575 }
2576
2577 if (name) {
2578 global_name = name;
2579 Py_INCREF(global_name);
2580 }
2581 else {
2582 global_name = PyObject_GetAttr(obj, name_str);
2583 if (global_name == NULL)
2584 goto error;
2585 }
2586
2587 module_name = whichmodule(obj, global_name);
2588 if (module_name == NULL)
2589 goto error;
2590
2591 /* XXX: Change to use the import C API directly with level=0 to disallow
2592 relative imports.
2593
2594 XXX: PyImport_ImportModuleLevel could be used. However, this bypasses
2595 builtins.__import__. Therefore, _pickle, unlike pickle.py, will ignore
2596 custom import functions (IMHO, this would be a nice security
2597 feature). The import C API would need to be extended to support the
2598 extra parameters of __import__ to fix that. */
2599 module = PyImport_Import(module_name);
2600 if (module == NULL) {
2601 PyErr_Format(PicklingError,
2602 "Can't pickle %R: import of module %R failed",
2603 obj, module_name);
2604 goto error;
2605 }
2606 cls = PyObject_GetAttr(module, global_name);
2607 if (cls == NULL) {
2608 PyErr_Format(PicklingError,
2609 "Can't pickle %R: attribute lookup %S.%S failed",
2610 obj, module_name, global_name);
2611 goto error;
2612 }
2613 if (cls != obj) {
2614 Py_DECREF(cls);
2615 PyErr_Format(PicklingError,
2616 "Can't pickle %R: it's not the same object as %S.%S",
2617 obj, module_name, global_name);
2618 goto error;
2619 }
2620 Py_DECREF(cls);
2621
2622 if (self->proto >= 2) {
2623 /* See whether this is in the extension registry, and if
2624 * so generate an EXT opcode.
2625 */
2626 PyObject *code_obj; /* extension code as Python object */
2627 long code; /* extension code as C value */
2628 char pdata[5];
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002629 Py_ssize_t n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002630
2631 PyTuple_SET_ITEM(two_tuple, 0, module_name);
2632 PyTuple_SET_ITEM(two_tuple, 1, global_name);
2633 code_obj = PyDict_GetItem(extension_registry, two_tuple);
2634 /* The object is not registered in the extension registry.
2635 This is the most likely code path. */
2636 if (code_obj == NULL)
2637 goto gen_global;
2638
2639 /* XXX: pickle.py doesn't check neither the type, nor the range
2640 of the value returned by the extension_registry. It should for
2641 consistency. */
2642
2643 /* Verify code_obj has the right type and value. */
2644 if (!PyLong_Check(code_obj)) {
2645 PyErr_Format(PicklingError,
2646 "Can't pickle %R: extension code %R isn't an integer",
2647 obj, code_obj);
2648 goto error;
2649 }
2650 code = PyLong_AS_LONG(code_obj);
2651 if (code <= 0 || code > 0x7fffffffL) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02002652 if (!PyErr_Occurred())
2653 PyErr_Format(PicklingError,
2654 "Can't pickle %R: extension code %ld is out of range",
2655 obj, code);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002656 goto error;
2657 }
2658
2659 /* Generate an EXT opcode. */
2660 if (code <= 0xff) {
2661 pdata[0] = EXT1;
2662 pdata[1] = (unsigned char)code;
2663 n = 2;
2664 }
2665 else if (code <= 0xffff) {
2666 pdata[0] = EXT2;
2667 pdata[1] = (unsigned char)(code & 0xff);
2668 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2669 n = 3;
2670 }
2671 else {
2672 pdata[0] = EXT4;
2673 pdata[1] = (unsigned char)(code & 0xff);
2674 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2675 pdata[3] = (unsigned char)((code >> 16) & 0xff);
2676 pdata[4] = (unsigned char)((code >> 24) & 0xff);
2677 n = 5;
2678 }
2679
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002680 if (_Pickler_Write(self, pdata, n) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002681 goto error;
2682 }
2683 else {
2684 /* Generate a normal global opcode if we are using a pickle
2685 protocol <= 2, or if the object is not registered in the
2686 extension registry. */
2687 PyObject *encoded;
2688 PyObject *(*unicode_encoder)(PyObject *);
2689
2690 gen_global:
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002691 if (_Pickler_Write(self, &global_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002692 goto error;
2693
2694 /* Since Python 3.0 now supports non-ASCII identifiers, we encode both
2695 the module name and the global name using UTF-8. We do so only when
2696 we are using the pickle protocol newer than version 3. This is to
2697 ensure compatibility with older Unpickler running on Python 2.x. */
2698 if (self->proto >= 3) {
2699 unicode_encoder = PyUnicode_AsUTF8String;
2700 }
2701 else {
2702 unicode_encoder = PyUnicode_AsASCIIString;
2703 }
2704
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00002705 /* For protocol < 3 and if the user didn't request against doing so,
2706 we convert module names to the old 2.x module names. */
2707 if (self->fix_imports) {
2708 PyObject *key;
2709 PyObject *item;
2710
2711 key = PyTuple_Pack(2, module_name, global_name);
2712 if (key == NULL)
2713 goto error;
2714 item = PyDict_GetItemWithError(name_mapping_3to2, key);
2715 Py_DECREF(key);
2716 if (item) {
2717 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
2718 PyErr_Format(PyExc_RuntimeError,
2719 "_compat_pickle.REVERSE_NAME_MAPPING values "
2720 "should be 2-tuples, not %.200s",
2721 Py_TYPE(item)->tp_name);
2722 goto error;
2723 }
2724 Py_CLEAR(module_name);
2725 Py_CLEAR(global_name);
2726 module_name = PyTuple_GET_ITEM(item, 0);
2727 global_name = PyTuple_GET_ITEM(item, 1);
2728 if (!PyUnicode_Check(module_name) ||
2729 !PyUnicode_Check(global_name)) {
2730 PyErr_Format(PyExc_RuntimeError,
2731 "_compat_pickle.REVERSE_NAME_MAPPING values "
2732 "should be pairs of str, not (%.200s, %.200s)",
2733 Py_TYPE(module_name)->tp_name,
2734 Py_TYPE(global_name)->tp_name);
2735 goto error;
2736 }
2737 Py_INCREF(module_name);
2738 Py_INCREF(global_name);
2739 }
2740 else if (PyErr_Occurred()) {
2741 goto error;
2742 }
2743
2744 item = PyDict_GetItemWithError(import_mapping_3to2, module_name);
2745 if (item) {
2746 if (!PyUnicode_Check(item)) {
2747 PyErr_Format(PyExc_RuntimeError,
2748 "_compat_pickle.REVERSE_IMPORT_MAPPING values "
2749 "should be strings, not %.200s",
2750 Py_TYPE(item)->tp_name);
2751 goto error;
2752 }
2753 Py_CLEAR(module_name);
2754 module_name = item;
2755 Py_INCREF(module_name);
2756 }
2757 else if (PyErr_Occurred()) {
2758 goto error;
2759 }
2760 }
2761
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002762 /* Save the name of the module. */
2763 encoded = unicode_encoder(module_name);
2764 if (encoded == NULL) {
2765 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2766 PyErr_Format(PicklingError,
2767 "can't pickle module identifier '%S' using "
2768 "pickle protocol %i", module_name, self->proto);
2769 goto error;
2770 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002771 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002772 PyBytes_GET_SIZE(encoded)) < 0) {
2773 Py_DECREF(encoded);
2774 goto error;
2775 }
2776 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002777 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002778 goto error;
2779
2780 /* Save the name of the module. */
2781 encoded = unicode_encoder(global_name);
2782 if (encoded == NULL) {
2783 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2784 PyErr_Format(PicklingError,
2785 "can't pickle global identifier '%S' using "
2786 "pickle protocol %i", global_name, self->proto);
2787 goto error;
2788 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002789 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002790 PyBytes_GET_SIZE(encoded)) < 0) {
2791 Py_DECREF(encoded);
2792 goto error;
2793 }
2794 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002795 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002796 goto error;
2797
2798 /* Memoize the object. */
2799 if (memo_put(self, obj) < 0)
2800 goto error;
2801 }
2802
2803 if (0) {
2804 error:
2805 status = -1;
2806 }
2807 Py_XDECREF(module_name);
2808 Py_XDECREF(global_name);
2809 Py_XDECREF(module);
2810
2811 return status;
2812}
2813
2814static int
2815save_pers(PicklerObject *self, PyObject *obj, PyObject *func)
2816{
2817 PyObject *pid = NULL;
2818 int status = 0;
2819
2820 const char persid_op = PERSID;
2821 const char binpersid_op = BINPERSID;
2822
2823 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002824 pid = _Pickler_FastCall(self, func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002825 if (pid == NULL)
2826 return -1;
2827
2828 if (pid != Py_None) {
2829 if (self->bin) {
2830 if (save(self, pid, 1) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002831 _Pickler_Write(self, &binpersid_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002832 goto error;
2833 }
2834 else {
2835 PyObject *pid_str = NULL;
2836 char *pid_ascii_bytes;
2837 Py_ssize_t size;
2838
2839 pid_str = PyObject_Str(pid);
2840 if (pid_str == NULL)
2841 goto error;
2842
2843 /* XXX: Should it check whether the persistent id only contains
2844 ASCII characters? And what if the pid contains embedded
2845 newlines? */
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00002846 pid_ascii_bytes = _PyUnicode_AsStringAndSize(pid_str, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002847 Py_DECREF(pid_str);
2848 if (pid_ascii_bytes == NULL)
2849 goto error;
2850
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002851 if (_Pickler_Write(self, &persid_op, 1) < 0 ||
2852 _Pickler_Write(self, pid_ascii_bytes, size) < 0 ||
2853 _Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002854 goto error;
2855 }
2856 status = 1;
2857 }
2858
2859 if (0) {
2860 error:
2861 status = -1;
2862 }
2863 Py_XDECREF(pid);
2864
2865 return status;
2866}
2867
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002868static PyObject *
2869get_class(PyObject *obj)
2870{
2871 PyObject *cls;
2872 static PyObject *str_class;
2873
2874 if (str_class == NULL) {
2875 str_class = PyUnicode_InternFromString("__class__");
2876 if (str_class == NULL)
2877 return NULL;
2878 }
2879 cls = PyObject_GetAttr(obj, str_class);
2880 if (cls == NULL) {
2881 if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
2882 PyErr_Clear();
2883 cls = (PyObject *) Py_TYPE(obj);
2884 Py_INCREF(cls);
2885 }
2886 }
2887 return cls;
2888}
2889
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002890/* We're saving obj, and args is the 2-thru-5 tuple returned by the
2891 * appropriate __reduce__ method for obj.
2892 */
2893static int
2894save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
2895{
2896 PyObject *callable;
2897 PyObject *argtup;
2898 PyObject *state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002899 PyObject *listitems = Py_None;
2900 PyObject *dictitems = Py_None;
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002901 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002902
2903 int use_newobj = self->proto >= 2;
2904
2905 const char reduce_op = REDUCE;
2906 const char build_op = BUILD;
2907 const char newobj_op = NEWOBJ;
2908
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002909 size = PyTuple_Size(args);
2910 if (size < 2 || size > 5) {
2911 PyErr_SetString(PicklingError, "tuple returned by "
2912 "__reduce__ must contain 2 through 5 elements");
2913 return -1;
2914 }
2915
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002916 if (!PyArg_UnpackTuple(args, "save_reduce", 2, 5,
2917 &callable, &argtup, &state, &listitems, &dictitems))
2918 return -1;
2919
2920 if (!PyCallable_Check(callable)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002921 PyErr_SetString(PicklingError, "first item of the tuple "
2922 "returned by __reduce__ must be callable");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002923 return -1;
2924 }
2925 if (!PyTuple_Check(argtup)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002926 PyErr_SetString(PicklingError, "second item of the tuple "
2927 "returned by __reduce__ must be a tuple");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002928 return -1;
2929 }
2930
2931 if (state == Py_None)
2932 state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002933
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002934 if (listitems == Py_None)
2935 listitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002936 else if (!PyIter_Check(listitems)) {
2937 PyErr_Format(PicklingError, "Fourth element of tuple"
2938 "returned by __reduce__ must be an iterator, not %s",
2939 Py_TYPE(listitems)->tp_name);
2940 return -1;
2941 }
2942
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002943 if (dictitems == Py_None)
2944 dictitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002945 else if (!PyIter_Check(dictitems)) {
2946 PyErr_Format(PicklingError, "Fifth element of tuple"
2947 "returned by __reduce__ must be an iterator, not %s",
2948 Py_TYPE(dictitems)->tp_name);
2949 return -1;
2950 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002951
2952 /* Protocol 2 special case: if callable's name is __newobj__, use
2953 NEWOBJ. */
2954 if (use_newobj) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002955 static PyObject *newobj_str = NULL, *name_str = NULL;
2956 PyObject *name;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002957
2958 if (newobj_str == NULL) {
2959 newobj_str = PyUnicode_InternFromString("__newobj__");
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002960 name_str = PyUnicode_InternFromString("__name__");
2961 if (newobj_str == NULL || name_str == NULL)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002962 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002963 }
2964
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002965 name = PyObject_GetAttr(callable, name_str);
2966 if (name == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002967 if (PyErr_ExceptionMatches(PyExc_AttributeError))
2968 PyErr_Clear();
2969 else
2970 return -1;
2971 use_newobj = 0;
2972 }
2973 else {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002974 use_newobj = PyUnicode_Check(name) &&
2975 PyUnicode_Compare(name, newobj_str) == 0;
2976 Py_DECREF(name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002977 }
2978 }
2979 if (use_newobj) {
2980 PyObject *cls;
2981 PyObject *newargtup;
2982 PyObject *obj_class;
2983 int p;
2984
2985 /* Sanity checks. */
2986 if (Py_SIZE(argtup) < 1) {
2987 PyErr_SetString(PicklingError, "__newobj__ arglist is empty");
2988 return -1;
2989 }
2990
2991 cls = PyTuple_GET_ITEM(argtup, 0);
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002992 if (!PyType_Check(cls)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002993 PyErr_SetString(PicklingError, "args[0] from "
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002994 "__newobj__ args is not a type");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002995 return -1;
2996 }
2997
2998 if (obj != NULL) {
Antoine Pitrou16c4ce12011-03-11 21:30:43 +01002999 obj_class = get_class(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003000 p = obj_class != cls; /* true iff a problem */
3001 Py_DECREF(obj_class);
3002 if (p) {
3003 PyErr_SetString(PicklingError, "args[0] from "
3004 "__newobj__ args has the wrong class");
3005 return -1;
3006 }
3007 }
3008 /* XXX: These calls save() are prone to infinite recursion. Imagine
3009 what happen if the value returned by the __reduce__() method of
3010 some extension type contains another object of the same type. Ouch!
3011
3012 Here is a quick example, that I ran into, to illustrate what I
3013 mean:
3014
3015 >>> import pickle, copyreg
3016 >>> copyreg.dispatch_table.pop(complex)
3017 >>> pickle.dumps(1+2j)
3018 Traceback (most recent call last):
3019 ...
3020 RuntimeError: maximum recursion depth exceeded
3021
3022 Removing the complex class from copyreg.dispatch_table made the
3023 __reduce_ex__() method emit another complex object:
3024
3025 >>> (1+1j).__reduce_ex__(2)
3026 (<function __newobj__ at 0xb7b71c3c>,
3027 (<class 'complex'>, (1+1j)), None, None, None)
3028
3029 Thus when save() was called on newargstup (the 2nd item) recursion
3030 ensued. Of course, the bug was in the complex class which had a
3031 broken __getnewargs__() that emitted another complex object. But,
3032 the point, here, is it is quite easy to end up with a broken reduce
3033 function. */
3034
3035 /* Save the class and its __new__ arguments. */
3036 if (save(self, cls, 0) < 0)
3037 return -1;
3038
3039 newargtup = PyTuple_GetSlice(argtup, 1, Py_SIZE(argtup));
3040 if (newargtup == NULL)
3041 return -1;
3042
3043 p = save(self, newargtup, 0);
3044 Py_DECREF(newargtup);
3045 if (p < 0)
3046 return -1;
3047
3048 /* Add NEWOBJ opcode. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003049 if (_Pickler_Write(self, &newobj_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003050 return -1;
3051 }
3052 else { /* Not using NEWOBJ. */
3053 if (save(self, callable, 0) < 0 ||
3054 save(self, argtup, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003055 _Pickler_Write(self, &reduce_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003056 return -1;
3057 }
3058
3059 /* obj can be NULL when save_reduce() is used directly. A NULL obj means
3060 the caller do not want to memoize the object. Not particularly useful,
3061 but that is to mimic the behavior save_reduce() in pickle.py when
3062 obj is None. */
3063 if (obj && memo_put(self, obj) < 0)
3064 return -1;
3065
3066 if (listitems && batch_list(self, listitems) < 0)
3067 return -1;
3068
3069 if (dictitems && batch_dict(self, dictitems) < 0)
3070 return -1;
3071
3072 if (state) {
Victor Stinner121aab42011-09-29 23:40:53 +02003073 if (save(self, state, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003074 _Pickler_Write(self, &build_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003075 return -1;
3076 }
3077
3078 return 0;
3079}
3080
3081static int
3082save(PicklerObject *self, PyObject *obj, int pers_save)
3083{
3084 PyTypeObject *type;
3085 PyObject *reduce_func = NULL;
3086 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003087 int status = 0;
3088
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00003089 if (Py_EnterRecursiveCall(" while pickling an object"))
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003090 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003091
3092 /* The extra pers_save argument is necessary to avoid calling save_pers()
3093 on its returned object. */
3094 if (!pers_save && self->pers_func) {
3095 /* save_pers() returns:
3096 -1 to signal an error;
3097 0 if it did nothing successfully;
3098 1 if a persistent id was saved.
3099 */
3100 if ((status = save_pers(self, obj, self->pers_func)) != 0)
3101 goto done;
3102 }
3103
3104 type = Py_TYPE(obj);
3105
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003106 /* The old cPickle had an optimization that used switch-case statement
3107 dispatching on the first letter of the type name. This has was removed
3108 since benchmarks shown that this optimization was actually slowing
3109 things down. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003110
3111 /* Atom types; these aren't memoized, so don't check the memo. */
3112
3113 if (obj == Py_None) {
3114 status = save_none(self, obj);
3115 goto done;
3116 }
3117 else if (obj == Py_False || obj == Py_True) {
3118 status = save_bool(self, obj);
3119 goto done;
3120 }
3121 else if (type == &PyLong_Type) {
3122 status = save_long(self, obj);
3123 goto done;
3124 }
3125 else if (type == &PyFloat_Type) {
3126 status = save_float(self, obj);
3127 goto done;
3128 }
3129
3130 /* Check the memo to see if it has the object. If so, generate
3131 a GET (or BINGET) opcode, instead of pickling the object
3132 once again. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003133 if (PyMemoTable_Get(self->memo, obj)) {
3134 if (memo_get(self, obj) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003135 goto error;
3136 goto done;
3137 }
3138
3139 if (type == &PyBytes_Type) {
3140 status = save_bytes(self, obj);
3141 goto done;
3142 }
3143 else if (type == &PyUnicode_Type) {
3144 status = save_unicode(self, obj);
3145 goto done;
3146 }
3147 else if (type == &PyDict_Type) {
3148 status = save_dict(self, obj);
3149 goto done;
3150 }
3151 else if (type == &PyList_Type) {
3152 status = save_list(self, obj);
3153 goto done;
3154 }
3155 else if (type == &PyTuple_Type) {
3156 status = save_tuple(self, obj);
3157 goto done;
3158 }
3159 else if (type == &PyType_Type) {
3160 status = save_global(self, obj, NULL);
3161 goto done;
3162 }
3163 else if (type == &PyFunction_Type) {
3164 status = save_global(self, obj, NULL);
3165 if (status < 0 && PyErr_ExceptionMatches(PickleError)) {
3166 /* fall back to reduce */
3167 PyErr_Clear();
3168 }
3169 else {
3170 goto done;
3171 }
3172 }
3173 else if (type == &PyCFunction_Type) {
3174 status = save_global(self, obj, NULL);
3175 goto done;
3176 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003177
3178 /* XXX: This part needs some unit tests. */
3179
3180 /* Get a reduction callable, and call it. This may come from
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003181 * self.dispatch_table, copyreg.dispatch_table, the object's
3182 * __reduce_ex__ method, or the object's __reduce__ method.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003183 */
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003184 if (self->dispatch_table == NULL) {
3185 reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
3186 /* PyDict_GetItem() unlike PyObject_GetItem() and
3187 PyObject_GetAttr() returns a borrowed ref */
3188 Py_XINCREF(reduce_func);
3189 } else {
3190 reduce_func = PyObject_GetItem(self->dispatch_table, (PyObject *)type);
3191 if (reduce_func == NULL) {
3192 if (PyErr_ExceptionMatches(PyExc_KeyError))
3193 PyErr_Clear();
3194 else
3195 goto error;
3196 }
3197 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003198 if (reduce_func != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003199 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003200 reduce_value = _Pickler_FastCall(self, reduce_func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003201 }
Antoine Pitrouffd41d92011-10-04 09:23:04 +02003202 else if (PyType_IsSubtype(type, &PyType_Type)) {
3203 status = save_global(self, obj, NULL);
3204 goto done;
3205 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003206 else {
3207 static PyObject *reduce_str = NULL;
3208 static PyObject *reduce_ex_str = NULL;
3209
3210 /* Cache the name of the reduce methods. */
3211 if (reduce_str == NULL) {
3212 reduce_str = PyUnicode_InternFromString("__reduce__");
3213 if (reduce_str == NULL)
3214 goto error;
3215 reduce_ex_str = PyUnicode_InternFromString("__reduce_ex__");
3216 if (reduce_ex_str == NULL)
3217 goto error;
3218 }
3219
3220 /* XXX: If the __reduce__ method is defined, __reduce_ex__ is
3221 automatically defined as __reduce__. While this is convenient, this
3222 make it impossible to know which method was actually called. Of
3223 course, this is not a big deal. But still, it would be nice to let
3224 the user know which method was called when something go
3225 wrong. Incidentally, this means if __reduce_ex__ is not defined, we
3226 don't actually have to check for a __reduce__ method. */
3227
3228 /* Check for a __reduce_ex__ method. */
3229 reduce_func = PyObject_GetAttr(obj, reduce_ex_str);
3230 if (reduce_func != NULL) {
3231 PyObject *proto;
3232 proto = PyLong_FromLong(self->proto);
3233 if (proto != NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003234 reduce_value = _Pickler_FastCall(self, reduce_func, proto);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003235 }
3236 }
3237 else {
3238 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3239 PyErr_Clear();
3240 else
3241 goto error;
3242 /* Check for a __reduce__ method. */
3243 reduce_func = PyObject_GetAttr(obj, reduce_str);
3244 if (reduce_func != NULL) {
3245 reduce_value = PyObject_Call(reduce_func, empty_tuple, NULL);
3246 }
3247 else {
3248 PyErr_Format(PicklingError, "can't pickle '%.200s' object: %R",
3249 type->tp_name, obj);
3250 goto error;
3251 }
3252 }
3253 }
3254
3255 if (reduce_value == NULL)
3256 goto error;
3257
3258 if (PyUnicode_Check(reduce_value)) {
3259 status = save_global(self, obj, reduce_value);
3260 goto done;
3261 }
3262
3263 if (!PyTuple_Check(reduce_value)) {
3264 PyErr_SetString(PicklingError,
3265 "__reduce__ must return a string or tuple");
3266 goto error;
3267 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003268
3269 status = save_reduce(self, reduce_value, obj);
3270
3271 if (0) {
3272 error:
3273 status = -1;
3274 }
3275 done:
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003276 Py_LeaveRecursiveCall();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003277 Py_XDECREF(reduce_func);
3278 Py_XDECREF(reduce_value);
3279
3280 return status;
3281}
3282
3283static int
3284dump(PicklerObject *self, PyObject *obj)
3285{
3286 const char stop_op = STOP;
3287
3288 if (self->proto >= 2) {
3289 char header[2];
3290
3291 header[0] = PROTO;
3292 assert(self->proto >= 0 && self->proto < 256);
3293 header[1] = (unsigned char)self->proto;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003294 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003295 return -1;
3296 }
3297
3298 if (save(self, obj, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003299 _Pickler_Write(self, &stop_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003300 return -1;
3301
3302 return 0;
3303}
3304
3305PyDoc_STRVAR(Pickler_clear_memo_doc,
3306"clear_memo() -> None. Clears the pickler's \"memo\"."
3307"\n"
3308"The memo is the data structure that remembers which objects the\n"
3309"pickler has already seen, so that shared or recursive objects are\n"
3310"pickled by reference and not by value. This method is useful when\n"
3311"re-using picklers.");
3312
3313static PyObject *
3314Pickler_clear_memo(PicklerObject *self)
3315{
3316 if (self->memo)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003317 PyMemoTable_Clear(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003318
3319 Py_RETURN_NONE;
3320}
3321
3322PyDoc_STRVAR(Pickler_dump_doc,
3323"dump(obj) -> None. Write a pickled representation of obj to the open file.");
3324
3325static PyObject *
3326Pickler_dump(PicklerObject *self, PyObject *args)
3327{
3328 PyObject *obj;
3329
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003330 /* Check whether the Pickler was initialized correctly (issue3664).
3331 Developers often forget to call __init__() in their subclasses, which
3332 would trigger a segfault without this check. */
3333 if (self->write == NULL) {
Victor Stinner121aab42011-09-29 23:40:53 +02003334 PyErr_Format(PicklingError,
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003335 "Pickler.__init__() was not called by %s.__init__()",
3336 Py_TYPE(self)->tp_name);
3337 return NULL;
3338 }
3339
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003340 if (!PyArg_ParseTuple(args, "O:dump", &obj))
3341 return NULL;
3342
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003343 if (_Pickler_ClearBuffer(self) < 0)
3344 return NULL;
3345
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003346 if (dump(self, obj) < 0)
3347 return NULL;
3348
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003349 if (_Pickler_FlushToFile(self) < 0)
3350 return NULL;
3351
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003352 Py_RETURN_NONE;
3353}
3354
3355static struct PyMethodDef Pickler_methods[] = {
3356 {"dump", (PyCFunction)Pickler_dump, METH_VARARGS,
3357 Pickler_dump_doc},
3358 {"clear_memo", (PyCFunction)Pickler_clear_memo, METH_NOARGS,
3359 Pickler_clear_memo_doc},
3360 {NULL, NULL} /* sentinel */
3361};
3362
3363static void
3364Pickler_dealloc(PicklerObject *self)
3365{
3366 PyObject_GC_UnTrack(self);
3367
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003368 Py_XDECREF(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003369 Py_XDECREF(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003370 Py_XDECREF(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003371 Py_XDECREF(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003372 Py_XDECREF(self->arg);
3373 Py_XDECREF(self->fast_memo);
3374
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003375 PyMemoTable_Del(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003376
3377 Py_TYPE(self)->tp_free((PyObject *)self);
3378}
3379
3380static int
3381Pickler_traverse(PicklerObject *self, visitproc visit, void *arg)
3382{
3383 Py_VISIT(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003384 Py_VISIT(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003385 Py_VISIT(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003386 Py_VISIT(self->arg);
3387 Py_VISIT(self->fast_memo);
3388 return 0;
3389}
3390
3391static int
3392Pickler_clear(PicklerObject *self)
3393{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003394 Py_CLEAR(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003395 Py_CLEAR(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003396 Py_CLEAR(self->pers_func);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003397 Py_CLEAR(self->dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003398 Py_CLEAR(self->arg);
3399 Py_CLEAR(self->fast_memo);
3400
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003401 if (self->memo != NULL) {
3402 PyMemoTable *memo = self->memo;
3403 self->memo = NULL;
3404 PyMemoTable_Del(memo);
3405 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003406 return 0;
3407}
3408
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003409
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003410PyDoc_STRVAR(Pickler_doc,
3411"Pickler(file, protocol=None)"
3412"\n"
3413"This takes a binary file for writing a pickle data stream.\n"
3414"\n"
3415"The optional protocol argument tells the pickler to use the\n"
3416"given protocol; supported protocols are 0, 1, 2, 3. The default\n"
3417"protocol is 3; a backward-incompatible protocol designed for\n"
3418"Python 3.0.\n"
3419"\n"
3420"Specifying a negative protocol version selects the highest\n"
3421"protocol version supported. The higher the protocol used, the\n"
3422"more recent the version of Python needed to read the pickle\n"
3423"produced.\n"
3424"\n"
3425"The file argument must have a write() method that accepts a single\n"
3426"bytes argument. It can thus be a file object opened for binary\n"
3427"writing, a io.BytesIO instance, or any other custom object that\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003428"meets this interface.\n"
3429"\n"
3430"If fix_imports is True and protocol is less than 3, pickle will try to\n"
3431"map the new Python 3.x names to the old module names used in Python\n"
3432"2.x, so that the pickle data stream is readable with Python 2.x.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003433
3434static int
3435Pickler_init(PicklerObject *self, PyObject *args, PyObject *kwds)
3436{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003437 static char *kwlist[] = {"file", "protocol", "fix_imports", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003438 PyObject *file;
3439 PyObject *proto_obj = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003440 PyObject *fix_imports = Py_True;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02003441 _Py_IDENTIFIER(persistent_id);
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003442 _Py_IDENTIFIER(dispatch_table);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003443
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003444 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:Pickler",
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003445 kwlist, &file, &proto_obj, &fix_imports))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003446 return -1;
3447
3448 /* In case of multiple __init__() calls, clear previous content. */
3449 if (self->write != NULL)
3450 (void)Pickler_clear(self);
3451
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003452 if (_Pickler_SetProtocol(self, proto_obj, fix_imports) < 0)
3453 return -1;
3454
3455 if (_Pickler_SetOutputStream(self, file) < 0)
3456 return -1;
3457
3458 /* memo and output_buffer may have already been created in _Pickler_New */
3459 if (self->memo == NULL) {
3460 self->memo = PyMemoTable_New();
3461 if (self->memo == NULL)
3462 return -1;
3463 }
3464 self->output_len = 0;
3465 if (self->output_buffer == NULL) {
3466 self->max_output_len = WRITE_BUF_SIZE;
3467 self->output_buffer = PyBytes_FromStringAndSize(NULL,
3468 self->max_output_len);
3469 if (self->output_buffer == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003470 return -1;
3471 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003472
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003473 self->arg = NULL;
3474 self->fast = 0;
3475 self->fast_nesting = 0;
3476 self->fast_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003477 self->pers_func = NULL;
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02003478 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_id)) {
3479 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
3480 &PyId_persistent_id);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003481 if (self->pers_func == NULL)
3482 return -1;
3483 }
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003484 self->dispatch_table = NULL;
3485 if (_PyObject_HasAttrId((PyObject *)self, &PyId_dispatch_table)) {
3486 self->dispatch_table = _PyObject_GetAttrId((PyObject *)self,
3487 &PyId_dispatch_table);
3488 if (self->dispatch_table == NULL)
3489 return -1;
3490 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003491 return 0;
3492}
3493
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003494/* Define a proxy object for the Pickler's internal memo object. This is to
3495 * avoid breaking code like:
3496 * pickler.memo.clear()
3497 * and
3498 * pickler.memo = saved_memo
3499 * Is this a good idea? Not really, but we don't want to break code that uses
3500 * it. Note that we don't implement the entire mapping API here. This is
3501 * intentional, as these should be treated as black-box implementation details.
3502 */
3503
3504typedef struct {
3505 PyObject_HEAD
3506 PicklerObject *pickler; /* Pickler whose memo table we're proxying. */
3507} PicklerMemoProxyObject;
3508
3509PyDoc_STRVAR(pmp_clear_doc,
3510"memo.clear() -> None. Remove all items from memo.");
3511
3512static PyObject *
3513pmp_clear(PicklerMemoProxyObject *self)
3514{
3515 if (self->pickler->memo)
3516 PyMemoTable_Clear(self->pickler->memo);
3517 Py_RETURN_NONE;
3518}
3519
3520PyDoc_STRVAR(pmp_copy_doc,
3521"memo.copy() -> new_memo. Copy the memo to a new object.");
3522
3523static PyObject *
3524pmp_copy(PicklerMemoProxyObject *self)
3525{
3526 Py_ssize_t i;
3527 PyMemoTable *memo;
3528 PyObject *new_memo = PyDict_New();
3529 if (new_memo == NULL)
3530 return NULL;
3531
3532 memo = self->pickler->memo;
3533 for (i = 0; i < memo->mt_allocated; ++i) {
3534 PyMemoEntry entry = memo->mt_table[i];
3535 if (entry.me_key != NULL) {
3536 int status;
3537 PyObject *key, *value;
3538
3539 key = PyLong_FromVoidPtr(entry.me_key);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003540 value = Py_BuildValue("nO", entry.me_value, entry.me_key);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003541
3542 if (key == NULL || value == NULL) {
3543 Py_XDECREF(key);
3544 Py_XDECREF(value);
3545 goto error;
3546 }
3547 status = PyDict_SetItem(new_memo, key, value);
3548 Py_DECREF(key);
3549 Py_DECREF(value);
3550 if (status < 0)
3551 goto error;
3552 }
3553 }
3554 return new_memo;
3555
3556 error:
3557 Py_XDECREF(new_memo);
3558 return NULL;
3559}
3560
3561PyDoc_STRVAR(pmp_reduce_doc,
3562"memo.__reduce__(). Pickling support.");
3563
3564static PyObject *
3565pmp_reduce(PicklerMemoProxyObject *self, PyObject *args)
3566{
3567 PyObject *reduce_value, *dict_args;
3568 PyObject *contents = pmp_copy(self);
3569 if (contents == NULL)
3570 return NULL;
3571
3572 reduce_value = PyTuple_New(2);
3573 if (reduce_value == NULL) {
3574 Py_DECREF(contents);
3575 return NULL;
3576 }
3577 dict_args = PyTuple_New(1);
3578 if (dict_args == NULL) {
3579 Py_DECREF(contents);
3580 Py_DECREF(reduce_value);
3581 return NULL;
3582 }
3583 PyTuple_SET_ITEM(dict_args, 0, contents);
3584 Py_INCREF((PyObject *)&PyDict_Type);
3585 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
3586 PyTuple_SET_ITEM(reduce_value, 1, dict_args);
3587 return reduce_value;
3588}
3589
3590static PyMethodDef picklerproxy_methods[] = {
3591 {"clear", (PyCFunction)pmp_clear, METH_NOARGS, pmp_clear_doc},
3592 {"copy", (PyCFunction)pmp_copy, METH_NOARGS, pmp_copy_doc},
3593 {"__reduce__", (PyCFunction)pmp_reduce, METH_VARARGS, pmp_reduce_doc},
3594 {NULL, NULL} /* sentinel */
3595};
3596
3597static void
3598PicklerMemoProxy_dealloc(PicklerMemoProxyObject *self)
3599{
3600 PyObject_GC_UnTrack(self);
3601 Py_XDECREF(self->pickler);
3602 PyObject_GC_Del((PyObject *)self);
3603}
3604
3605static int
3606PicklerMemoProxy_traverse(PicklerMemoProxyObject *self,
3607 visitproc visit, void *arg)
3608{
3609 Py_VISIT(self->pickler);
3610 return 0;
3611}
3612
3613static int
3614PicklerMemoProxy_clear(PicklerMemoProxyObject *self)
3615{
3616 Py_CLEAR(self->pickler);
3617 return 0;
3618}
3619
3620static PyTypeObject PicklerMemoProxyType = {
3621 PyVarObject_HEAD_INIT(NULL, 0)
3622 "_pickle.PicklerMemoProxy", /*tp_name*/
3623 sizeof(PicklerMemoProxyObject), /*tp_basicsize*/
3624 0,
3625 (destructor)PicklerMemoProxy_dealloc, /* tp_dealloc */
3626 0, /* tp_print */
3627 0, /* tp_getattr */
3628 0, /* tp_setattr */
3629 0, /* tp_compare */
3630 0, /* tp_repr */
3631 0, /* tp_as_number */
3632 0, /* tp_as_sequence */
3633 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00003634 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003635 0, /* tp_call */
3636 0, /* tp_str */
3637 PyObject_GenericGetAttr, /* tp_getattro */
3638 PyObject_GenericSetAttr, /* tp_setattro */
3639 0, /* tp_as_buffer */
3640 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3641 0, /* tp_doc */
3642 (traverseproc)PicklerMemoProxy_traverse, /* tp_traverse */
3643 (inquiry)PicklerMemoProxy_clear, /* tp_clear */
3644 0, /* tp_richcompare */
3645 0, /* tp_weaklistoffset */
3646 0, /* tp_iter */
3647 0, /* tp_iternext */
3648 picklerproxy_methods, /* tp_methods */
3649};
3650
3651static PyObject *
3652PicklerMemoProxy_New(PicklerObject *pickler)
3653{
3654 PicklerMemoProxyObject *self;
3655
3656 self = PyObject_GC_New(PicklerMemoProxyObject, &PicklerMemoProxyType);
3657 if (self == NULL)
3658 return NULL;
3659 Py_INCREF(pickler);
3660 self->pickler = pickler;
3661 PyObject_GC_Track(self);
3662 return (PyObject *)self;
3663}
3664
3665/*****************************************************************************/
3666
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003667static PyObject *
3668Pickler_get_memo(PicklerObject *self)
3669{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003670 return PicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003671}
3672
3673static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003674Pickler_set_memo(PicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003675{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003676 PyMemoTable *new_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003677
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003678 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003679 PyErr_SetString(PyExc_TypeError,
3680 "attribute deletion is not supported");
3681 return -1;
3682 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003683
3684 if (Py_TYPE(obj) == &PicklerMemoProxyType) {
3685 PicklerObject *pickler =
3686 ((PicklerMemoProxyObject *)obj)->pickler;
3687
3688 new_memo = PyMemoTable_Copy(pickler->memo);
3689 if (new_memo == NULL)
3690 return -1;
3691 }
3692 else if (PyDict_Check(obj)) {
3693 Py_ssize_t i = 0;
3694 PyObject *key, *value;
3695
3696 new_memo = PyMemoTable_New();
3697 if (new_memo == NULL)
3698 return -1;
3699
3700 while (PyDict_Next(obj, &i, &key, &value)) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003701 Py_ssize_t memo_id;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003702 PyObject *memo_obj;
3703
3704 if (!PyTuple_Check(value) || Py_SIZE(value) != 2) {
3705 PyErr_SetString(PyExc_TypeError,
3706 "'memo' values must be 2-item tuples");
3707 goto error;
3708 }
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003709 memo_id = PyLong_AsSsize_t(PyTuple_GET_ITEM(value, 0));
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003710 if (memo_id == -1 && PyErr_Occurred())
3711 goto error;
3712 memo_obj = PyTuple_GET_ITEM(value, 1);
3713 if (PyMemoTable_Set(new_memo, memo_obj, memo_id) < 0)
3714 goto error;
3715 }
3716 }
3717 else {
3718 PyErr_Format(PyExc_TypeError,
3719 "'memo' attribute must be an PicklerMemoProxy object"
3720 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003721 return -1;
3722 }
3723
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003724 PyMemoTable_Del(self->memo);
3725 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003726
3727 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003728
3729 error:
3730 if (new_memo)
3731 PyMemoTable_Del(new_memo);
3732 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003733}
3734
3735static PyObject *
3736Pickler_get_persid(PicklerObject *self)
3737{
3738 if (self->pers_func == NULL)
3739 PyErr_SetString(PyExc_AttributeError, "persistent_id");
3740 else
3741 Py_INCREF(self->pers_func);
3742 return self->pers_func;
3743}
3744
3745static int
3746Pickler_set_persid(PicklerObject *self, PyObject *value)
3747{
3748 PyObject *tmp;
3749
3750 if (value == NULL) {
3751 PyErr_SetString(PyExc_TypeError,
3752 "attribute deletion is not supported");
3753 return -1;
3754 }
3755 if (!PyCallable_Check(value)) {
3756 PyErr_SetString(PyExc_TypeError,
3757 "persistent_id must be a callable taking one argument");
3758 return -1;
3759 }
3760
3761 tmp = self->pers_func;
3762 Py_INCREF(value);
3763 self->pers_func = value;
3764 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
3765
3766 return 0;
3767}
3768
3769static PyMemberDef Pickler_members[] = {
3770 {"bin", T_INT, offsetof(PicklerObject, bin)},
3771 {"fast", T_INT, offsetof(PicklerObject, fast)},
Antoine Pitrou8d3c2902012-03-04 18:31:48 +01003772 {"dispatch_table", T_OBJECT_EX, offsetof(PicklerObject, dispatch_table)},
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003773 {NULL}
3774};
3775
3776static PyGetSetDef Pickler_getsets[] = {
3777 {"memo", (getter)Pickler_get_memo,
3778 (setter)Pickler_set_memo},
3779 {"persistent_id", (getter)Pickler_get_persid,
3780 (setter)Pickler_set_persid},
3781 {NULL}
3782};
3783
3784static PyTypeObject Pickler_Type = {
3785 PyVarObject_HEAD_INIT(NULL, 0)
3786 "_pickle.Pickler" , /*tp_name*/
3787 sizeof(PicklerObject), /*tp_basicsize*/
3788 0, /*tp_itemsize*/
3789 (destructor)Pickler_dealloc, /*tp_dealloc*/
3790 0, /*tp_print*/
3791 0, /*tp_getattr*/
3792 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00003793 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003794 0, /*tp_repr*/
3795 0, /*tp_as_number*/
3796 0, /*tp_as_sequence*/
3797 0, /*tp_as_mapping*/
3798 0, /*tp_hash*/
3799 0, /*tp_call*/
3800 0, /*tp_str*/
3801 0, /*tp_getattro*/
3802 0, /*tp_setattro*/
3803 0, /*tp_as_buffer*/
3804 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3805 Pickler_doc, /*tp_doc*/
3806 (traverseproc)Pickler_traverse, /*tp_traverse*/
3807 (inquiry)Pickler_clear, /*tp_clear*/
3808 0, /*tp_richcompare*/
3809 0, /*tp_weaklistoffset*/
3810 0, /*tp_iter*/
3811 0, /*tp_iternext*/
3812 Pickler_methods, /*tp_methods*/
3813 Pickler_members, /*tp_members*/
3814 Pickler_getsets, /*tp_getset*/
3815 0, /*tp_base*/
3816 0, /*tp_dict*/
3817 0, /*tp_descr_get*/
3818 0, /*tp_descr_set*/
3819 0, /*tp_dictoffset*/
3820 (initproc)Pickler_init, /*tp_init*/
3821 PyType_GenericAlloc, /*tp_alloc*/
3822 PyType_GenericNew, /*tp_new*/
3823 PyObject_GC_Del, /*tp_free*/
3824 0, /*tp_is_gc*/
3825};
3826
Victor Stinner121aab42011-09-29 23:40:53 +02003827/* Temporary helper for calling self.find_class().
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003828
3829 XXX: It would be nice to able to avoid Python function call overhead, by
3830 using directly the C version of find_class(), when find_class() is not
3831 overridden by a subclass. Although, this could become rather hackish. A
3832 simpler optimization would be to call the C function when self is not a
3833 subclass instance. */
3834static PyObject *
3835find_class(UnpicklerObject *self, PyObject *module_name, PyObject *global_name)
3836{
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02003837 _Py_IDENTIFIER(find_class);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02003838
3839 return _PyObject_CallMethodId((PyObject *)self, &PyId_find_class, "OO",
3840 module_name, global_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003841}
3842
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003843static Py_ssize_t
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003844marker(UnpicklerObject *self)
3845{
3846 if (self->num_marks < 1) {
3847 PyErr_SetString(UnpicklingError, "could not find MARK");
3848 return -1;
3849 }
3850
3851 return self->marks[--self->num_marks];
3852}
3853
3854static int
3855load_none(UnpicklerObject *self)
3856{
3857 PDATA_APPEND(self->stack, Py_None, -1);
3858 return 0;
3859}
3860
3861static int
3862bad_readline(void)
3863{
3864 PyErr_SetString(UnpicklingError, "pickle data was truncated");
3865 return -1;
3866}
3867
3868static int
3869load_int(UnpicklerObject *self)
3870{
3871 PyObject *value;
3872 char *endptr, *s;
3873 Py_ssize_t len;
3874 long x;
3875
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003876 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003877 return -1;
3878 if (len < 2)
3879 return bad_readline();
3880
3881 errno = 0;
Victor Stinner121aab42011-09-29 23:40:53 +02003882 /* XXX: Should the base argument of strtol() be explicitly set to 10?
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003883 XXX(avassalotti): Should this uses PyOS_strtol()? */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003884 x = strtol(s, &endptr, 0);
3885
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003886 if (errno || (*endptr != '\n' && *endptr != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003887 /* Hm, maybe we've got something long. Let's try reading
3888 * it as a Python long object. */
3889 errno = 0;
3890 /* XXX: Same thing about the base here. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003891 value = PyLong_FromString(s, NULL, 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003892 if (value == NULL) {
3893 PyErr_SetString(PyExc_ValueError,
3894 "could not convert string to int");
3895 return -1;
3896 }
3897 }
3898 else {
3899 if (len == 3 && (x == 0 || x == 1)) {
3900 if ((value = PyBool_FromLong(x)) == NULL)
3901 return -1;
3902 }
3903 else {
3904 if ((value = PyLong_FromLong(x)) == NULL)
3905 return -1;
3906 }
3907 }
3908
3909 PDATA_PUSH(self->stack, value, -1);
3910 return 0;
3911}
3912
3913static int
3914load_bool(UnpicklerObject *self, PyObject *boolean)
3915{
3916 assert(boolean == Py_True || boolean == Py_False);
3917 PDATA_APPEND(self->stack, boolean, -1);
3918 return 0;
3919}
3920
Antoine Pitrou82be19f2011-08-29 23:09:33 +02003921/* s contains x bytes of an unsigned little-endian integer. Return its value
3922 * as a C Py_ssize_t, or -1 if it's higher than PY_SSIZE_T_MAX.
3923 */
3924static Py_ssize_t
3925calc_binsize(char *bytes, int size)
3926{
3927 unsigned char *s = (unsigned char *)bytes;
3928 size_t x = 0;
3929
3930 assert(size == 4);
3931
3932 x = (size_t) s[0];
3933 x |= (size_t) s[1] << 8;
3934 x |= (size_t) s[2] << 16;
3935 x |= (size_t) s[3] << 24;
3936
3937 if (x > PY_SSIZE_T_MAX)
3938 return -1;
3939 else
3940 return (Py_ssize_t) x;
3941}
3942
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003943/* s contains x bytes of a little-endian integer. Return its value as a
3944 * C int. Obscure: when x is 1 or 2, this is an unsigned little-endian
3945 * int, but when x is 4 it's a signed one. This is an historical source
3946 * of x-platform bugs.
3947 */
3948static long
3949calc_binint(char *bytes, int size)
3950{
3951 unsigned char *s = (unsigned char *)bytes;
3952 int i = size;
3953 long x = 0;
3954
3955 for (i = 0; i < size; i++) {
3956 x |= (long)s[i] << (i * 8);
3957 }
3958
3959 /* Unlike BININT1 and BININT2, BININT (more accurately BININT4)
3960 * is signed, so on a box with longs bigger than 4 bytes we need
3961 * to extend a BININT's sign bit to the full width.
3962 */
3963 if (SIZEOF_LONG > 4 && size == 4) {
3964 x |= -(x & (1L << 31));
3965 }
3966
3967 return x;
3968}
3969
3970static int
3971load_binintx(UnpicklerObject *self, char *s, int size)
3972{
3973 PyObject *value;
3974 long x;
3975
3976 x = calc_binint(s, size);
3977
3978 if ((value = PyLong_FromLong(x)) == NULL)
3979 return -1;
3980
3981 PDATA_PUSH(self->stack, value, -1);
3982 return 0;
3983}
3984
3985static int
3986load_binint(UnpicklerObject *self)
3987{
3988 char *s;
3989
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003990 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003991 return -1;
3992
3993 return load_binintx(self, s, 4);
3994}
3995
3996static int
3997load_binint1(UnpicklerObject *self)
3998{
3999 char *s;
4000
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004001 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004002 return -1;
4003
4004 return load_binintx(self, s, 1);
4005}
4006
4007static int
4008load_binint2(UnpicklerObject *self)
4009{
4010 char *s;
4011
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004012 if (_Unpickler_Read(self, &s, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004013 return -1;
4014
4015 return load_binintx(self, s, 2);
4016}
4017
4018static int
4019load_long(UnpicklerObject *self)
4020{
4021 PyObject *value;
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004022 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004023 Py_ssize_t len;
4024
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004025 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004026 return -1;
4027 if (len < 2)
4028 return bad_readline();
4029
Mark Dickinson8dd05142009-01-20 20:43:58 +00004030 /* s[len-2] will usually be 'L' (and s[len-1] is '\n'); we need to remove
4031 the 'L' before calling PyLong_FromString. In order to maintain
4032 compatibility with Python 3.0.0, we don't actually *require*
4033 the 'L' to be present. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004034 if (s[len-2] == 'L')
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00004035 s[len-2] = '\0';
Alexandre Vassalottie4bccb72009-01-24 01:47:57 +00004036 /* XXX: Should the base argument explicitly set to 10? */
4037 value = PyLong_FromString(s, NULL, 0);
Mark Dickinson8dd05142009-01-20 20:43:58 +00004038 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004039 return -1;
4040
4041 PDATA_PUSH(self->stack, value, -1);
4042 return 0;
4043}
4044
4045/* 'size' bytes contain the # of bytes of little-endian 256's-complement
4046 * data following.
4047 */
4048static int
4049load_counted_long(UnpicklerObject *self, int size)
4050{
4051 PyObject *value;
4052 char *nbytes;
4053 char *pdata;
4054
4055 assert(size == 1 || size == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004056 if (_Unpickler_Read(self, &nbytes, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004057 return -1;
4058
4059 size = calc_binint(nbytes, size);
4060 if (size < 0) {
4061 /* Corrupt or hostile pickle -- we never write one like this */
4062 PyErr_SetString(UnpicklingError,
4063 "LONG pickle has negative byte count");
4064 return -1;
4065 }
4066
4067 if (size == 0)
4068 value = PyLong_FromLong(0L);
4069 else {
4070 /* Read the raw little-endian bytes and convert. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004071 if (_Unpickler_Read(self, &pdata, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004072 return -1;
4073 value = _PyLong_FromByteArray((unsigned char *)pdata, (size_t)size,
4074 1 /* little endian */ , 1 /* signed */ );
4075 }
4076 if (value == NULL)
4077 return -1;
4078 PDATA_PUSH(self->stack, value, -1);
4079 return 0;
4080}
4081
4082static int
4083load_float(UnpicklerObject *self)
4084{
4085 PyObject *value;
4086 char *endptr, *s;
4087 Py_ssize_t len;
4088 double d;
4089
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004090 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004091 return -1;
4092 if (len < 2)
4093 return bad_readline();
4094
4095 errno = 0;
Mark Dickinson725bfd82009-05-03 20:33:40 +00004096 d = PyOS_string_to_double(s, &endptr, PyExc_OverflowError);
4097 if (d == -1.0 && PyErr_Occurred())
4098 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004099 if ((endptr[0] != '\n') && (endptr[0] != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004100 PyErr_SetString(PyExc_ValueError, "could not convert string to float");
4101 return -1;
4102 }
Mark Dickinson725bfd82009-05-03 20:33:40 +00004103 value = PyFloat_FromDouble(d);
4104 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004105 return -1;
4106
4107 PDATA_PUSH(self->stack, value, -1);
4108 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004109}
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004110
4111static int
4112load_binfloat(UnpicklerObject *self)
4113{
4114 PyObject *value;
4115 double x;
4116 char *s;
4117
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004118 if (_Unpickler_Read(self, &s, 8) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004119 return -1;
4120
4121 x = _PyFloat_Unpack8((unsigned char *)s, 0);
4122 if (x == -1.0 && PyErr_Occurred())
4123 return -1;
4124
4125 if ((value = PyFloat_FromDouble(x)) == NULL)
4126 return -1;
4127
4128 PDATA_PUSH(self->stack, value, -1);
4129 return 0;
4130}
4131
4132static int
4133load_string(UnpicklerObject *self)
4134{
4135 PyObject *bytes;
4136 PyObject *str = NULL;
4137 Py_ssize_t len;
4138 char *s, *p;
4139
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004140 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004141 return -1;
4142 if (len < 3)
4143 return bad_readline();
4144 if ((s = strdup(s)) == NULL) {
4145 PyErr_NoMemory();
4146 return -1;
4147 }
4148
4149 /* Strip outermost quotes */
4150 while (s[len - 1] <= ' ')
4151 len--;
4152 if (s[0] == '"' && s[len - 1] == '"') {
4153 s[len - 1] = '\0';
4154 p = s + 1;
4155 len -= 2;
4156 }
4157 else if (s[0] == '\'' && s[len - 1] == '\'') {
4158 s[len - 1] = '\0';
4159 p = s + 1;
4160 len -= 2;
4161 }
4162 else {
4163 free(s);
4164 PyErr_SetString(PyExc_ValueError, "insecure string pickle");
4165 return -1;
4166 }
4167
4168 /* Use the PyBytes API to decode the string, since that is what is used
4169 to encode, and then coerce the result to Unicode. */
4170 bytes = PyBytes_DecodeEscape(p, len, NULL, 0, NULL);
4171 free(s);
4172 if (bytes == NULL)
4173 return -1;
4174 str = PyUnicode_FromEncodedObject(bytes, self->encoding, self->errors);
4175 Py_DECREF(bytes);
4176 if (str == NULL)
4177 return -1;
4178
4179 PDATA_PUSH(self->stack, str, -1);
4180 return 0;
4181}
4182
4183static int
4184load_binbytes(UnpicklerObject *self)
4185{
4186 PyObject *bytes;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004187 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004188 char *s;
4189
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004190 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004191 return -1;
4192
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004193 x = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004194 if (x < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004195 PyErr_Format(PyExc_OverflowError,
4196 "BINBYTES exceeds system's maximum size of %zd bytes",
4197 PY_SSIZE_T_MAX
4198 );
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004199 return -1;
4200 }
4201
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004202 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004203 return -1;
4204 bytes = PyBytes_FromStringAndSize(s, x);
4205 if (bytes == NULL)
4206 return -1;
4207
4208 PDATA_PUSH(self->stack, bytes, -1);
4209 return 0;
4210}
4211
4212static int
4213load_short_binbytes(UnpicklerObject *self)
4214{
4215 PyObject *bytes;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004216 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004217 char *s;
4218
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004219 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004220 return -1;
4221
4222 x = (unsigned char)s[0];
4223
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004224 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004225 return -1;
4226
4227 bytes = PyBytes_FromStringAndSize(s, x);
4228 if (bytes == NULL)
4229 return -1;
4230
4231 PDATA_PUSH(self->stack, bytes, -1);
4232 return 0;
4233}
4234
4235static int
4236load_binstring(UnpicklerObject *self)
4237{
4238 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004239 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004240 char *s;
4241
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004242 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004243 return -1;
4244
4245 x = calc_binint(s, 4);
4246 if (x < 0) {
Victor Stinner121aab42011-09-29 23:40:53 +02004247 PyErr_SetString(UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004248 "BINSTRING pickle has negative byte count");
4249 return -1;
4250 }
4251
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004252 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004253 return -1;
4254
4255 /* Convert Python 2.x strings to unicode. */
4256 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4257 if (str == NULL)
4258 return -1;
4259
4260 PDATA_PUSH(self->stack, str, -1);
4261 return 0;
4262}
4263
4264static int
4265load_short_binstring(UnpicklerObject *self)
4266{
4267 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004268 Py_ssize_t x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004269 char *s;
4270
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004271 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004272 return -1;
4273
4274 x = (unsigned char)s[0];
4275
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004276 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004277 return -1;
4278
4279 /* Convert Python 2.x strings to unicode. */
4280 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4281 if (str == NULL)
4282 return -1;
4283
4284 PDATA_PUSH(self->stack, str, -1);
4285 return 0;
4286}
4287
4288static int
4289load_unicode(UnpicklerObject *self)
4290{
4291 PyObject *str;
4292 Py_ssize_t len;
4293 char *s;
4294
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004295 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004296 return -1;
4297 if (len < 1)
4298 return bad_readline();
4299
4300 str = PyUnicode_DecodeRawUnicodeEscape(s, len - 1, NULL);
4301 if (str == NULL)
4302 return -1;
4303
4304 PDATA_PUSH(self->stack, str, -1);
4305 return 0;
4306}
4307
4308static int
4309load_binunicode(UnpicklerObject *self)
4310{
4311 PyObject *str;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004312 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004313 char *s;
4314
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004315 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004316 return -1;
4317
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004318 size = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004319 if (size < 0) {
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004320 PyErr_Format(PyExc_OverflowError,
4321 "BINUNICODE exceeds system's maximum size of %zd bytes",
4322 PY_SSIZE_T_MAX
4323 );
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004324 return -1;
4325 }
4326
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004327
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004328 if (_Unpickler_Read(self, &s, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004329 return -1;
4330
Victor Stinner485fb562010-04-13 11:07:24 +00004331 str = PyUnicode_DecodeUTF8(s, size, "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004332 if (str == NULL)
4333 return -1;
4334
4335 PDATA_PUSH(self->stack, str, -1);
4336 return 0;
4337}
4338
4339static int
4340load_tuple(UnpicklerObject *self)
4341{
4342 PyObject *tuple;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004343 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004344
4345 if ((i = marker(self)) < 0)
4346 return -1;
4347
4348 tuple = Pdata_poptuple(self->stack, i);
4349 if (tuple == NULL)
4350 return -1;
4351 PDATA_PUSH(self->stack, tuple, -1);
4352 return 0;
4353}
4354
4355static int
4356load_counted_tuple(UnpicklerObject *self, int len)
4357{
4358 PyObject *tuple;
4359
4360 tuple = PyTuple_New(len);
4361 if (tuple == NULL)
4362 return -1;
4363
4364 while (--len >= 0) {
4365 PyObject *item;
4366
4367 PDATA_POP(self->stack, item);
4368 if (item == NULL)
4369 return -1;
4370 PyTuple_SET_ITEM(tuple, len, item);
4371 }
4372 PDATA_PUSH(self->stack, tuple, -1);
4373 return 0;
4374}
4375
4376static int
4377load_empty_list(UnpicklerObject *self)
4378{
4379 PyObject *list;
4380
4381 if ((list = PyList_New(0)) == NULL)
4382 return -1;
4383 PDATA_PUSH(self->stack, list, -1);
4384 return 0;
4385}
4386
4387static int
4388load_empty_dict(UnpicklerObject *self)
4389{
4390 PyObject *dict;
4391
4392 if ((dict = PyDict_New()) == NULL)
4393 return -1;
4394 PDATA_PUSH(self->stack, dict, -1);
4395 return 0;
4396}
4397
4398static int
4399load_list(UnpicklerObject *self)
4400{
4401 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004402 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004403
4404 if ((i = marker(self)) < 0)
4405 return -1;
4406
4407 list = Pdata_poplist(self->stack, i);
4408 if (list == NULL)
4409 return -1;
4410 PDATA_PUSH(self->stack, list, -1);
4411 return 0;
4412}
4413
4414static int
4415load_dict(UnpicklerObject *self)
4416{
4417 PyObject *dict, *key, *value;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004418 Py_ssize_t i, j, k;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004419
4420 if ((i = marker(self)) < 0)
4421 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004422 j = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004423
4424 if ((dict = PyDict_New()) == NULL)
4425 return -1;
4426
4427 for (k = i + 1; k < j; k += 2) {
4428 key = self->stack->data[k - 1];
4429 value = self->stack->data[k];
4430 if (PyDict_SetItem(dict, key, value) < 0) {
4431 Py_DECREF(dict);
4432 return -1;
4433 }
4434 }
4435 Pdata_clear(self->stack, i);
4436 PDATA_PUSH(self->stack, dict, -1);
4437 return 0;
4438}
4439
4440static PyObject *
4441instantiate(PyObject *cls, PyObject *args)
4442{
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004443 PyObject *result = NULL;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02004444 _Py_IDENTIFIER(__getinitargs__);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004445 /* Caller must assure args are a tuple. Normally, args come from
4446 Pdata_poptuple which packs objects from the top of the stack
4447 into a newly created tuple. */
4448 assert(PyTuple_Check(args));
4449 if (Py_SIZE(args) > 0 || !PyType_Check(cls) ||
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02004450 _PyObject_HasAttrId(cls, &PyId___getinitargs__)) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004451 result = PyObject_CallObject(cls, args);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004452 }
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004453 else {
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02004454 _Py_IDENTIFIER(__new__);
Martin v. Löwisafe55bb2011-10-09 10:38:36 +02004455
4456 result = _PyObject_CallMethodId(cls, &PyId___new__, "O", cls);
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004457 }
4458 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004459}
4460
4461static int
4462load_obj(UnpicklerObject *self)
4463{
4464 PyObject *cls, *args, *obj = NULL;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004465 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004466
4467 if ((i = marker(self)) < 0)
4468 return -1;
4469
4470 args = Pdata_poptuple(self->stack, i + 1);
4471 if (args == NULL)
4472 return -1;
4473
4474 PDATA_POP(self->stack, cls);
4475 if (cls) {
4476 obj = instantiate(cls, args);
4477 Py_DECREF(cls);
4478 }
4479 Py_DECREF(args);
4480 if (obj == NULL)
4481 return -1;
4482
4483 PDATA_PUSH(self->stack, obj, -1);
4484 return 0;
4485}
4486
4487static int
4488load_inst(UnpicklerObject *self)
4489{
4490 PyObject *cls = NULL;
4491 PyObject *args = NULL;
4492 PyObject *obj = NULL;
4493 PyObject *module_name;
4494 PyObject *class_name;
4495 Py_ssize_t len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004496 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004497 char *s;
4498
4499 if ((i = marker(self)) < 0)
4500 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004501 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004502 return -1;
4503 if (len < 2)
4504 return bad_readline();
4505
4506 /* Here it is safe to use PyUnicode_DecodeASCII(), even though non-ASCII
4507 identifiers are permitted in Python 3.0, since the INST opcode is only
4508 supported by older protocols on Python 2.x. */
4509 module_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
4510 if (module_name == NULL)
4511 return -1;
4512
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004513 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004514 if (len < 2)
4515 return bad_readline();
4516 class_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004517 if (class_name != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004518 cls = find_class(self, module_name, class_name);
4519 Py_DECREF(class_name);
4520 }
4521 }
4522 Py_DECREF(module_name);
4523
4524 if (cls == NULL)
4525 return -1;
4526
4527 if ((args = Pdata_poptuple(self->stack, i)) != NULL) {
4528 obj = instantiate(cls, args);
4529 Py_DECREF(args);
4530 }
4531 Py_DECREF(cls);
4532
4533 if (obj == NULL)
4534 return -1;
4535
4536 PDATA_PUSH(self->stack, obj, -1);
4537 return 0;
4538}
4539
4540static int
4541load_newobj(UnpicklerObject *self)
4542{
4543 PyObject *args = NULL;
4544 PyObject *clsraw = NULL;
4545 PyTypeObject *cls; /* clsraw cast to its true type */
4546 PyObject *obj;
4547
4548 /* Stack is ... cls argtuple, and we want to call
4549 * cls.__new__(cls, *argtuple).
4550 */
4551 PDATA_POP(self->stack, args);
4552 if (args == NULL)
4553 goto error;
4554 if (!PyTuple_Check(args)) {
4555 PyErr_SetString(UnpicklingError, "NEWOBJ expected an arg " "tuple.");
4556 goto error;
4557 }
4558
4559 PDATA_POP(self->stack, clsraw);
4560 cls = (PyTypeObject *)clsraw;
4561 if (cls == NULL)
4562 goto error;
4563 if (!PyType_Check(cls)) {
4564 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4565 "isn't a type object");
4566 goto error;
4567 }
4568 if (cls->tp_new == NULL) {
4569 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4570 "has NULL tp_new");
4571 goto error;
4572 }
4573
4574 /* Call __new__. */
4575 obj = cls->tp_new(cls, args, NULL);
4576 if (obj == NULL)
4577 goto error;
4578
4579 Py_DECREF(args);
4580 Py_DECREF(clsraw);
4581 PDATA_PUSH(self->stack, obj, -1);
4582 return 0;
4583
4584 error:
4585 Py_XDECREF(args);
4586 Py_XDECREF(clsraw);
4587 return -1;
4588}
4589
4590static int
4591load_global(UnpicklerObject *self)
4592{
4593 PyObject *global = NULL;
4594 PyObject *module_name;
4595 PyObject *global_name;
4596 Py_ssize_t len;
4597 char *s;
4598
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004599 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004600 return -1;
4601 if (len < 2)
4602 return bad_readline();
4603 module_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4604 if (!module_name)
4605 return -1;
4606
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004607 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004608 if (len < 2) {
4609 Py_DECREF(module_name);
4610 return bad_readline();
4611 }
4612 global_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4613 if (global_name) {
4614 global = find_class(self, module_name, global_name);
4615 Py_DECREF(global_name);
4616 }
4617 }
4618 Py_DECREF(module_name);
4619
4620 if (global == NULL)
4621 return -1;
4622 PDATA_PUSH(self->stack, global, -1);
4623 return 0;
4624}
4625
4626static int
4627load_persid(UnpicklerObject *self)
4628{
4629 PyObject *pid;
4630 Py_ssize_t len;
4631 char *s;
4632
4633 if (self->pers_func) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004634 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004635 return -1;
4636 if (len < 2)
4637 return bad_readline();
4638
4639 pid = PyBytes_FromStringAndSize(s, len - 1);
4640 if (pid == NULL)
4641 return -1;
4642
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004643 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004644 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004645 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004646 if (pid == NULL)
4647 return -1;
4648
4649 PDATA_PUSH(self->stack, pid, -1);
4650 return 0;
4651 }
4652 else {
4653 PyErr_SetString(UnpicklingError,
4654 "A load persistent id instruction was encountered,\n"
4655 "but no persistent_load function was specified.");
4656 return -1;
4657 }
4658}
4659
4660static int
4661load_binpersid(UnpicklerObject *self)
4662{
4663 PyObject *pid;
4664
4665 if (self->pers_func) {
4666 PDATA_POP(self->stack, pid);
4667 if (pid == NULL)
4668 return -1;
4669
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004670 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004671 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004672 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004673 if (pid == NULL)
4674 return -1;
4675
4676 PDATA_PUSH(self->stack, pid, -1);
4677 return 0;
4678 }
4679 else {
4680 PyErr_SetString(UnpicklingError,
4681 "A load persistent id instruction was encountered,\n"
4682 "but no persistent_load function was specified.");
4683 return -1;
4684 }
4685}
4686
4687static int
4688load_pop(UnpicklerObject *self)
4689{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004690 Py_ssize_t len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004691
4692 /* Note that we split the (pickle.py) stack into two stacks,
4693 * an object stack and a mark stack. We have to be clever and
4694 * pop the right one. We do this by looking at the top of the
Collin Winter8ca69de2009-05-26 16:53:41 +00004695 * mark stack first, and only signalling a stack underflow if
4696 * the object stack is empty and the mark stack doesn't match
4697 * our expectations.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004698 */
Collin Winter8ca69de2009-05-26 16:53:41 +00004699 if (self->num_marks > 0 && self->marks[self->num_marks - 1] == len) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004700 self->num_marks--;
Antoine Pitrou01a15ea2010-01-07 17:57:31 +00004701 } else if (len > 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004702 len--;
4703 Py_DECREF(self->stack->data[len]);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004704 Py_SIZE(self->stack) = len;
Collin Winter8ca69de2009-05-26 16:53:41 +00004705 } else {
4706 return stack_underflow();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004707 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004708 return 0;
4709}
4710
4711static int
4712load_pop_mark(UnpicklerObject *self)
4713{
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004714 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004715
4716 if ((i = marker(self)) < 0)
4717 return -1;
4718
4719 Pdata_clear(self->stack, i);
4720
4721 return 0;
4722}
4723
4724static int
4725load_dup(UnpicklerObject *self)
4726{
4727 PyObject *last;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004728 Py_ssize_t len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004729
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004730 if ((len = Py_SIZE(self->stack)) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004731 return stack_underflow();
4732 last = self->stack->data[len - 1];
4733 PDATA_APPEND(self->stack, last, -1);
4734 return 0;
4735}
4736
4737static int
4738load_get(UnpicklerObject *self)
4739{
4740 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004741 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004742 Py_ssize_t len;
4743 char *s;
4744
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004745 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004746 return -1;
4747 if (len < 2)
4748 return bad_readline();
4749
4750 key = PyLong_FromString(s, NULL, 10);
4751 if (key == NULL)
4752 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004753 idx = PyLong_AsSsize_t(key);
4754 if (idx == -1 && PyErr_Occurred()) {
4755 Py_DECREF(key);
4756 return -1;
4757 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004758
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004759 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004760 if (value == NULL) {
4761 if (!PyErr_Occurred())
4762 PyErr_SetObject(PyExc_KeyError, key);
4763 Py_DECREF(key);
4764 return -1;
4765 }
4766 Py_DECREF(key);
4767
4768 PDATA_APPEND(self->stack, value, -1);
4769 return 0;
4770}
4771
4772static int
4773load_binget(UnpicklerObject *self)
4774{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004775 PyObject *value;
4776 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004777 char *s;
4778
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004779 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004780 return -1;
4781
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004782 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004783
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004784 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004785 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004786 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004787 if (!PyErr_Occurred())
4788 PyErr_SetObject(PyExc_KeyError, key);
4789 Py_DECREF(key);
4790 return -1;
4791 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004792
4793 PDATA_APPEND(self->stack, value, -1);
4794 return 0;
4795}
4796
4797static int
4798load_long_binget(UnpicklerObject *self)
4799{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004800 PyObject *value;
4801 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004802 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004803
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004804 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004805 return -1;
4806
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004807 idx = calc_binsize(s, 4);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004808
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004809 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004810 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004811 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004812 if (!PyErr_Occurred())
4813 PyErr_SetObject(PyExc_KeyError, key);
4814 Py_DECREF(key);
4815 return -1;
4816 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004817
4818 PDATA_APPEND(self->stack, value, -1);
4819 return 0;
4820}
4821
4822/* Push an object from the extension registry (EXT[124]). nbytes is
4823 * the number of bytes following the opcode, holding the index (code) value.
4824 */
4825static int
4826load_extension(UnpicklerObject *self, int nbytes)
4827{
4828 char *codebytes; /* the nbytes bytes after the opcode */
4829 long code; /* calc_binint returns long */
4830 PyObject *py_code; /* code as a Python int */
4831 PyObject *obj; /* the object to push */
4832 PyObject *pair; /* (module_name, class_name) */
4833 PyObject *module_name, *class_name;
4834
4835 assert(nbytes == 1 || nbytes == 2 || nbytes == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004836 if (_Unpickler_Read(self, &codebytes, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004837 return -1;
4838 code = calc_binint(codebytes, nbytes);
4839 if (code <= 0) { /* note that 0 is forbidden */
4840 /* Corrupt or hostile pickle. */
4841 PyErr_SetString(UnpicklingError, "EXT specifies code <= 0");
4842 return -1;
4843 }
4844
4845 /* Look for the code in the cache. */
4846 py_code = PyLong_FromLong(code);
4847 if (py_code == NULL)
4848 return -1;
4849 obj = PyDict_GetItem(extension_cache, py_code);
4850 if (obj != NULL) {
4851 /* Bingo. */
4852 Py_DECREF(py_code);
4853 PDATA_APPEND(self->stack, obj, -1);
4854 return 0;
4855 }
4856
4857 /* Look up the (module_name, class_name) pair. */
4858 pair = PyDict_GetItem(inverted_registry, py_code);
4859 if (pair == NULL) {
4860 Py_DECREF(py_code);
4861 PyErr_Format(PyExc_ValueError, "unregistered extension "
4862 "code %ld", code);
4863 return -1;
4864 }
4865 /* Since the extension registry is manipulable via Python code,
4866 * confirm that pair is really a 2-tuple of strings.
4867 */
4868 if (!PyTuple_Check(pair) || PyTuple_Size(pair) != 2 ||
4869 !PyUnicode_Check(module_name = PyTuple_GET_ITEM(pair, 0)) ||
4870 !PyUnicode_Check(class_name = PyTuple_GET_ITEM(pair, 1))) {
4871 Py_DECREF(py_code);
4872 PyErr_Format(PyExc_ValueError, "_inverted_registry[%ld] "
4873 "isn't a 2-tuple of strings", code);
4874 return -1;
4875 }
4876 /* Load the object. */
4877 obj = find_class(self, module_name, class_name);
4878 if (obj == NULL) {
4879 Py_DECREF(py_code);
4880 return -1;
4881 }
4882 /* Cache code -> obj. */
4883 code = PyDict_SetItem(extension_cache, py_code, obj);
4884 Py_DECREF(py_code);
4885 if (code < 0) {
4886 Py_DECREF(obj);
4887 return -1;
4888 }
4889 PDATA_PUSH(self->stack, obj, -1);
4890 return 0;
4891}
4892
4893static int
4894load_put(UnpicklerObject *self)
4895{
4896 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004897 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004898 Py_ssize_t len;
4899 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004900
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004901 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004902 return -1;
4903 if (len < 2)
4904 return bad_readline();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004905 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004906 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004907 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004908
4909 key = PyLong_FromString(s, NULL, 10);
4910 if (key == NULL)
4911 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004912 idx = PyLong_AsSsize_t(key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004913 Py_DECREF(key);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02004914 if (idx < 0) {
4915 if (!PyErr_Occurred())
4916 PyErr_SetString(PyExc_ValueError,
4917 "negative PUT argument");
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004918 return -1;
Antoine Pitrou55549ec2011-08-30 00:27:10 +02004919 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004920
4921 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004922}
4923
4924static int
4925load_binput(UnpicklerObject *self)
4926{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004927 PyObject *value;
4928 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004929 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004930
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004931 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004932 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004933
4934 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004935 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004936 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004937
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004938 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004939
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004940 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004941}
4942
4943static int
4944load_long_binput(UnpicklerObject *self)
4945{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004946 PyObject *value;
4947 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004948 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004949
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004950 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004951 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004952
4953 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004954 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004955 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004956
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004957 idx = calc_binsize(s, 4);
Antoine Pitrou55549ec2011-08-30 00:27:10 +02004958 if (idx < 0) {
4959 PyErr_SetString(PyExc_ValueError,
4960 "negative LONG_BINPUT argument");
4961 return -1;
4962 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004963
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004964 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004965}
4966
4967static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004968do_append(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004969{
4970 PyObject *value;
4971 PyObject *list;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004972 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004973
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004974 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004975 if (x > len || x <= 0)
4976 return stack_underflow();
4977 if (len == x) /* nothing to do */
4978 return 0;
4979
4980 list = self->stack->data[x - 1];
4981
4982 if (PyList_Check(list)) {
4983 PyObject *slice;
4984 Py_ssize_t list_len;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004985 int ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004986
4987 slice = Pdata_poplist(self->stack, x);
4988 if (!slice)
4989 return -1;
4990 list_len = PyList_GET_SIZE(list);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004991 ret = PyList_SetSlice(list, list_len, list_len, slice);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004992 Py_DECREF(slice);
Antoine Pitrou82be19f2011-08-29 23:09:33 +02004993 return ret;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004994 }
4995 else {
4996 PyObject *append_func;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02004997 _Py_IDENTIFIER(append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004998
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02004999 append_func = _PyObject_GetAttrId(list, &PyId_append);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005000 if (append_func == NULL)
5001 return -1;
5002 for (i = x; i < len; i++) {
5003 PyObject *result;
5004
5005 value = self->stack->data[i];
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005006 result = _Unpickler_FastCall(self, append_func, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005007 if (result == NULL) {
5008 Pdata_clear(self->stack, i + 1);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005009 Py_SIZE(self->stack) = x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005010 return -1;
5011 }
5012 Py_DECREF(result);
5013 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005014 Py_SIZE(self->stack) = x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005015 }
5016
5017 return 0;
5018}
5019
5020static int
5021load_append(UnpicklerObject *self)
5022{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005023 return do_append(self, Py_SIZE(self->stack) - 1);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005024}
5025
5026static int
5027load_appends(UnpicklerObject *self)
5028{
5029 return do_append(self, marker(self));
5030}
5031
5032static int
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005033do_setitems(UnpicklerObject *self, Py_ssize_t x)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005034{
5035 PyObject *value, *key;
5036 PyObject *dict;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005037 Py_ssize_t len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005038 int status = 0;
5039
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;
Victor Stinner121aab42011-09-29 23:40:53 +02005045 if ((len - x) % 2 != 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005046 /* Currupt or hostile pickle -- we never write one like this. */
5047 PyErr_SetString(UnpicklingError, "odd number of items for SETITEMS");
5048 return -1;
5049 }
5050
5051 /* Here, dict does not actually need to be a PyDict; it could be anything
5052 that supports the __setitem__ attribute. */
5053 dict = self->stack->data[x - 1];
5054
5055 for (i = x + 1; i < len; i += 2) {
5056 key = self->stack->data[i - 1];
5057 value = self->stack->data[i];
5058 if (PyObject_SetItem(dict, key, value) < 0) {
5059 status = -1;
5060 break;
5061 }
5062 }
5063
5064 Pdata_clear(self->stack, x);
5065 return status;
5066}
5067
5068static int
5069load_setitem(UnpicklerObject *self)
5070{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005071 return do_setitems(self, Py_SIZE(self->stack) - 2);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005072}
5073
5074static int
5075load_setitems(UnpicklerObject *self)
5076{
5077 return do_setitems(self, marker(self));
5078}
5079
5080static int
5081load_build(UnpicklerObject *self)
5082{
5083 PyObject *state, *inst, *slotstate;
5084 PyObject *setstate;
5085 int status = 0;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005086 _Py_IDENTIFIER(__setstate__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005087
5088 /* Stack is ... instance, state. We want to leave instance at
5089 * the stack top, possibly mutated via instance.__setstate__(state).
5090 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005091 if (Py_SIZE(self->stack) < 2)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005092 return stack_underflow();
5093
5094 PDATA_POP(self->stack, state);
5095 if (state == NULL)
5096 return -1;
5097
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005098 inst = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005099
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005100 setstate = _PyObject_GetAttrId(inst, &PyId___setstate__);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005101 if (setstate == NULL) {
5102 if (PyErr_ExceptionMatches(PyExc_AttributeError))
5103 PyErr_Clear();
Antoine Pitroud79dc622008-09-05 00:03:33 +00005104 else {
5105 Py_DECREF(state);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005106 return -1;
Antoine Pitroud79dc622008-09-05 00:03:33 +00005107 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005108 }
5109 else {
5110 PyObject *result;
5111
5112 /* The explicit __setstate__ is responsible for everything. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005113 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Antoine Pitroud79dc622008-09-05 00:03:33 +00005114 reference to state first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005115 result = _Unpickler_FastCall(self, setstate, state);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005116 Py_DECREF(setstate);
5117 if (result == NULL)
5118 return -1;
5119 Py_DECREF(result);
5120 return 0;
5121 }
5122
5123 /* A default __setstate__. First see whether state embeds a
5124 * slot state dict too (a proto 2 addition).
5125 */
5126 if (PyTuple_Check(state) && Py_SIZE(state) == 2) {
5127 PyObject *tmp = state;
5128
5129 state = PyTuple_GET_ITEM(tmp, 0);
5130 slotstate = PyTuple_GET_ITEM(tmp, 1);
5131 Py_INCREF(state);
5132 Py_INCREF(slotstate);
5133 Py_DECREF(tmp);
5134 }
5135 else
5136 slotstate = NULL;
5137
5138 /* Set inst.__dict__ from the state dict (if any). */
5139 if (state != Py_None) {
5140 PyObject *dict;
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005141 PyObject *d_key, *d_value;
5142 Py_ssize_t i;
Martin v. Löwisbd928fe2011-10-14 10:20:37 +02005143 _Py_IDENTIFIER(__dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005144
5145 if (!PyDict_Check(state)) {
5146 PyErr_SetString(UnpicklingError, "state is not a dictionary");
5147 goto error;
5148 }
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005149 dict = _PyObject_GetAttrId(inst, &PyId___dict__);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005150 if (dict == NULL)
5151 goto error;
5152
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005153 i = 0;
5154 while (PyDict_Next(state, &i, &d_key, &d_value)) {
5155 /* normally the keys for instance attributes are
5156 interned. we should try to do that here. */
5157 Py_INCREF(d_key);
5158 if (PyUnicode_CheckExact(d_key))
5159 PyUnicode_InternInPlace(&d_key);
5160 if (PyObject_SetItem(dict, d_key, d_value) < 0) {
5161 Py_DECREF(d_key);
5162 goto error;
5163 }
5164 Py_DECREF(d_key);
5165 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005166 Py_DECREF(dict);
5167 }
5168
5169 /* Also set instance attributes from the slotstate dict (if any). */
5170 if (slotstate != NULL) {
5171 PyObject *d_key, *d_value;
5172 Py_ssize_t i;
5173
5174 if (!PyDict_Check(slotstate)) {
5175 PyErr_SetString(UnpicklingError,
5176 "slot state is not a dictionary");
5177 goto error;
5178 }
5179 i = 0;
5180 while (PyDict_Next(slotstate, &i, &d_key, &d_value)) {
5181 if (PyObject_SetAttr(inst, d_key, d_value) < 0)
5182 goto error;
5183 }
5184 }
5185
5186 if (0) {
5187 error:
5188 status = -1;
5189 }
5190
5191 Py_DECREF(state);
5192 Py_XDECREF(slotstate);
5193 return status;
5194}
5195
5196static int
5197load_mark(UnpicklerObject *self)
5198{
5199
5200 /* Note that we split the (pickle.py) stack into two stacks, an
5201 * object stack and a mark stack. Here we push a mark onto the
5202 * mark stack.
5203 */
5204
5205 if ((self->num_marks + 1) >= self->marks_size) {
5206 size_t alloc;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005207 Py_ssize_t *marks;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005208
5209 /* Use the size_t type to check for overflow. */
5210 alloc = ((size_t)self->num_marks << 1) + 20;
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005211 if (alloc > (PY_SSIZE_T_MAX / sizeof(Py_ssize_t)) ||
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00005212 alloc <= ((size_t)self->num_marks + 1)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005213 PyErr_NoMemory();
5214 return -1;
5215 }
5216
5217 if (self->marks == NULL)
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005218 marks = (Py_ssize_t *) PyMem_Malloc(alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005219 else
Antoine Pitrou82be19f2011-08-29 23:09:33 +02005220 marks = (Py_ssize_t *) PyMem_Realloc(self->marks,
5221 alloc * sizeof(Py_ssize_t));
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005222 if (marks == NULL) {
5223 PyErr_NoMemory();
5224 return -1;
5225 }
5226 self->marks = marks;
5227 self->marks_size = (Py_ssize_t)alloc;
5228 }
5229
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005230 self->marks[self->num_marks++] = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005231
5232 return 0;
5233}
5234
5235static int
5236load_reduce(UnpicklerObject *self)
5237{
5238 PyObject *callable = NULL;
5239 PyObject *argtup = NULL;
5240 PyObject *obj = NULL;
5241
5242 PDATA_POP(self->stack, argtup);
5243 if (argtup == NULL)
5244 return -1;
5245 PDATA_POP(self->stack, callable);
5246 if (callable) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005247 obj = PyObject_CallObject(callable, argtup);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005248 Py_DECREF(callable);
5249 }
5250 Py_DECREF(argtup);
5251
5252 if (obj == NULL)
5253 return -1;
5254
5255 PDATA_PUSH(self->stack, obj, -1);
5256 return 0;
5257}
5258
5259/* Just raises an error if we don't know the protocol specified. PROTO
5260 * is the first opcode for protocols >= 2.
5261 */
5262static int
5263load_proto(UnpicklerObject *self)
5264{
5265 char *s;
5266 int i;
5267
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005268 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005269 return -1;
5270
5271 i = (unsigned char)s[0];
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005272 if (i <= HIGHEST_PROTOCOL) {
5273 self->proto = i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005274 return 0;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005275 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005276
5277 PyErr_Format(PyExc_ValueError, "unsupported pickle protocol: %d", i);
5278 return -1;
5279}
5280
5281static PyObject *
5282load(UnpicklerObject *self)
5283{
5284 PyObject *err;
5285 PyObject *value = NULL;
5286 char *s;
5287
5288 self->num_marks = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005289 if (Py_SIZE(self->stack))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005290 Pdata_clear(self->stack, 0);
5291
5292 /* Convenient macros for the dispatch while-switch loop just below. */
5293#define OP(opcode, load_func) \
5294 case opcode: if (load_func(self) < 0) break; continue;
5295
5296#define OP_ARG(opcode, load_func, arg) \
5297 case opcode: if (load_func(self, (arg)) < 0) break; continue;
5298
5299 while (1) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005300 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005301 break;
5302
5303 switch ((enum opcode)s[0]) {
5304 OP(NONE, load_none)
5305 OP(BININT, load_binint)
5306 OP(BININT1, load_binint1)
5307 OP(BININT2, load_binint2)
5308 OP(INT, load_int)
5309 OP(LONG, load_long)
5310 OP_ARG(LONG1, load_counted_long, 1)
5311 OP_ARG(LONG4, load_counted_long, 4)
5312 OP(FLOAT, load_float)
5313 OP(BINFLOAT, load_binfloat)
5314 OP(BINBYTES, load_binbytes)
5315 OP(SHORT_BINBYTES, load_short_binbytes)
5316 OP(BINSTRING, load_binstring)
5317 OP(SHORT_BINSTRING, load_short_binstring)
5318 OP(STRING, load_string)
5319 OP(UNICODE, load_unicode)
5320 OP(BINUNICODE, load_binunicode)
5321 OP_ARG(EMPTY_TUPLE, load_counted_tuple, 0)
5322 OP_ARG(TUPLE1, load_counted_tuple, 1)
5323 OP_ARG(TUPLE2, load_counted_tuple, 2)
5324 OP_ARG(TUPLE3, load_counted_tuple, 3)
5325 OP(TUPLE, load_tuple)
5326 OP(EMPTY_LIST, load_empty_list)
5327 OP(LIST, load_list)
5328 OP(EMPTY_DICT, load_empty_dict)
5329 OP(DICT, load_dict)
5330 OP(OBJ, load_obj)
5331 OP(INST, load_inst)
5332 OP(NEWOBJ, load_newobj)
5333 OP(GLOBAL, load_global)
5334 OP(APPEND, load_append)
5335 OP(APPENDS, load_appends)
5336 OP(BUILD, load_build)
5337 OP(DUP, load_dup)
5338 OP(BINGET, load_binget)
5339 OP(LONG_BINGET, load_long_binget)
5340 OP(GET, load_get)
5341 OP(MARK, load_mark)
5342 OP(BINPUT, load_binput)
5343 OP(LONG_BINPUT, load_long_binput)
5344 OP(PUT, load_put)
5345 OP(POP, load_pop)
5346 OP(POP_MARK, load_pop_mark)
5347 OP(SETITEM, load_setitem)
5348 OP(SETITEMS, load_setitems)
5349 OP(PERSID, load_persid)
5350 OP(BINPERSID, load_binpersid)
5351 OP(REDUCE, load_reduce)
5352 OP(PROTO, load_proto)
5353 OP_ARG(EXT1, load_extension, 1)
5354 OP_ARG(EXT2, load_extension, 2)
5355 OP_ARG(EXT4, load_extension, 4)
5356 OP_ARG(NEWTRUE, load_bool, Py_True)
5357 OP_ARG(NEWFALSE, load_bool, Py_False)
5358
5359 case STOP:
5360 break;
5361
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005362 default:
Benjamin Petersonadde86d2011-09-23 13:41:41 -04005363 if (s[0] == '\0')
5364 PyErr_SetNone(PyExc_EOFError);
5365 else
5366 PyErr_Format(UnpicklingError,
5367 "invalid load key, '%c'.", s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005368 return NULL;
5369 }
5370
5371 break; /* and we are done! */
5372 }
5373
Antoine Pitrou04248a82010-10-12 20:51:21 +00005374 if (_Unpickler_SkipConsumed(self) < 0)
5375 return NULL;
5376
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005377 /* XXX: It is not clear what this is actually for. */
5378 if ((err = PyErr_Occurred())) {
5379 if (err == PyExc_EOFError) {
5380 PyErr_SetNone(PyExc_EOFError);
5381 }
5382 return NULL;
5383 }
5384
5385 PDATA_POP(self->stack, value);
5386 return value;
5387}
5388
5389PyDoc_STRVAR(Unpickler_load_doc,
5390"load() -> object. Load a pickle."
5391"\n"
5392"Read a pickled object representation from the open file object given in\n"
5393"the constructor, and return the reconstituted object hierarchy specified\n"
5394"therein.\n");
5395
5396static PyObject *
5397Unpickler_load(UnpicklerObject *self)
5398{
5399 /* Check whether the Unpickler was initialized correctly. This prevents
5400 segfaulting if a subclass overridden __init__ with a function that does
5401 not call Unpickler.__init__(). Here, we simply ensure that self->read
5402 is not NULL. */
5403 if (self->read == NULL) {
Victor Stinner121aab42011-09-29 23:40:53 +02005404 PyErr_Format(UnpicklingError,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005405 "Unpickler.__init__() was not called by %s.__init__()",
5406 Py_TYPE(self)->tp_name);
5407 return NULL;
5408 }
5409
5410 return load(self);
5411}
5412
5413/* The name of find_class() is misleading. In newer pickle protocols, this
5414 function is used for loading any global (i.e., functions), not just
5415 classes. The name is kept only for backward compatibility. */
5416
5417PyDoc_STRVAR(Unpickler_find_class_doc,
5418"find_class(module_name, global_name) -> object.\n"
5419"\n"
5420"Return an object from a specified module, importing the module if\n"
5421"necessary. Subclasses may override this method (e.g. to restrict\n"
5422"unpickling of arbitrary classes and functions).\n"
5423"\n"
5424"This method is called whenever a class or a function object is\n"
5425"needed. Both arguments passed are str objects.\n");
5426
5427static PyObject *
5428Unpickler_find_class(UnpicklerObject *self, PyObject *args)
5429{
5430 PyObject *global;
5431 PyObject *modules_dict;
5432 PyObject *module;
5433 PyObject *module_name, *global_name;
5434
5435 if (!PyArg_UnpackTuple(args, "find_class", 2, 2,
5436 &module_name, &global_name))
5437 return NULL;
5438
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005439 /* Try to map the old names used in Python 2.x to the new ones used in
5440 Python 3.x. We do this only with old pickle protocols and when the
5441 user has not disabled the feature. */
5442 if (self->proto < 3 && self->fix_imports) {
5443 PyObject *key;
5444 PyObject *item;
5445
5446 /* Check if the global (i.e., a function or a class) was renamed
5447 or moved to another module. */
5448 key = PyTuple_Pack(2, module_name, global_name);
5449 if (key == NULL)
5450 return NULL;
5451 item = PyDict_GetItemWithError(name_mapping_2to3, key);
5452 Py_DECREF(key);
5453 if (item) {
5454 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
5455 PyErr_Format(PyExc_RuntimeError,
5456 "_compat_pickle.NAME_MAPPING values should be "
5457 "2-tuples, not %.200s", Py_TYPE(item)->tp_name);
5458 return NULL;
5459 }
5460 module_name = PyTuple_GET_ITEM(item, 0);
5461 global_name = PyTuple_GET_ITEM(item, 1);
5462 if (!PyUnicode_Check(module_name) ||
5463 !PyUnicode_Check(global_name)) {
5464 PyErr_Format(PyExc_RuntimeError,
5465 "_compat_pickle.NAME_MAPPING values should be "
5466 "pairs of str, not (%.200s, %.200s)",
5467 Py_TYPE(module_name)->tp_name,
5468 Py_TYPE(global_name)->tp_name);
5469 return NULL;
5470 }
5471 }
5472 else if (PyErr_Occurred()) {
5473 return NULL;
5474 }
5475
5476 /* Check if the module was renamed. */
5477 item = PyDict_GetItemWithError(import_mapping_2to3, module_name);
5478 if (item) {
5479 if (!PyUnicode_Check(item)) {
5480 PyErr_Format(PyExc_RuntimeError,
5481 "_compat_pickle.IMPORT_MAPPING values should be "
5482 "strings, not %.200s", Py_TYPE(item)->tp_name);
5483 return NULL;
5484 }
5485 module_name = item;
5486 }
5487 else if (PyErr_Occurred()) {
5488 return NULL;
5489 }
5490 }
5491
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005492 modules_dict = PySys_GetObject("modules");
5493 if (modules_dict == NULL)
5494 return NULL;
5495
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005496 module = PyDict_GetItemWithError(modules_dict, module_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005497 if (module == NULL) {
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005498 if (PyErr_Occurred())
5499 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005500 module = PyImport_Import(module_name);
5501 if (module == NULL)
5502 return NULL;
5503 global = PyObject_GetAttr(module, global_name);
5504 Py_DECREF(module);
5505 }
Victor Stinner121aab42011-09-29 23:40:53 +02005506 else {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005507 global = PyObject_GetAttr(module, global_name);
5508 }
5509 return global;
5510}
5511
5512static struct PyMethodDef Unpickler_methods[] = {
5513 {"load", (PyCFunction)Unpickler_load, METH_NOARGS,
5514 Unpickler_load_doc},
5515 {"find_class", (PyCFunction)Unpickler_find_class, METH_VARARGS,
5516 Unpickler_find_class_doc},
5517 {NULL, NULL} /* sentinel */
5518};
5519
5520static void
5521Unpickler_dealloc(UnpicklerObject *self)
5522{
5523 PyObject_GC_UnTrack((PyObject *)self);
5524 Py_XDECREF(self->readline);
5525 Py_XDECREF(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005526 Py_XDECREF(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005527 Py_XDECREF(self->stack);
5528 Py_XDECREF(self->pers_func);
5529 Py_XDECREF(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005530 if (self->buffer.buf != NULL) {
5531 PyBuffer_Release(&self->buffer);
5532 self->buffer.buf = NULL;
5533 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005534
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005535 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005536 PyMem_Free(self->marks);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005537 PyMem_Free(self->input_line);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005538 free(self->encoding);
5539 free(self->errors);
5540
5541 Py_TYPE(self)->tp_free((PyObject *)self);
5542}
5543
5544static int
5545Unpickler_traverse(UnpicklerObject *self, visitproc visit, void *arg)
5546{
5547 Py_VISIT(self->readline);
5548 Py_VISIT(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005549 Py_VISIT(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005550 Py_VISIT(self->stack);
5551 Py_VISIT(self->pers_func);
5552 Py_VISIT(self->arg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005553 return 0;
5554}
5555
5556static int
5557Unpickler_clear(UnpicklerObject *self)
5558{
5559 Py_CLEAR(self->readline);
5560 Py_CLEAR(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005561 Py_CLEAR(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005562 Py_CLEAR(self->stack);
5563 Py_CLEAR(self->pers_func);
5564 Py_CLEAR(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005565 if (self->buffer.buf != NULL) {
5566 PyBuffer_Release(&self->buffer);
5567 self->buffer.buf = NULL;
5568 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005569
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005570 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005571 PyMem_Free(self->marks);
5572 self->marks = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005573 PyMem_Free(self->input_line);
5574 self->input_line = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005575 free(self->encoding);
5576 self->encoding = NULL;
5577 free(self->errors);
5578 self->errors = NULL;
5579
5580 return 0;
5581}
5582
5583PyDoc_STRVAR(Unpickler_doc,
5584"Unpickler(file, *, encoding='ASCII', errors='strict')"
5585"\n"
5586"This takes a binary file for reading a pickle data stream.\n"
5587"\n"
5588"The protocol version of the pickle is detected automatically, so no\n"
5589"proto argument is needed.\n"
5590"\n"
5591"The file-like object must have two methods, a read() method\n"
5592"that takes an integer argument, and a readline() method that\n"
5593"requires no arguments. Both methods should return bytes.\n"
5594"Thus file-like object can be a binary file object opened for\n"
5595"reading, a BytesIO object, or any other custom object that\n"
5596"meets this interface.\n"
5597"\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005598"Optional keyword arguments are *fix_imports*, *encoding* and *errors*,\n"
5599"which are used to control compatiblity support for pickle stream\n"
5600"generated by Python 2.x. If *fix_imports* is True, pickle will try to\n"
5601"map the old Python 2.x names to the new names used in Python 3.x. The\n"
5602"*encoding* and *errors* tell pickle how to decode 8-bit string\n"
5603"instances pickled by Python 2.x; these default to 'ASCII' and\n"
5604"'strict', respectively.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005605
5606static int
5607Unpickler_init(UnpicklerObject *self, PyObject *args, PyObject *kwds)
5608{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005609 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005610 PyObject *file;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005611 PyObject *fix_imports = Py_True;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005612 char *encoding = NULL;
5613 char *errors = NULL;
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005614 _Py_IDENTIFIER(persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005615
5616 /* XXX: That is an horrible error message. But, I don't know how to do
5617 better... */
5618 if (Py_SIZE(args) != 1) {
5619 PyErr_Format(PyExc_TypeError,
5620 "%s takes exactly one positional argument (%zd given)",
5621 Py_TYPE(self)->tp_name, Py_SIZE(args));
5622 return -1;
5623 }
5624
5625 /* Arguments parsing needs to be done in the __init__() method to allow
5626 subclasses to define their own __init__() method, which may (or may
5627 not) support Unpickler arguments. However, this means we need to be
5628 extra careful in the other Unpickler methods, since a subclass could
5629 forget to call Unpickler.__init__() thus breaking our internal
5630 invariants. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005631 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:Unpickler", kwlist,
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005632 &file, &fix_imports, &encoding, &errors))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005633 return -1;
5634
5635 /* In case of multiple __init__() calls, clear previous content. */
5636 if (self->read != NULL)
5637 (void)Unpickler_clear(self);
5638
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005639 if (_Unpickler_SetInputStream(self, file) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005640 return -1;
5641
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005642 if (_Unpickler_SetInputEncoding(self, encoding, errors) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005643 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005644
5645 self->fix_imports = PyObject_IsTrue(fix_imports);
5646 if (self->fix_imports == -1)
5647 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005648
Martin v. Löwis1c67dd92011-10-14 15:16:45 +02005649 if (_PyObject_HasAttrId((PyObject *)self, &PyId_persistent_load)) {
Martin v. Löwis1ee1b6f2011-10-10 18:11:30 +02005650 self->pers_func = _PyObject_GetAttrId((PyObject *)self,
5651 &PyId_persistent_load);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005652 if (self->pers_func == NULL)
5653 return -1;
5654 }
5655 else {
5656 self->pers_func = NULL;
5657 }
5658
5659 self->stack = (Pdata *)Pdata_New();
5660 if (self->stack == NULL)
5661 return -1;
5662
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005663 self->memo_size = 32;
5664 self->memo = _Unpickler_NewMemo(self->memo_size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005665 if (self->memo == NULL)
5666 return -1;
5667
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005668 self->arg = NULL;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005669 self->proto = 0;
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005670
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005671 return 0;
5672}
5673
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005674/* Define a proxy object for the Unpickler's internal memo object. This is to
5675 * avoid breaking code like:
5676 * unpickler.memo.clear()
5677 * and
5678 * unpickler.memo = saved_memo
5679 * Is this a good idea? Not really, but we don't want to break code that uses
5680 * it. Note that we don't implement the entire mapping API here. This is
5681 * intentional, as these should be treated as black-box implementation details.
5682 *
5683 * We do, however, have to implement pickling/unpickling support because of
Victor Stinner121aab42011-09-29 23:40:53 +02005684 * real-world code like cvs2svn.
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005685 */
5686
5687typedef struct {
5688 PyObject_HEAD
5689 UnpicklerObject *unpickler;
5690} UnpicklerMemoProxyObject;
5691
5692PyDoc_STRVAR(ump_clear_doc,
5693"memo.clear() -> None. Remove all items from memo.");
5694
5695static PyObject *
5696ump_clear(UnpicklerMemoProxyObject *self)
5697{
5698 _Unpickler_MemoCleanup(self->unpickler);
5699 self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size);
5700 if (self->unpickler->memo == NULL)
5701 return NULL;
5702 Py_RETURN_NONE;
5703}
5704
5705PyDoc_STRVAR(ump_copy_doc,
5706"memo.copy() -> new_memo. Copy the memo to a new object.");
5707
5708static PyObject *
5709ump_copy(UnpicklerMemoProxyObject *self)
5710{
5711 Py_ssize_t i;
5712 PyObject *new_memo = PyDict_New();
5713 if (new_memo == NULL)
5714 return NULL;
5715
5716 for (i = 0; i < self->unpickler->memo_size; i++) {
5717 int status;
5718 PyObject *key, *value;
5719
5720 value = self->unpickler->memo[i];
5721 if (value == NULL)
5722 continue;
5723
5724 key = PyLong_FromSsize_t(i);
5725 if (key == NULL)
5726 goto error;
5727 status = PyDict_SetItem(new_memo, key, value);
5728 Py_DECREF(key);
5729 if (status < 0)
5730 goto error;
5731 }
5732 return new_memo;
5733
5734error:
5735 Py_DECREF(new_memo);
5736 return NULL;
5737}
5738
5739PyDoc_STRVAR(ump_reduce_doc,
5740"memo.__reduce__(). Pickling support.");
5741
5742static PyObject *
5743ump_reduce(UnpicklerMemoProxyObject *self, PyObject *args)
5744{
5745 PyObject *reduce_value;
5746 PyObject *constructor_args;
5747 PyObject *contents = ump_copy(self);
5748 if (contents == NULL)
5749 return NULL;
5750
5751 reduce_value = PyTuple_New(2);
5752 if (reduce_value == NULL) {
5753 Py_DECREF(contents);
5754 return NULL;
5755 }
5756 constructor_args = PyTuple_New(1);
5757 if (constructor_args == NULL) {
5758 Py_DECREF(contents);
5759 Py_DECREF(reduce_value);
5760 return NULL;
5761 }
5762 PyTuple_SET_ITEM(constructor_args, 0, contents);
5763 Py_INCREF((PyObject *)&PyDict_Type);
5764 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
5765 PyTuple_SET_ITEM(reduce_value, 1, constructor_args);
5766 return reduce_value;
5767}
5768
5769static PyMethodDef unpicklerproxy_methods[] = {
5770 {"clear", (PyCFunction)ump_clear, METH_NOARGS, ump_clear_doc},
5771 {"copy", (PyCFunction)ump_copy, METH_NOARGS, ump_copy_doc},
5772 {"__reduce__", (PyCFunction)ump_reduce, METH_VARARGS, ump_reduce_doc},
5773 {NULL, NULL} /* sentinel */
5774};
5775
5776static void
5777UnpicklerMemoProxy_dealloc(UnpicklerMemoProxyObject *self)
5778{
5779 PyObject_GC_UnTrack(self);
5780 Py_XDECREF(self->unpickler);
5781 PyObject_GC_Del((PyObject *)self);
5782}
5783
5784static int
5785UnpicklerMemoProxy_traverse(UnpicklerMemoProxyObject *self,
5786 visitproc visit, void *arg)
5787{
5788 Py_VISIT(self->unpickler);
5789 return 0;
5790}
5791
5792static int
5793UnpicklerMemoProxy_clear(UnpicklerMemoProxyObject *self)
5794{
5795 Py_CLEAR(self->unpickler);
5796 return 0;
5797}
5798
5799static PyTypeObject UnpicklerMemoProxyType = {
5800 PyVarObject_HEAD_INIT(NULL, 0)
5801 "_pickle.UnpicklerMemoProxy", /*tp_name*/
5802 sizeof(UnpicklerMemoProxyObject), /*tp_basicsize*/
5803 0,
5804 (destructor)UnpicklerMemoProxy_dealloc, /* tp_dealloc */
5805 0, /* tp_print */
5806 0, /* tp_getattr */
5807 0, /* tp_setattr */
5808 0, /* tp_compare */
5809 0, /* tp_repr */
5810 0, /* tp_as_number */
5811 0, /* tp_as_sequence */
5812 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00005813 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005814 0, /* tp_call */
5815 0, /* tp_str */
5816 PyObject_GenericGetAttr, /* tp_getattro */
5817 PyObject_GenericSetAttr, /* tp_setattro */
5818 0, /* tp_as_buffer */
5819 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5820 0, /* tp_doc */
5821 (traverseproc)UnpicklerMemoProxy_traverse, /* tp_traverse */
5822 (inquiry)UnpicklerMemoProxy_clear, /* tp_clear */
5823 0, /* tp_richcompare */
5824 0, /* tp_weaklistoffset */
5825 0, /* tp_iter */
5826 0, /* tp_iternext */
5827 unpicklerproxy_methods, /* tp_methods */
5828};
5829
5830static PyObject *
5831UnpicklerMemoProxy_New(UnpicklerObject *unpickler)
5832{
5833 UnpicklerMemoProxyObject *self;
5834
5835 self = PyObject_GC_New(UnpicklerMemoProxyObject,
5836 &UnpicklerMemoProxyType);
5837 if (self == NULL)
5838 return NULL;
5839 Py_INCREF(unpickler);
5840 self->unpickler = unpickler;
5841 PyObject_GC_Track(self);
5842 return (PyObject *)self;
5843}
5844
5845/*****************************************************************************/
5846
5847
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005848static PyObject *
5849Unpickler_get_memo(UnpicklerObject *self)
5850{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005851 return UnpicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005852}
5853
5854static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005855Unpickler_set_memo(UnpicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005856{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005857 PyObject **new_memo;
5858 Py_ssize_t new_memo_size = 0;
5859 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005860
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005861 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005862 PyErr_SetString(PyExc_TypeError,
5863 "attribute deletion is not supported");
5864 return -1;
5865 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005866
5867 if (Py_TYPE(obj) == &UnpicklerMemoProxyType) {
5868 UnpicklerObject *unpickler =
5869 ((UnpicklerMemoProxyObject *)obj)->unpickler;
5870
5871 new_memo_size = unpickler->memo_size;
5872 new_memo = _Unpickler_NewMemo(new_memo_size);
5873 if (new_memo == NULL)
5874 return -1;
5875
5876 for (i = 0; i < new_memo_size; i++) {
5877 Py_XINCREF(unpickler->memo[i]);
5878 new_memo[i] = unpickler->memo[i];
5879 }
5880 }
5881 else if (PyDict_Check(obj)) {
5882 Py_ssize_t i = 0;
5883 PyObject *key, *value;
5884
5885 new_memo_size = PyDict_Size(obj);
5886 new_memo = _Unpickler_NewMemo(new_memo_size);
5887 if (new_memo == NULL)
5888 return -1;
5889
5890 while (PyDict_Next(obj, &i, &key, &value)) {
5891 Py_ssize_t idx;
5892 if (!PyLong_Check(key)) {
5893 PyErr_SetString(PyExc_TypeError,
5894 "memo key must be integers");
5895 goto error;
5896 }
5897 idx = PyLong_AsSsize_t(key);
5898 if (idx == -1 && PyErr_Occurred())
5899 goto error;
5900 if (_Unpickler_MemoPut(self, idx, value) < 0)
5901 goto error;
5902 }
5903 }
5904 else {
5905 PyErr_Format(PyExc_TypeError,
5906 "'memo' attribute must be an UnpicklerMemoProxy object"
5907 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005908 return -1;
5909 }
5910
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005911 _Unpickler_MemoCleanup(self);
5912 self->memo_size = new_memo_size;
5913 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005914
5915 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005916
5917 error:
5918 if (new_memo_size) {
5919 i = new_memo_size;
5920 while (--i >= 0) {
5921 Py_XDECREF(new_memo[i]);
5922 }
5923 PyMem_FREE(new_memo);
5924 }
5925 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005926}
5927
5928static PyObject *
5929Unpickler_get_persload(UnpicklerObject *self)
5930{
5931 if (self->pers_func == NULL)
5932 PyErr_SetString(PyExc_AttributeError, "persistent_load");
5933 else
5934 Py_INCREF(self->pers_func);
5935 return self->pers_func;
5936}
5937
5938static int
5939Unpickler_set_persload(UnpicklerObject *self, PyObject *value)
5940{
5941 PyObject *tmp;
5942
5943 if (value == NULL) {
5944 PyErr_SetString(PyExc_TypeError,
5945 "attribute deletion is not supported");
5946 return -1;
5947 }
5948 if (!PyCallable_Check(value)) {
5949 PyErr_SetString(PyExc_TypeError,
5950 "persistent_load must be a callable taking "
5951 "one argument");
5952 return -1;
5953 }
5954
5955 tmp = self->pers_func;
5956 Py_INCREF(value);
5957 self->pers_func = value;
5958 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
5959
5960 return 0;
5961}
5962
5963static PyGetSetDef Unpickler_getsets[] = {
5964 {"memo", (getter)Unpickler_get_memo, (setter)Unpickler_set_memo},
5965 {"persistent_load", (getter)Unpickler_get_persload,
5966 (setter)Unpickler_set_persload},
5967 {NULL}
5968};
5969
5970static PyTypeObject Unpickler_Type = {
5971 PyVarObject_HEAD_INIT(NULL, 0)
5972 "_pickle.Unpickler", /*tp_name*/
5973 sizeof(UnpicklerObject), /*tp_basicsize*/
5974 0, /*tp_itemsize*/
5975 (destructor)Unpickler_dealloc, /*tp_dealloc*/
5976 0, /*tp_print*/
5977 0, /*tp_getattr*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005978 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00005979 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005980 0, /*tp_repr*/
5981 0, /*tp_as_number*/
5982 0, /*tp_as_sequence*/
5983 0, /*tp_as_mapping*/
5984 0, /*tp_hash*/
5985 0, /*tp_call*/
5986 0, /*tp_str*/
5987 0, /*tp_getattro*/
5988 0, /*tp_setattro*/
5989 0, /*tp_as_buffer*/
5990 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5991 Unpickler_doc, /*tp_doc*/
5992 (traverseproc)Unpickler_traverse, /*tp_traverse*/
5993 (inquiry)Unpickler_clear, /*tp_clear*/
5994 0, /*tp_richcompare*/
5995 0, /*tp_weaklistoffset*/
5996 0, /*tp_iter*/
5997 0, /*tp_iternext*/
5998 Unpickler_methods, /*tp_methods*/
5999 0, /*tp_members*/
6000 Unpickler_getsets, /*tp_getset*/
6001 0, /*tp_base*/
6002 0, /*tp_dict*/
6003 0, /*tp_descr_get*/
6004 0, /*tp_descr_set*/
6005 0, /*tp_dictoffset*/
6006 (initproc)Unpickler_init, /*tp_init*/
6007 PyType_GenericAlloc, /*tp_alloc*/
6008 PyType_GenericNew, /*tp_new*/
6009 PyObject_GC_Del, /*tp_free*/
6010 0, /*tp_is_gc*/
6011};
6012
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006013PyDoc_STRVAR(pickle_dump_doc,
6014"dump(obj, file, protocol=None, *, fix_imports=True) -> None\n"
6015"\n"
6016"Write a pickled representation of obj to the open file object file. This\n"
6017"is equivalent to ``Pickler(file, protocol).dump(obj)``, but may be more\n"
6018"efficient.\n"
6019"\n"
6020"The optional protocol argument tells the pickler to use the given protocol;\n"
6021"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
6022"backward-incompatible protocol designed for Python 3.0.\n"
6023"\n"
6024"Specifying a negative protocol version selects the highest protocol version\n"
6025"supported. The higher the protocol used, the more recent the version of\n"
6026"Python needed to read the pickle produced.\n"
6027"\n"
6028"The file argument must have a write() method that accepts a single bytes\n"
6029"argument. It can thus be a file object opened for binary writing, a\n"
6030"io.BytesIO instance, or any other custom object that meets this interface.\n"
6031"\n"
6032"If fix_imports is True and protocol is less than 3, pickle will try to\n"
6033"map the new Python 3.x names to the old module names used in Python 2.x,\n"
6034"so that the pickle data stream is readable with Python 2.x.\n");
6035
6036static PyObject *
6037pickle_dump(PyObject *self, PyObject *args, PyObject *kwds)
6038{
6039 static char *kwlist[] = {"obj", "file", "protocol", "fix_imports", 0};
6040 PyObject *obj;
6041 PyObject *file;
6042 PyObject *proto = NULL;
6043 PyObject *fix_imports = Py_True;
6044 PicklerObject *pickler;
6045
6046 /* fix_imports is a keyword-only argument. */
6047 if (Py_SIZE(args) > 3) {
6048 PyErr_Format(PyExc_TypeError,
6049 "pickle.dump() takes at most 3 positional "
6050 "argument (%zd given)", Py_SIZE(args));
6051 return NULL;
6052 }
6053
6054 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|OO:dump", kwlist,
6055 &obj, &file, &proto, &fix_imports))
6056 return NULL;
6057
6058 pickler = _Pickler_New();
6059 if (pickler == NULL)
6060 return NULL;
6061
6062 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6063 goto error;
6064
6065 if (_Pickler_SetOutputStream(pickler, file) < 0)
6066 goto error;
6067
6068 if (dump(pickler, obj) < 0)
6069 goto error;
6070
6071 if (_Pickler_FlushToFile(pickler) < 0)
6072 goto error;
6073
6074 Py_DECREF(pickler);
6075 Py_RETURN_NONE;
6076
6077 error:
6078 Py_XDECREF(pickler);
6079 return NULL;
6080}
6081
6082PyDoc_STRVAR(pickle_dumps_doc,
6083"dumps(obj, protocol=None, *, fix_imports=True) -> bytes\n"
6084"\n"
6085"Return the pickled representation of the object as a bytes\n"
6086"object, instead of writing it to a file.\n"
6087"\n"
6088"The optional protocol argument tells the pickler to use the given protocol;\n"
6089"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
6090"backward-incompatible protocol designed for Python 3.0.\n"
6091"\n"
6092"Specifying a negative protocol version selects the highest protocol version\n"
6093"supported. The higher the protocol used, the more recent the version of\n"
6094"Python needed to read the pickle produced.\n"
6095"\n"
6096"If fix_imports is True and *protocol* is less than 3, pickle will try to\n"
6097"map the new Python 3.x names to the old module names used in Python 2.x,\n"
6098"so that the pickle data stream is readable with Python 2.x.\n");
6099
6100static PyObject *
6101pickle_dumps(PyObject *self, PyObject *args, PyObject *kwds)
6102{
6103 static char *kwlist[] = {"obj", "protocol", "fix_imports", 0};
6104 PyObject *obj;
6105 PyObject *proto = NULL;
6106 PyObject *result;
6107 PyObject *fix_imports = Py_True;
6108 PicklerObject *pickler;
6109
6110 /* fix_imports is a keyword-only argument. */
6111 if (Py_SIZE(args) > 2) {
6112 PyErr_Format(PyExc_TypeError,
6113 "pickle.dumps() takes at most 2 positional "
6114 "argument (%zd given)", Py_SIZE(args));
6115 return NULL;
6116 }
6117
6118 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:dumps", kwlist,
6119 &obj, &proto, &fix_imports))
6120 return NULL;
6121
6122 pickler = _Pickler_New();
6123 if (pickler == NULL)
6124 return NULL;
6125
6126 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6127 goto error;
6128
6129 if (dump(pickler, obj) < 0)
6130 goto error;
6131
6132 result = _Pickler_GetString(pickler);
6133 Py_DECREF(pickler);
6134 return result;
6135
6136 error:
6137 Py_XDECREF(pickler);
6138 return NULL;
6139}
6140
6141PyDoc_STRVAR(pickle_load_doc,
6142"load(file, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6143"\n"
6144"Read a pickled object representation from the open file object file and\n"
6145"return the reconstituted object hierarchy specified therein. This is\n"
6146"equivalent to ``Unpickler(file).load()``, but may be more efficient.\n"
6147"\n"
6148"The protocol version of the pickle is detected automatically, so no protocol\n"
6149"argument is needed. Bytes past the pickled object's representation are\n"
6150"ignored.\n"
6151"\n"
6152"The argument file must have two methods, a read() method that takes an\n"
6153"integer argument, and a readline() method that requires no arguments. Both\n"
6154"methods should return bytes. Thus *file* can be a binary file object opened\n"
6155"for reading, a BytesIO object, or any other custom object that meets this\n"
6156"interface.\n"
6157"\n"
6158"Optional keyword arguments are fix_imports, encoding and errors,\n"
6159"which are used to control compatiblity support for pickle stream generated\n"
6160"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6161"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6162"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6163"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6164
6165static PyObject *
6166pickle_load(PyObject *self, PyObject *args, PyObject *kwds)
6167{
6168 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
6169 PyObject *file;
6170 PyObject *fix_imports = Py_True;
6171 PyObject *result;
6172 char *encoding = NULL;
6173 char *errors = NULL;
6174 UnpicklerObject *unpickler;
6175
6176 /* fix_imports, encoding and errors are a keyword-only argument. */
6177 if (Py_SIZE(args) != 1) {
6178 PyErr_Format(PyExc_TypeError,
6179 "pickle.load() takes exactly one positional "
6180 "argument (%zd given)", Py_SIZE(args));
6181 return NULL;
6182 }
6183
6184 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:load", kwlist,
6185 &file, &fix_imports, &encoding, &errors))
6186 return NULL;
6187
6188 unpickler = _Unpickler_New();
6189 if (unpickler == NULL)
6190 return NULL;
6191
6192 if (_Unpickler_SetInputStream(unpickler, file) < 0)
6193 goto error;
6194
6195 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6196 goto error;
6197
6198 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6199 if (unpickler->fix_imports == -1)
6200 goto error;
6201
6202 result = load(unpickler);
6203 Py_DECREF(unpickler);
6204 return result;
6205
6206 error:
6207 Py_XDECREF(unpickler);
6208 return NULL;
6209}
6210
6211PyDoc_STRVAR(pickle_loads_doc,
6212"loads(input, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6213"\n"
6214"Read a pickled object hierarchy from a bytes object and return the\n"
6215"reconstituted object hierarchy specified therein\n"
6216"\n"
6217"The protocol version of the pickle is detected automatically, so no protocol\n"
6218"argument is needed. Bytes past the pickled object's representation are\n"
6219"ignored.\n"
6220"\n"
6221"Optional keyword arguments are fix_imports, encoding and errors, which\n"
6222"are used to control compatiblity support for pickle stream generated\n"
6223"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6224"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6225"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6226"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6227
6228static PyObject *
6229pickle_loads(PyObject *self, PyObject *args, PyObject *kwds)
6230{
6231 static char *kwlist[] = {"input", "fix_imports", "encoding", "errors", 0};
6232 PyObject *input;
6233 PyObject *fix_imports = Py_True;
6234 PyObject *result;
6235 char *encoding = NULL;
6236 char *errors = NULL;
6237 UnpicklerObject *unpickler;
6238
6239 /* fix_imports, encoding and errors are a keyword-only argument. */
6240 if (Py_SIZE(args) != 1) {
6241 PyErr_Format(PyExc_TypeError,
6242 "pickle.loads() takes exactly one positional "
6243 "argument (%zd given)", Py_SIZE(args));
6244 return NULL;
6245 }
6246
6247 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:loads", kwlist,
6248 &input, &fix_imports, &encoding, &errors))
6249 return NULL;
6250
6251 unpickler = _Unpickler_New();
6252 if (unpickler == NULL)
6253 return NULL;
6254
6255 if (_Unpickler_SetStringInput(unpickler, input) < 0)
6256 goto error;
6257
6258 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6259 goto error;
6260
6261 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6262 if (unpickler->fix_imports == -1)
6263 goto error;
6264
6265 result = load(unpickler);
6266 Py_DECREF(unpickler);
6267 return result;
6268
6269 error:
6270 Py_XDECREF(unpickler);
6271 return NULL;
6272}
6273
6274
6275static struct PyMethodDef pickle_methods[] = {
6276 {"dump", (PyCFunction)pickle_dump, METH_VARARGS|METH_KEYWORDS,
6277 pickle_dump_doc},
6278 {"dumps", (PyCFunction)pickle_dumps, METH_VARARGS|METH_KEYWORDS,
6279 pickle_dumps_doc},
6280 {"load", (PyCFunction)pickle_load, METH_VARARGS|METH_KEYWORDS,
6281 pickle_load_doc},
6282 {"loads", (PyCFunction)pickle_loads, METH_VARARGS|METH_KEYWORDS,
6283 pickle_loads_doc},
6284 {NULL, NULL} /* sentinel */
6285};
6286
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006287static int
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006288initmodule(void)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006289{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006290 PyObject *copyreg = NULL;
6291 PyObject *compat_pickle = NULL;
6292
6293 /* XXX: We should ensure that the types of the dictionaries imported are
6294 exactly PyDict objects. Otherwise, it is possible to crash the pickle
6295 since we use the PyDict API directly to access these dictionaries. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006296
6297 copyreg = PyImport_ImportModule("copyreg");
6298 if (!copyreg)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006299 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006300 dispatch_table = PyObject_GetAttrString(copyreg, "dispatch_table");
6301 if (!dispatch_table)
6302 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006303 extension_registry = \
6304 PyObject_GetAttrString(copyreg, "_extension_registry");
6305 if (!extension_registry)
6306 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006307 inverted_registry = PyObject_GetAttrString(copyreg, "_inverted_registry");
6308 if (!inverted_registry)
6309 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006310 extension_cache = PyObject_GetAttrString(copyreg, "_extension_cache");
6311 if (!extension_cache)
6312 goto error;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006313 Py_CLEAR(copyreg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006314
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006315 /* Load the 2.x -> 3.x stdlib module mapping tables */
6316 compat_pickle = PyImport_ImportModule("_compat_pickle");
6317 if (!compat_pickle)
6318 goto error;
6319 name_mapping_2to3 = PyObject_GetAttrString(compat_pickle, "NAME_MAPPING");
6320 if (!name_mapping_2to3)
6321 goto error;
6322 if (!PyDict_CheckExact(name_mapping_2to3)) {
6323 PyErr_Format(PyExc_RuntimeError,
6324 "_compat_pickle.NAME_MAPPING should be a dict, not %.200s",
6325 Py_TYPE(name_mapping_2to3)->tp_name);
6326 goto error;
6327 }
6328 import_mapping_2to3 = PyObject_GetAttrString(compat_pickle,
6329 "IMPORT_MAPPING");
6330 if (!import_mapping_2to3)
6331 goto error;
6332 if (!PyDict_CheckExact(import_mapping_2to3)) {
6333 PyErr_Format(PyExc_RuntimeError,
6334 "_compat_pickle.IMPORT_MAPPING should be a dict, "
6335 "not %.200s", Py_TYPE(import_mapping_2to3)->tp_name);
6336 goto error;
6337 }
6338 /* ... and the 3.x -> 2.x mapping tables */
6339 name_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6340 "REVERSE_NAME_MAPPING");
6341 if (!name_mapping_3to2)
6342 goto error;
6343 if (!PyDict_CheckExact(name_mapping_3to2)) {
6344 PyErr_Format(PyExc_RuntimeError,
Ezio Melotti13925002011-03-16 11:05:33 +02006345 "_compat_pickle.REVERSE_NAME_MAPPING should be a dict, "
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006346 "not %.200s", Py_TYPE(name_mapping_3to2)->tp_name);
6347 goto error;
6348 }
6349 import_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6350 "REVERSE_IMPORT_MAPPING");
6351 if (!import_mapping_3to2)
6352 goto error;
6353 if (!PyDict_CheckExact(import_mapping_3to2)) {
6354 PyErr_Format(PyExc_RuntimeError,
6355 "_compat_pickle.REVERSE_IMPORT_MAPPING should be a dict, "
6356 "not %.200s", Py_TYPE(import_mapping_3to2)->tp_name);
6357 goto error;
6358 }
6359 Py_CLEAR(compat_pickle);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006360
6361 empty_tuple = PyTuple_New(0);
6362 if (empty_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006363 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006364 two_tuple = PyTuple_New(2);
6365 if (two_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006366 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006367 /* We use this temp container with no regard to refcounts, or to
6368 * keeping containees alive. Exempt from GC, because we don't
6369 * want anything looking at two_tuple() by magic.
6370 */
6371 PyObject_GC_UnTrack(two_tuple);
6372
6373 return 0;
6374
6375 error:
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006376 Py_CLEAR(copyreg);
6377 Py_CLEAR(dispatch_table);
6378 Py_CLEAR(extension_registry);
6379 Py_CLEAR(inverted_registry);
6380 Py_CLEAR(extension_cache);
6381 Py_CLEAR(compat_pickle);
6382 Py_CLEAR(name_mapping_2to3);
6383 Py_CLEAR(import_mapping_2to3);
6384 Py_CLEAR(name_mapping_3to2);
6385 Py_CLEAR(import_mapping_3to2);
6386 Py_CLEAR(empty_tuple);
6387 Py_CLEAR(two_tuple);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006388 return -1;
6389}
6390
6391static struct PyModuleDef _picklemodule = {
6392 PyModuleDef_HEAD_INIT,
6393 "_pickle",
6394 pickle_module_doc,
6395 -1,
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006396 pickle_methods,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006397 NULL,
6398 NULL,
6399 NULL,
6400 NULL
6401};
6402
6403PyMODINIT_FUNC
6404PyInit__pickle(void)
6405{
6406 PyObject *m;
6407
6408 if (PyType_Ready(&Unpickler_Type) < 0)
6409 return NULL;
6410 if (PyType_Ready(&Pickler_Type) < 0)
6411 return NULL;
6412 if (PyType_Ready(&Pdata_Type) < 0)
6413 return NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006414 if (PyType_Ready(&PicklerMemoProxyType) < 0)
6415 return NULL;
6416 if (PyType_Ready(&UnpicklerMemoProxyType) < 0)
6417 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006418
6419 /* Create the module and add the functions. */
6420 m = PyModule_Create(&_picklemodule);
6421 if (m == NULL)
6422 return NULL;
6423
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006424 Py_INCREF(&Pickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006425 if (PyModule_AddObject(m, "Pickler", (PyObject *)&Pickler_Type) < 0)
6426 return NULL;
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006427 Py_INCREF(&Unpickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006428 if (PyModule_AddObject(m, "Unpickler", (PyObject *)&Unpickler_Type) < 0)
6429 return NULL;
6430
6431 /* Initialize the exceptions. */
6432 PickleError = PyErr_NewException("_pickle.PickleError", NULL, NULL);
6433 if (PickleError == NULL)
6434 return NULL;
6435 PicklingError = \
6436 PyErr_NewException("_pickle.PicklingError", PickleError, NULL);
6437 if (PicklingError == NULL)
6438 return NULL;
6439 UnpicklingError = \
6440 PyErr_NewException("_pickle.UnpicklingError", PickleError, NULL);
6441 if (UnpicklingError == NULL)
6442 return NULL;
6443
6444 if (PyModule_AddObject(m, "PickleError", PickleError) < 0)
6445 return NULL;
6446 if (PyModule_AddObject(m, "PicklingError", PicklingError) < 0)
6447 return NULL;
6448 if (PyModule_AddObject(m, "UnpicklingError", UnpicklingError) < 0)
6449 return NULL;
6450
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006451 if (initmodule() < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006452 return NULL;
6453
6454 return m;
6455}