blob: 001360b8610931351f2886e5cdf73a56301fb62c [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 Pitrouea99c5c2010-09-09 18:33:21 +0000156 int i = Py_SIZE(self);
157 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
193Pdata_clear(Pdata *self, int clearto)
194{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000195 int 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;
306 long me_value;
307} 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 */
322 PyObject *arg;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000323
324 PyObject *write; /* write() method of the output stream. */
325 PyObject *output_buffer; /* Write into a local bytearray buffer before
326 flushing to the stream. */
327 Py_ssize_t output_len; /* Length of output_buffer. */
328 Py_ssize_t max_output_len; /* Allocation size of output_buffer. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000329 int proto; /* Pickle protocol number, >= 0 */
330 int bin; /* Boolean, true if proto > 0 */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000331 int buf_size; /* Size of the current buffered pickle data */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000332 int fast; /* Enable fast mode if set to a true value.
333 The fast mode disable the usage of memo,
334 therefore speeding the pickling process by
335 not generating superfluous PUT opcodes. It
336 should not be used if with self-referential
337 objects. */
338 int fast_nesting;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000339 int fix_imports; /* Indicate whether Pickler should fix
340 the name of globals for Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000341 PyObject *fast_memo;
342} PicklerObject;
343
344typedef struct UnpicklerObject {
345 PyObject_HEAD
346 Pdata *stack; /* Pickle data stack, store unpickled objects. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000347
348 /* The unpickler memo is just an array of PyObject *s. Using a dict
349 is unnecessary, since the keys are contiguous ints. */
350 PyObject **memo;
351 Py_ssize_t memo_size;
352
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000353 PyObject *arg;
354 PyObject *pers_func; /* persistent_load() method, can be NULL. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000355
356 Py_buffer buffer;
357 char *input_buffer;
358 char *input_line;
359 Py_ssize_t input_len;
360 Py_ssize_t next_read_idx;
Antoine Pitrou04248a82010-10-12 20:51:21 +0000361 Py_ssize_t prefetched_idx; /* index of first prefetched byte */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000362 PyObject *read; /* read() method of the input stream. */
363 PyObject *readline; /* readline() method of the input stream. */
Antoine Pitrou04248a82010-10-12 20:51:21 +0000364 PyObject *peek; /* peek() method of the input stream, or NULL */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000365
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000366 char *encoding; /* Name of the encoding to be used for
367 decoding strings pickled using Python
368 2.x. The default value is "ASCII" */
369 char *errors; /* Name of errors handling scheme to used when
370 decoding strings. The default value is
371 "strict". */
372 int *marks; /* Mark stack, used for unpickling container
373 objects. */
374 Py_ssize_t num_marks; /* Number of marks in the mark stack. */
375 Py_ssize_t marks_size; /* Current allocated size of the mark stack. */
Antoine Pitroud9dfaa92009-06-04 20:32:06 +0000376 int proto; /* Protocol of the pickle loaded. */
377 int fix_imports; /* Indicate whether Unpickler should fix
378 the name of globals pickled by Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000379} UnpicklerObject;
380
381/* Forward declarations */
382static int save(PicklerObject *, PyObject *, int);
383static int save_reduce(PicklerObject *, PyObject *, PyObject *);
384static PyTypeObject Pickler_Type;
385static PyTypeObject Unpickler_Type;
386
387
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000388/*************************************************************************
389 A custom hashtable mapping void* to longs. This is used by the pickler for
390 memoization. Using a custom hashtable rather than PyDict allows us to skip
391 a bunch of unnecessary object creation. This makes a huge performance
392 difference. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000393
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000394#define MT_MINSIZE 8
395#define PERTURB_SHIFT 5
396
397
398static PyMemoTable *
399PyMemoTable_New(void)
400{
401 PyMemoTable *memo = PyMem_MALLOC(sizeof(PyMemoTable));
402 if (memo == NULL) {
403 PyErr_NoMemory();
404 return NULL;
405 }
406
407 memo->mt_used = 0;
408 memo->mt_allocated = MT_MINSIZE;
409 memo->mt_mask = MT_MINSIZE - 1;
410 memo->mt_table = PyMem_MALLOC(MT_MINSIZE * sizeof(PyMemoEntry));
411 if (memo->mt_table == NULL) {
412 PyMem_FREE(memo);
413 PyErr_NoMemory();
414 return NULL;
415 }
416 memset(memo->mt_table, 0, MT_MINSIZE * sizeof(PyMemoEntry));
417
418 return memo;
419}
420
421static PyMemoTable *
422PyMemoTable_Copy(PyMemoTable *self)
423{
424 Py_ssize_t i;
425 PyMemoTable *new = PyMemoTable_New();
426 if (new == NULL)
427 return NULL;
428
429 new->mt_used = self->mt_used;
430 new->mt_allocated = self->mt_allocated;
431 new->mt_mask = self->mt_mask;
432 /* The table we get from _New() is probably smaller than we wanted.
433 Free it and allocate one that's the right size. */
434 PyMem_FREE(new->mt_table);
435 new->mt_table = PyMem_MALLOC(self->mt_allocated * sizeof(PyMemoEntry));
436 if (new->mt_table == NULL) {
437 PyMem_FREE(new);
438 return NULL;
439 }
440 for (i = 0; i < self->mt_allocated; i++) {
441 Py_XINCREF(self->mt_table[i].me_key);
442 }
443 memcpy(new->mt_table, self->mt_table,
444 sizeof(PyMemoEntry) * self->mt_allocated);
445
446 return new;
447}
448
449static Py_ssize_t
450PyMemoTable_Size(PyMemoTable *self)
451{
452 return self->mt_used;
453}
454
455static int
456PyMemoTable_Clear(PyMemoTable *self)
457{
458 Py_ssize_t i = self->mt_allocated;
459
460 while (--i >= 0) {
461 Py_XDECREF(self->mt_table[i].me_key);
462 }
463 self->mt_used = 0;
464 memset(self->mt_table, 0, self->mt_allocated * sizeof(PyMemoEntry));
465 return 0;
466}
467
468static void
469PyMemoTable_Del(PyMemoTable *self)
470{
471 if (self == NULL)
472 return;
473 PyMemoTable_Clear(self);
474
475 PyMem_FREE(self->mt_table);
476 PyMem_FREE(self);
477}
478
479/* Since entries cannot be deleted from this hashtable, _PyMemoTable_Lookup()
480 can be considerably simpler than dictobject.c's lookdict(). */
481static PyMemoEntry *
482_PyMemoTable_Lookup(PyMemoTable *self, PyObject *key)
483{
484 size_t i;
485 size_t perturb;
486 size_t mask = (size_t)self->mt_mask;
487 PyMemoEntry *table = self->mt_table;
488 PyMemoEntry *entry;
Benjamin Peterson8f67d082010-10-17 20:54:53 +0000489 Py_hash_t hash = (Py_hash_t)key >> 3;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000490
491 i = hash & mask;
492 entry = &table[i];
493 if (entry->me_key == NULL || entry->me_key == key)
494 return entry;
495
496 for (perturb = hash; ; perturb >>= PERTURB_SHIFT) {
497 i = (i << 2) + i + perturb + 1;
498 entry = &table[i & mask];
499 if (entry->me_key == NULL || entry->me_key == key)
500 return entry;
501 }
502 assert(0); /* Never reached */
503 return NULL;
504}
505
506/* Returns -1 on failure, 0 on success. */
507static int
508_PyMemoTable_ResizeTable(PyMemoTable *self, Py_ssize_t min_size)
509{
510 PyMemoEntry *oldtable = NULL;
511 PyMemoEntry *oldentry, *newentry;
512 Py_ssize_t new_size = MT_MINSIZE;
513 Py_ssize_t to_process;
514
515 assert(min_size > 0);
516
517 /* Find the smallest valid table size >= min_size. */
518 while (new_size < min_size && new_size > 0)
519 new_size <<= 1;
520 if (new_size <= 0) {
521 PyErr_NoMemory();
522 return -1;
523 }
524 /* new_size needs to be a power of two. */
525 assert((new_size & (new_size - 1)) == 0);
526
527 /* Allocate new table. */
528 oldtable = self->mt_table;
529 self->mt_table = PyMem_MALLOC(new_size * sizeof(PyMemoEntry));
530 if (self->mt_table == NULL) {
531 PyMem_FREE(oldtable);
532 PyErr_NoMemory();
533 return -1;
534 }
535 self->mt_allocated = new_size;
536 self->mt_mask = new_size - 1;
537 memset(self->mt_table, 0, sizeof(PyMemoEntry) * new_size);
538
539 /* Copy entries from the old table. */
540 to_process = self->mt_used;
541 for (oldentry = oldtable; to_process > 0; oldentry++) {
542 if (oldentry->me_key != NULL) {
543 to_process--;
544 /* newentry is a pointer to a chunk of the new
545 mt_table, so we're setting the key:value pair
546 in-place. */
547 newentry = _PyMemoTable_Lookup(self, oldentry->me_key);
548 newentry->me_key = oldentry->me_key;
549 newentry->me_value = oldentry->me_value;
550 }
551 }
552
553 /* Deallocate the old table. */
554 PyMem_FREE(oldtable);
555 return 0;
556}
557
558/* Returns NULL on failure, a pointer to the value otherwise. */
559static long *
560PyMemoTable_Get(PyMemoTable *self, PyObject *key)
561{
562 PyMemoEntry *entry = _PyMemoTable_Lookup(self, key);
563 if (entry->me_key == NULL)
564 return NULL;
565 return &entry->me_value;
566}
567
568/* Returns -1 on failure, 0 on success. */
569static int
570PyMemoTable_Set(PyMemoTable *self, PyObject *key, long value)
571{
572 PyMemoEntry *entry;
573
574 assert(key != NULL);
575
576 entry = _PyMemoTable_Lookup(self, key);
577 if (entry->me_key != NULL) {
578 entry->me_value = value;
579 return 0;
580 }
581 Py_INCREF(key);
582 entry->me_key = key;
583 entry->me_value = value;
584 self->mt_used++;
585
586 /* If we added a key, we can safely resize. Otherwise just return!
587 * If used >= 2/3 size, adjust size. Normally, this quaduples the size.
588 *
589 * Quadrupling the size improves average table sparseness
590 * (reducing collisions) at the cost of some memory. It also halves
591 * the number of expensive resize operations in a growing memo table.
592 *
593 * Very large memo tables (over 50K items) use doubling instead.
594 * This may help applications with severe memory constraints.
595 */
596 if (!(self->mt_used * 3 >= (self->mt_mask + 1) * 2))
597 return 0;
598 return _PyMemoTable_ResizeTable(self,
599 (self->mt_used > 50000 ? 2 : 4) * self->mt_used);
600}
601
602#undef MT_MINSIZE
603#undef PERTURB_SHIFT
604
605/*************************************************************************/
606
607/* Helpers for creating the argument tuple passed to functions. This has the
608 performance advantage of calling PyTuple_New() only once.
609
610 XXX(avassalotti): Inline directly in _Pickler_FastCall() and
611 _Unpickler_FastCall(). */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000612#define ARG_TUP(self, obj) do { \
613 if ((self)->arg || ((self)->arg=PyTuple_New(1))) { \
614 Py_XDECREF(PyTuple_GET_ITEM((self)->arg, 0)); \
615 PyTuple_SET_ITEM((self)->arg, 0, (obj)); \
616 } \
617 else { \
618 Py_DECREF((obj)); \
619 } \
620 } while (0)
621
622#define FREE_ARG_TUP(self) do { \
623 if ((self)->arg->ob_refcnt > 1) \
624 Py_CLEAR((self)->arg); \
625 } while (0)
626
627/* A temporary cleaner API for fast single argument function call.
628
629 XXX: Does caching the argument tuple provides any real performance benefits?
630
631 A quick benchmark, on a 2.0GHz Athlon64 3200+ running Linux 2.6.24 with
632 glibc 2.7, tells me that it takes roughly 20,000,000 PyTuple_New(1) calls
633 when the tuple is retrieved from the freelist (i.e, call PyTuple_New() then
634 immediately DECREF it) and 1,200,000 calls when allocating brand new tuples
635 (i.e, call PyTuple_New() and store the returned value in an array), to save
636 one second (wall clock time). Either ways, the loading time a pickle stream
637 large enough to generate this number of calls would be massively
638 overwhelmed by other factors, like I/O throughput, the GC traversal and
639 object allocation overhead. So, I really doubt these functions provide any
640 real benefits.
641
642 On the other hand, oprofile reports that pickle spends a lot of time in
643 these functions. But, that is probably more related to the function call
644 overhead, than the argument tuple allocation.
645
646 XXX: And, what is the reference behavior of these? Steal, borrow? At first
647 glance, it seems to steal the reference of 'arg' and borrow the reference
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000648 of 'func'. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000649static PyObject *
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000650_Pickler_FastCall(PicklerObject *self, PyObject *func, PyObject *arg)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000651{
652 PyObject *result = NULL;
653
654 ARG_TUP(self, arg);
655 if (self->arg) {
656 result = PyObject_Call(func, self->arg, NULL);
657 FREE_ARG_TUP(self);
658 }
659 return result;
660}
661
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000662static int
663_Pickler_ClearBuffer(PicklerObject *self)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000664{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000665 Py_CLEAR(self->output_buffer);
666 self->output_buffer =
667 PyBytes_FromStringAndSize(NULL, self->max_output_len);
668 if (self->output_buffer == NULL)
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +0000669 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000670 self->output_len = 0;
671 return 0;
672}
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +0000673
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000674static PyObject *
675_Pickler_GetString(PicklerObject *self)
676{
677 PyObject *output_buffer = self->output_buffer;
678
679 assert(self->output_buffer != NULL);
680 self->output_buffer = NULL;
681 /* Resize down to exact size */
682 if (_PyBytes_Resize(&output_buffer, self->output_len) < 0)
683 return NULL;
684 return output_buffer;
685}
686
687static int
688_Pickler_FlushToFile(PicklerObject *self)
689{
690 PyObject *output, *result;
691
692 assert(self->write != NULL);
693
694 output = _Pickler_GetString(self);
695 if (output == NULL)
696 return -1;
697
698 result = _Pickler_FastCall(self, self->write, output);
699 Py_XDECREF(result);
700 return (result == NULL) ? -1 : 0;
701}
702
703static int
704_Pickler_Write(PicklerObject *self, const char *s, Py_ssize_t n)
705{
706 Py_ssize_t i, required;
707 char *buffer;
708
709 assert(s != NULL);
710
711 required = self->output_len + n;
712 if (required > self->max_output_len) {
713 if (self->write != NULL && required > MAX_WRITE_BUF_SIZE) {
714 /* XXX This reallocates a new buffer every time, which is a bit
715 wasteful. */
716 if (_Pickler_FlushToFile(self) < 0)
717 return -1;
718 if (_Pickler_ClearBuffer(self) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000719 return -1;
720 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000721 if (self->write != NULL && n > MAX_WRITE_BUF_SIZE) {
722 /* we already flushed above, so the buffer is empty */
723 PyObject *result;
724 /* XXX we could spare an intermediate copy and pass
725 a memoryview instead */
726 PyObject *output = PyBytes_FromStringAndSize(s, n);
727 if (s == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000728 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000729 result = _Pickler_FastCall(self, self->write, output);
730 Py_XDECREF(result);
731 return (result == NULL) ? -1 : 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000732 }
733 else {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000734 if (self->output_len >= PY_SSIZE_T_MAX / 2 - n) {
735 PyErr_NoMemory();
736 return -1;
737 }
738 self->max_output_len = (self->output_len + n) * 2;
739 if (_PyBytes_Resize(&self->output_buffer, self->max_output_len) < 0)
740 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000741 }
742 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000743 buffer = PyBytes_AS_STRING(self->output_buffer);
744 if (n < 8) {
745 /* This is faster than memcpy when the string is short. */
746 for (i = 0; i < n; i++) {
747 buffer[self->output_len + i] = s[i];
748 }
749 }
750 else {
751 memcpy(buffer + self->output_len, s, n);
752 }
753 self->output_len += n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000754 return n;
755}
756
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000757static PicklerObject *
758_Pickler_New(void)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000759{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000760 PicklerObject *self;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000761
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000762 self = PyObject_GC_New(PicklerObject, &Pickler_Type);
763 if (self == NULL)
764 return NULL;
765
766 self->pers_func = NULL;
767 self->arg = NULL;
768 self->write = NULL;
769 self->proto = 0;
770 self->bin = 0;
771 self->fast = 0;
772 self->fast_nesting = 0;
773 self->fix_imports = 0;
774 self->fast_memo = NULL;
775
776 self->memo = PyMemoTable_New();
777 if (self->memo == NULL) {
778 Py_DECREF(self);
779 return NULL;
780 }
781 self->max_output_len = WRITE_BUF_SIZE;
782 self->output_len = 0;
783 self->output_buffer = PyBytes_FromStringAndSize(NULL,
784 self->max_output_len);
785 if (self->output_buffer == NULL) {
786 Py_DECREF(self);
787 return NULL;
788 }
789 return self;
790}
791
792static int
793_Pickler_SetProtocol(PicklerObject *self, PyObject *proto_obj,
794 PyObject *fix_imports_obj)
795{
796 long proto = 0;
797 int fix_imports;
798
799 if (proto_obj == NULL || proto_obj == Py_None)
800 proto = DEFAULT_PROTOCOL;
801 else {
802 proto = PyLong_AsLong(proto_obj);
803 if (proto == -1 && PyErr_Occurred())
804 return -1;
805 }
806 if (proto < 0)
807 proto = HIGHEST_PROTOCOL;
808 if (proto > HIGHEST_PROTOCOL) {
809 PyErr_Format(PyExc_ValueError, "pickle protocol must be <= %d",
810 HIGHEST_PROTOCOL);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000811 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000812 }
813 fix_imports = PyObject_IsTrue(fix_imports_obj);
814 if (fix_imports == -1)
815 return -1;
816
817 self->proto = proto;
818 self->bin = proto > 0;
819 self->fix_imports = fix_imports && proto < 3;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000820
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000821 return 0;
822}
823
824/* Returns -1 (with an exception set) on failure, 0 on success. This may
825 be called once on a freshly created Pickler. */
826static int
827_Pickler_SetOutputStream(PicklerObject *self, PyObject *file)
828{
829 assert(file != NULL);
830 self->write = PyObject_GetAttrString(file, "write");
831 if (self->write == NULL) {
832 if (PyErr_ExceptionMatches(PyExc_AttributeError))
833 PyErr_SetString(PyExc_TypeError,
834 "file must have a 'write' attribute");
835 return -1;
836 }
837
838 return 0;
839}
840
841/* See documentation for _Pickler_FastCall(). */
842static PyObject *
843_Unpickler_FastCall(UnpicklerObject *self, PyObject *func, PyObject *arg)
844{
845 PyObject *result = NULL;
846
847 ARG_TUP(self, arg);
848 if (self->arg) {
849 result = PyObject_Call(func, self->arg, NULL);
850 FREE_ARG_TUP(self);
851 }
852 return result;
853}
854
855/* Returns the size of the input on success, -1 on failure. This takes its
856 own reference to `input`. */
857static Py_ssize_t
858_Unpickler_SetStringInput(UnpicklerObject *self, PyObject *input)
859{
860 if (self->buffer.buf != NULL)
861 PyBuffer_Release(&self->buffer);
862 if (PyObject_GetBuffer(input, &self->buffer, PyBUF_CONTIG_RO) < 0)
863 return -1;
864 self->input_buffer = self->buffer.buf;
865 self->input_len = self->buffer.len;
866 self->next_read_idx = 0;
Antoine Pitrou04248a82010-10-12 20:51:21 +0000867 self->prefetched_idx = self->input_len;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000868 return self->input_len;
869}
870
Antoine Pitrou04248a82010-10-12 20:51:21 +0000871static int
872_Unpickler_SkipConsumed(UnpicklerObject *self)
873{
874 Py_ssize_t consumed = self->next_read_idx - self->prefetched_idx;
875
876 if (consumed > 0) {
877 PyObject *r;
878 assert(self->peek); /* otherwise we did something wrong */
879 /* This makes an useless copy... */
880 r = PyObject_CallFunction(self->read, "n", consumed);
881 if (r == NULL)
882 return -1;
883 Py_DECREF(r);
884 self->prefetched_idx = self->next_read_idx;
885 }
886 return 0;
887}
888
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000889static const Py_ssize_t READ_WHOLE_LINE = -1;
890
891/* If reading from a file, we need to only pull the bytes we need, since there
892 may be multiple pickle objects arranged contiguously in the same input
893 buffer.
894
895 If `n` is READ_WHOLE_LINE, read a whole line. Otherwise, read up to `n`
896 bytes from the input stream/buffer.
897
898 Update the unpickler's input buffer with the newly-read data. Returns -1 on
899 failure; on success, returns the number of bytes read from the file.
900
901 On success, self->input_len will be 0; this is intentional so that when
902 unpickling from a file, the "we've run out of data" code paths will trigger,
903 causing the Unpickler to go back to the file for more data. Use the returned
904 size to tell you how much data you can process. */
905static Py_ssize_t
906_Unpickler_ReadFromFile(UnpicklerObject *self, Py_ssize_t n)
907{
908 PyObject *data;
Antoine Pitrou04248a82010-10-12 20:51:21 +0000909 Py_ssize_t read_size, prefetched_size = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000910
911 assert(self->read != NULL);
Antoine Pitrou04248a82010-10-12 20:51:21 +0000912
913 if (_Unpickler_SkipConsumed(self) < 0)
914 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000915
916 if (n == READ_WHOLE_LINE)
917 data = PyObject_Call(self->readline, empty_tuple, NULL);
918 else {
919 PyObject *len = PyLong_FromSsize_t(n);
920 if (len == NULL)
921 return -1;
922 data = _Unpickler_FastCall(self, self->read, len);
923 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000924 if (data == NULL)
925 return -1;
926
Antoine Pitrou04248a82010-10-12 20:51:21 +0000927 /* Prefetch some data without advancing the file pointer, if possible */
928 if (self->peek) {
929 PyObject *len, *prefetched;
930 len = PyLong_FromSsize_t(PREFETCH);
931 if (len == NULL) {
932 Py_DECREF(data);
933 return -1;
934 }
935 prefetched = _Unpickler_FastCall(self, self->peek, len);
936 if (prefetched == NULL) {
937 if (PyErr_ExceptionMatches(PyExc_NotImplementedError)) {
938 /* peek() is probably not supported by the given file object */
939 PyErr_Clear();
940 Py_CLEAR(self->peek);
941 }
942 else {
943 Py_DECREF(data);
944 return -1;
945 }
946 }
947 else {
948 assert(PyBytes_Check(prefetched));
949 prefetched_size = PyBytes_GET_SIZE(prefetched);
950 PyBytes_ConcatAndDel(&data, prefetched);
951 if (data == NULL)
952 return -1;
953 }
954 }
955
956 read_size = _Unpickler_SetStringInput(self, data) - prefetched_size;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000957 Py_DECREF(data);
Antoine Pitrou04248a82010-10-12 20:51:21 +0000958 self->prefetched_idx = read_size;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000959 return read_size;
960}
961
962/* Read `n` bytes from the unpickler's data source, storing the result in `*s`.
963
964 This should be used for all data reads, rather than accessing the unpickler's
965 input buffer directly. This method deals correctly with reading from input
966 streams, which the input buffer doesn't deal with.
967
968 Note that when reading from a file-like object, self->next_read_idx won't
969 be updated (it should remain at 0 for the entire unpickling process). You
970 should use this function's return value to know how many bytes you can
971 consume.
972
973 Returns -1 (with an exception set) on failure. On success, return the
974 number of chars read. */
975static Py_ssize_t
976_Unpickler_Read(UnpicklerObject *self, char **s, Py_ssize_t n)
977{
Antoine Pitrou04248a82010-10-12 20:51:21 +0000978 Py_ssize_t num_read;
979
Antoine Pitrou04248a82010-10-12 20:51:21 +0000980 if (self->next_read_idx + n <= self->input_len) {
981 *s = self->input_buffer + self->next_read_idx;
982 self->next_read_idx += n;
983 return n;
984 }
985 if (!self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000986 PyErr_Format(PyExc_EOFError, "Ran out of input");
Amaury Forgeot d'Arc3e4e72f2008-11-11 20:05:06 +0000987 return -1;
988 }
Antoine Pitrou04248a82010-10-12 20:51:21 +0000989 num_read = _Unpickler_ReadFromFile(self, n);
990 if (num_read < 0)
991 return -1;
992 if (num_read < n) {
993 PyErr_Format(PyExc_EOFError, "Ran out of input");
994 return -1;
995 }
996 *s = self->input_buffer;
997 self->next_read_idx = n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000998 return n;
999}
1000
1001static Py_ssize_t
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001002_Unpickler_CopyLine(UnpicklerObject *self, char *line, Py_ssize_t len,
1003 char **result)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001004{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001005 char *input_line = PyMem_Realloc(self->input_line, len + 1);
1006 if (input_line == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001007 return -1;
1008
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001009 memcpy(input_line, line, len);
1010 input_line[len] = '\0';
1011 self->input_line = input_line;
1012 *result = self->input_line;
1013 return len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001014}
1015
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001016/* Read a line from the input stream/buffer. If we run off the end of the input
1017 before hitting \n, return the data we found.
1018
1019 Returns the number of chars read, or -1 on failure. */
1020static Py_ssize_t
1021_Unpickler_Readline(UnpicklerObject *self, char **result)
1022{
1023 Py_ssize_t i, num_read;
1024
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001025 for (i = self->next_read_idx; i < self->input_len; i++) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001026 if (self->input_buffer[i] == '\n') {
1027 char *line_start = self->input_buffer + self->next_read_idx;
1028 num_read = i - self->next_read_idx + 1;
1029 self->next_read_idx = i + 1;
1030 return _Unpickler_CopyLine(self, line_start, num_read, result);
1031 }
1032 }
1033 if (self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001034 num_read = _Unpickler_ReadFromFile(self, READ_WHOLE_LINE);
1035 if (num_read < 0)
1036 return -1;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001037 self->next_read_idx = num_read;
Antoine Pitrouf6c7a852011-08-11 21:04:02 +02001038 return _Unpickler_CopyLine(self, self->input_buffer, num_read, result);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001039 }
1040
1041 /* If we get here, we've run off the end of the input string. Return the
1042 remaining string and let the caller figure it out. */
1043 *result = self->input_buffer + self->next_read_idx;
1044 num_read = i - self->next_read_idx;
1045 self->next_read_idx = i;
1046 return num_read;
1047}
1048
1049/* Returns -1 (with an exception set) on failure, 0 on success. The memo array
1050 will be modified in place. */
1051static int
1052_Unpickler_ResizeMemoList(UnpicklerObject *self, Py_ssize_t new_size)
1053{
1054 Py_ssize_t i;
1055 PyObject **memo;
1056
1057 assert(new_size > self->memo_size);
1058
1059 memo = PyMem_REALLOC(self->memo, new_size * sizeof(PyObject *));
1060 if (memo == NULL) {
1061 PyErr_NoMemory();
1062 return -1;
1063 }
1064 self->memo = memo;
1065 for (i = self->memo_size; i < new_size; i++)
1066 self->memo[i] = NULL;
1067 self->memo_size = new_size;
1068 return 0;
1069}
1070
1071/* Returns NULL if idx is out of bounds. */
1072static PyObject *
1073_Unpickler_MemoGet(UnpicklerObject *self, Py_ssize_t idx)
1074{
1075 if (idx < 0 || idx >= self->memo_size)
1076 return NULL;
1077
1078 return self->memo[idx];
1079}
1080
1081/* Returns -1 (with an exception set) on failure, 0 on success.
1082 This takes its own reference to `value`. */
1083static int
1084_Unpickler_MemoPut(UnpicklerObject *self, Py_ssize_t idx, PyObject *value)
1085{
1086 PyObject *old_item;
1087
1088 if (idx >= self->memo_size) {
1089 if (_Unpickler_ResizeMemoList(self, idx * 2) < 0)
1090 return -1;
1091 assert(idx < self->memo_size);
1092 }
1093 Py_INCREF(value);
1094 old_item = self->memo[idx];
1095 self->memo[idx] = value;
1096 Py_XDECREF(old_item);
1097 return 0;
1098}
1099
1100static PyObject **
1101_Unpickler_NewMemo(Py_ssize_t new_size)
1102{
1103 PyObject **memo = PyMem_MALLOC(new_size * sizeof(PyObject *));
1104 if (memo == NULL)
1105 return NULL;
1106 memset(memo, 0, new_size * sizeof(PyObject *));
1107 return memo;
1108}
1109
1110/* Free the unpickler's memo, taking care to decref any items left in it. */
1111static void
1112_Unpickler_MemoCleanup(UnpicklerObject *self)
1113{
1114 Py_ssize_t i;
1115 PyObject **memo = self->memo;
1116
1117 if (self->memo == NULL)
1118 return;
1119 self->memo = NULL;
1120 i = self->memo_size;
1121 while (--i >= 0) {
1122 Py_XDECREF(memo[i]);
1123 }
1124 PyMem_FREE(memo);
1125}
1126
1127static UnpicklerObject *
1128_Unpickler_New(void)
1129{
1130 UnpicklerObject *self;
1131
1132 self = PyObject_GC_New(UnpicklerObject, &Unpickler_Type);
1133 if (self == NULL)
1134 return NULL;
1135
1136 self->stack = (Pdata *)Pdata_New();
1137 if (self->stack == NULL) {
1138 Py_DECREF(self);
1139 return NULL;
1140 }
1141 memset(&self->buffer, 0, sizeof(Py_buffer));
1142
1143 self->memo_size = 32;
1144 self->memo = _Unpickler_NewMemo(self->memo_size);
1145 if (self->memo == NULL) {
1146 Py_DECREF(self);
1147 return NULL;
1148 }
1149
1150 self->arg = NULL;
1151 self->pers_func = NULL;
1152 self->input_buffer = NULL;
1153 self->input_line = NULL;
1154 self->input_len = 0;
1155 self->next_read_idx = 0;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001156 self->prefetched_idx = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001157 self->read = NULL;
1158 self->readline = NULL;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001159 self->peek = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001160 self->encoding = NULL;
1161 self->errors = NULL;
1162 self->marks = NULL;
1163 self->num_marks = 0;
1164 self->marks_size = 0;
1165 self->proto = 0;
1166 self->fix_imports = 0;
1167
1168 return self;
1169}
1170
1171/* Returns -1 (with an exception set) on failure, 0 on success. This may
1172 be called once on a freshly created Pickler. */
1173static int
1174_Unpickler_SetInputStream(UnpicklerObject *self, PyObject *file)
1175{
Antoine Pitrou04248a82010-10-12 20:51:21 +00001176 self->peek = PyObject_GetAttrString(file, "peek");
1177 if (self->peek == NULL) {
1178 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1179 PyErr_Clear();
1180 else
1181 return -1;
1182 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001183 self->read = PyObject_GetAttrString(file, "read");
1184 self->readline = PyObject_GetAttrString(file, "readline");
1185 if (self->readline == NULL || self->read == NULL) {
1186 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1187 PyErr_SetString(PyExc_TypeError,
1188 "file must have 'read' and 'readline' attributes");
1189 Py_CLEAR(self->read);
1190 Py_CLEAR(self->readline);
Antoine Pitrou04248a82010-10-12 20:51:21 +00001191 Py_CLEAR(self->peek);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001192 return -1;
1193 }
1194 return 0;
1195}
1196
1197/* Returns -1 (with an exception set) on failure, 0 on success. This may
1198 be called once on a freshly created Pickler. */
1199static int
1200_Unpickler_SetInputEncoding(UnpicklerObject *self,
1201 const char *encoding,
1202 const char *errors)
1203{
1204 if (encoding == NULL)
1205 encoding = "ASCII";
1206 if (errors == NULL)
1207 errors = "strict";
1208
1209 self->encoding = strdup(encoding);
1210 self->errors = strdup(errors);
1211 if (self->encoding == NULL || self->errors == NULL) {
1212 PyErr_NoMemory();
1213 return -1;
1214 }
1215 return 0;
1216}
1217
1218/* Generate a GET opcode for an object stored in the memo. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001219static int
1220memo_get(PicklerObject *self, PyObject *key)
1221{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001222 long *value;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001223 char pdata[30];
1224 int len;
1225
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001226 value = PyMemoTable_Get(self->memo, key);
1227 if (value == NULL) {
1228 PyErr_SetObject(PyExc_KeyError, key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001229 return -1;
1230 }
1231
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001232 if (!self->bin) {
1233 pdata[0] = GET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001234 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1, "%ld\n", *value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001235 len = (int)strlen(pdata);
1236 }
1237 else {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001238 if (*value < 256) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001239 pdata[0] = BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001240 pdata[1] = (unsigned char)(*value & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001241 len = 2;
1242 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001243 else if (*value <= 0xffffffffL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001244 pdata[0] = LONG_BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001245 pdata[1] = (unsigned char)(*value & 0xff);
1246 pdata[2] = (unsigned char)((*value >> 8) & 0xff);
1247 pdata[3] = (unsigned char)((*value >> 16) & 0xff);
1248 pdata[4] = (unsigned char)((*value >> 24) & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001249 len = 5;
1250 }
1251 else { /* unlikely */
1252 PyErr_SetString(PicklingError,
1253 "memo id too large for LONG_BINGET");
1254 return -1;
1255 }
1256 }
1257
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001258 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001259 return -1;
1260
1261 return 0;
1262}
1263
1264/* Store an object in the memo, assign it a new unique ID based on the number
1265 of objects currently stored in the memo and generate a PUT opcode. */
1266static int
1267memo_put(PicklerObject *self, PyObject *obj)
1268{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001269 long x;
1270 char pdata[30];
1271 int len;
1272 int status = 0;
1273
1274 if (self->fast)
1275 return 0;
1276
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001277 x = PyMemoTable_Size(self->memo);
1278 if (PyMemoTable_Set(self->memo, obj, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001279 goto error;
1280
1281 if (!self->bin) {
1282 pdata[0] = PUT;
1283 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1, "%ld\n", x);
1284 len = strlen(pdata);
1285 }
1286 else {
1287 if (x < 256) {
1288 pdata[0] = BINPUT;
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00001289 pdata[1] = (unsigned char)x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001290 len = 2;
1291 }
1292 else if (x <= 0xffffffffL) {
1293 pdata[0] = LONG_BINPUT;
1294 pdata[1] = (unsigned char)(x & 0xff);
1295 pdata[2] = (unsigned char)((x >> 8) & 0xff);
1296 pdata[3] = (unsigned char)((x >> 16) & 0xff);
1297 pdata[4] = (unsigned char)((x >> 24) & 0xff);
1298 len = 5;
1299 }
1300 else { /* unlikely */
1301 PyErr_SetString(PicklingError,
1302 "memo id too large for LONG_BINPUT");
1303 return -1;
1304 }
1305 }
1306
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001307 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001308 goto error;
1309
1310 if (0) {
1311 error:
1312 status = -1;
1313 }
1314
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001315 return status;
1316}
1317
1318static PyObject *
1319whichmodule(PyObject *global, PyObject *global_name)
1320{
1321 Py_ssize_t i, j;
1322 static PyObject *module_str = NULL;
1323 static PyObject *main_str = NULL;
1324 PyObject *module_name;
1325 PyObject *modules_dict;
1326 PyObject *module;
1327 PyObject *obj;
1328
1329 if (module_str == NULL) {
1330 module_str = PyUnicode_InternFromString("__module__");
1331 if (module_str == NULL)
1332 return NULL;
1333 main_str = PyUnicode_InternFromString("__main__");
1334 if (main_str == NULL)
1335 return NULL;
1336 }
1337
1338 module_name = PyObject_GetAttr(global, module_str);
1339
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00001340 /* In some rare cases (e.g., bound methods of extension types),
1341 __module__ can be None. If it is so, then search sys.modules
1342 for the module of global. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001343 if (module_name == Py_None) {
1344 Py_DECREF(module_name);
1345 goto search;
1346 }
1347
1348 if (module_name) {
1349 return module_name;
1350 }
1351 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1352 PyErr_Clear();
1353 else
1354 return NULL;
1355
1356 search:
1357 modules_dict = PySys_GetObject("modules");
1358 if (modules_dict == NULL)
1359 return NULL;
1360
1361 i = 0;
1362 module_name = NULL;
1363 while ((j = PyDict_Next(modules_dict, &i, &module_name, &module))) {
Mark Dickinson211c6252009-02-01 10:28:51 +00001364 if (PyObject_RichCompareBool(module_name, main_str, Py_EQ) == 1)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001365 continue;
1366
1367 obj = PyObject_GetAttr(module, global_name);
1368 if (obj == NULL) {
1369 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1370 PyErr_Clear();
1371 else
1372 return NULL;
1373 continue;
1374 }
1375
1376 if (obj != global) {
1377 Py_DECREF(obj);
1378 continue;
1379 }
1380
1381 Py_DECREF(obj);
1382 break;
1383 }
1384
1385 /* If no module is found, use __main__. */
1386 if (!j) {
1387 module_name = main_str;
1388 }
1389
1390 Py_INCREF(module_name);
1391 return module_name;
1392}
1393
1394/* fast_save_enter() and fast_save_leave() are guards against recursive
1395 objects when Pickler is used with the "fast mode" (i.e., with object
1396 memoization disabled). If the nesting of a list or dict object exceed
1397 FAST_NESTING_LIMIT, these guards will start keeping an internal
1398 reference to the seen list or dict objects and check whether these objects
1399 are recursive. These are not strictly necessary, since save() has a
1400 hard-coded recursion limit, but they give a nicer error message than the
1401 typical RuntimeError. */
1402static int
1403fast_save_enter(PicklerObject *self, PyObject *obj)
1404{
1405 /* if fast_nesting < 0, we're doing an error exit. */
1406 if (++self->fast_nesting >= FAST_NESTING_LIMIT) {
1407 PyObject *key = NULL;
1408 if (self->fast_memo == NULL) {
1409 self->fast_memo = PyDict_New();
1410 if (self->fast_memo == NULL) {
1411 self->fast_nesting = -1;
1412 return 0;
1413 }
1414 }
1415 key = PyLong_FromVoidPtr(obj);
1416 if (key == NULL)
1417 return 0;
1418 if (PyDict_GetItem(self->fast_memo, key)) {
1419 Py_DECREF(key);
1420 PyErr_Format(PyExc_ValueError,
1421 "fast mode: can't pickle cyclic objects "
1422 "including object type %.200s at %p",
1423 obj->ob_type->tp_name, obj);
1424 self->fast_nesting = -1;
1425 return 0;
1426 }
1427 if (PyDict_SetItem(self->fast_memo, key, Py_None) < 0) {
1428 Py_DECREF(key);
1429 self->fast_nesting = -1;
1430 return 0;
1431 }
1432 Py_DECREF(key);
1433 }
1434 return 1;
1435}
1436
1437static int
1438fast_save_leave(PicklerObject *self, PyObject *obj)
1439{
1440 if (self->fast_nesting-- >= FAST_NESTING_LIMIT) {
1441 PyObject *key = PyLong_FromVoidPtr(obj);
1442 if (key == NULL)
1443 return 0;
1444 if (PyDict_DelItem(self->fast_memo, key) < 0) {
1445 Py_DECREF(key);
1446 return 0;
1447 }
1448 Py_DECREF(key);
1449 }
1450 return 1;
1451}
1452
1453static int
1454save_none(PicklerObject *self, PyObject *obj)
1455{
1456 const char none_op = NONE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001457 if (_Pickler_Write(self, &none_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001458 return -1;
1459
1460 return 0;
1461}
1462
1463static int
1464save_bool(PicklerObject *self, PyObject *obj)
1465{
1466 static const char *buf[2] = { FALSE, TRUE };
1467 const char len[2] = {sizeof(FALSE) - 1, sizeof(TRUE) - 1};
1468 int p = (obj == Py_True);
1469
1470 if (self->proto >= 2) {
1471 const char bool_op = p ? NEWTRUE : NEWFALSE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001472 if (_Pickler_Write(self, &bool_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001473 return -1;
1474 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001475 else if (_Pickler_Write(self, buf[p], len[p]) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001476 return -1;
1477
1478 return 0;
1479}
1480
1481static int
1482save_int(PicklerObject *self, long x)
1483{
1484 char pdata[32];
1485 int len = 0;
1486
1487 if (!self->bin
1488#if SIZEOF_LONG > 4
1489 || x > 0x7fffffffL || x < -0x80000000L
1490#endif
1491 ) {
1492 /* Text-mode pickle, or long too big to fit in the 4-byte
1493 * signed BININT format: store as a string.
1494 */
Mark Dickinson8dd05142009-01-20 20:43:58 +00001495 pdata[0] = LONG; /* use LONG for consistency with pickle.py */
1496 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1, "%ldL\n", x);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001497 if (_Pickler_Write(self, pdata, strlen(pdata)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001498 return -1;
1499 }
1500 else {
1501 /* Binary pickle and x fits in a signed 4-byte int. */
1502 pdata[1] = (unsigned char)(x & 0xff);
1503 pdata[2] = (unsigned char)((x >> 8) & 0xff);
1504 pdata[3] = (unsigned char)((x >> 16) & 0xff);
1505 pdata[4] = (unsigned char)((x >> 24) & 0xff);
1506
1507 if ((pdata[4] == 0) && (pdata[3] == 0)) {
1508 if (pdata[2] == 0) {
1509 pdata[0] = BININT1;
1510 len = 2;
1511 }
1512 else {
1513 pdata[0] = BININT2;
1514 len = 3;
1515 }
1516 }
1517 else {
1518 pdata[0] = BININT;
1519 len = 5;
1520 }
1521
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001522 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001523 return -1;
1524 }
1525
1526 return 0;
1527}
1528
1529static int
1530save_long(PicklerObject *self, PyObject *obj)
1531{
1532 PyObject *repr = NULL;
1533 Py_ssize_t size;
1534 long val = PyLong_AsLong(obj);
1535 int status = 0;
1536
1537 const char long_op = LONG;
1538
1539 if (val == -1 && PyErr_Occurred()) {
1540 /* out of range for int pickling */
1541 PyErr_Clear();
1542 }
1543 else
1544 return save_int(self, val);
1545
1546 if (self->proto >= 2) {
1547 /* Linear-time pickling. */
1548 size_t nbits;
1549 size_t nbytes;
1550 unsigned char *pdata;
1551 char header[5];
1552 int i;
1553 int sign = _PyLong_Sign(obj);
1554
1555 if (sign == 0) {
1556 header[0] = LONG1;
1557 header[1] = 0; /* It's 0 -- an empty bytestring. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001558 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001559 goto error;
1560 return 0;
1561 }
1562 nbits = _PyLong_NumBits(obj);
1563 if (nbits == (size_t)-1 && PyErr_Occurred())
1564 goto error;
1565 /* How many bytes do we need? There are nbits >> 3 full
1566 * bytes of data, and nbits & 7 leftover bits. If there
1567 * are any leftover bits, then we clearly need another
1568 * byte. Wnat's not so obvious is that we *probably*
1569 * need another byte even if there aren't any leftovers:
1570 * the most-significant bit of the most-significant byte
1571 * acts like a sign bit, and it's usually got a sense
1572 * opposite of the one we need. The exception is longs
1573 * of the form -(2**(8*j-1)) for j > 0. Such a long is
1574 * its own 256's-complement, so has the right sign bit
1575 * even without the extra byte. That's a pain to check
1576 * for in advance, though, so we always grab an extra
1577 * byte at the start, and cut it back later if possible.
1578 */
1579 nbytes = (nbits >> 3) + 1;
1580 if (nbytes > INT_MAX) {
1581 PyErr_SetString(PyExc_OverflowError,
1582 "long too large to pickle");
1583 goto error;
1584 }
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001585 repr = PyBytes_FromStringAndSize(NULL, (Py_ssize_t)nbytes);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001586 if (repr == NULL)
1587 goto error;
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001588 pdata = (unsigned char *)PyBytes_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001589 i = _PyLong_AsByteArray((PyLongObject *)obj,
1590 pdata, nbytes,
1591 1 /* little endian */ , 1 /* signed */ );
1592 if (i < 0)
1593 goto error;
1594 /* If the long is negative, this may be a byte more than
1595 * needed. This is so iff the MSB is all redundant sign
1596 * bits.
1597 */
1598 if (sign < 0 &&
1599 nbytes > 1 &&
1600 pdata[nbytes - 1] == 0xff &&
1601 (pdata[nbytes - 2] & 0x80) != 0) {
1602 nbytes--;
1603 }
1604
1605 if (nbytes < 256) {
1606 header[0] = LONG1;
1607 header[1] = (unsigned char)nbytes;
1608 size = 2;
1609 }
1610 else {
1611 header[0] = LONG4;
1612 size = (int)nbytes;
1613 for (i = 1; i < 5; i++) {
1614 header[i] = (unsigned char)(size & 0xff);
1615 size >>= 8;
1616 }
1617 size = 5;
1618 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001619 if (_Pickler_Write(self, header, size) < 0 ||
1620 _Pickler_Write(self, (char *)pdata, (int)nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001621 goto error;
1622 }
1623 else {
1624 char *string;
1625
Mark Dickinson8dd05142009-01-20 20:43:58 +00001626 /* proto < 2: write the repr and newline. This is quadratic-time (in
1627 the number of digits), in both directions. We add a trailing 'L'
1628 to the repr, for compatibility with Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001629
1630 repr = PyObject_Repr(obj);
1631 if (repr == NULL)
1632 goto error;
1633
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00001634 string = _PyUnicode_AsStringAndSize(repr, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001635 if (string == NULL)
1636 goto error;
1637
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001638 if (_Pickler_Write(self, &long_op, 1) < 0 ||
1639 _Pickler_Write(self, string, size) < 0 ||
1640 _Pickler_Write(self, "L\n", 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001641 goto error;
1642 }
1643
1644 if (0) {
1645 error:
1646 status = -1;
1647 }
1648 Py_XDECREF(repr);
1649
1650 return status;
1651}
1652
1653static int
1654save_float(PicklerObject *self, PyObject *obj)
1655{
1656 double x = PyFloat_AS_DOUBLE((PyFloatObject *)obj);
1657
1658 if (self->bin) {
1659 char pdata[9];
1660 pdata[0] = BINFLOAT;
1661 if (_PyFloat_Pack8(x, (unsigned char *)&pdata[1], 0) < 0)
1662 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001663 if (_Pickler_Write(self, pdata, 9) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001664 return -1;
Eric Smith0923d1d2009-04-16 20:16:10 +00001665 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001666 else {
Eric Smith0923d1d2009-04-16 20:16:10 +00001667 int result = -1;
1668 char *buf = NULL;
1669 char op = FLOAT;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001670
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001671 if (_Pickler_Write(self, &op, 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001672 goto done;
1673
Mark Dickinson3e09f432009-04-17 08:41:23 +00001674 buf = PyOS_double_to_string(x, 'g', 17, 0, NULL);
Eric Smith0923d1d2009-04-16 20:16:10 +00001675 if (!buf) {
1676 PyErr_NoMemory();
1677 goto done;
1678 }
1679
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001680 if (_Pickler_Write(self, buf, strlen(buf)) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001681 goto done;
1682
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001683 if (_Pickler_Write(self, "\n", 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001684 goto done;
1685
1686 result = 0;
1687done:
1688 PyMem_Free(buf);
1689 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001690 }
1691
1692 return 0;
1693}
1694
1695static int
1696save_bytes(PicklerObject *self, PyObject *obj)
1697{
1698 if (self->proto < 3) {
1699 /* Older pickle protocols do not have an opcode for pickling bytes
1700 objects. Therefore, we need to fake the copy protocol (i.e.,
1701 the __reduce__ method) to permit bytes object unpickling. */
1702 PyObject *reduce_value = NULL;
1703 PyObject *bytelist = NULL;
1704 int status;
1705
1706 bytelist = PySequence_List(obj);
1707 if (bytelist == NULL)
1708 return -1;
1709
1710 reduce_value = Py_BuildValue("(O(O))", (PyObject *)&PyBytes_Type,
1711 bytelist);
1712 if (reduce_value == NULL) {
1713 Py_DECREF(bytelist);
1714 return -1;
1715 }
1716
1717 /* save_reduce() will memoize the object automatically. */
1718 status = save_reduce(self, reduce_value, obj);
1719 Py_DECREF(reduce_value);
1720 Py_DECREF(bytelist);
1721 return status;
1722 }
1723 else {
1724 Py_ssize_t size;
1725 char header[5];
1726 int len;
1727
1728 size = PyBytes_Size(obj);
1729 if (size < 0)
1730 return -1;
1731
1732 if (size < 256) {
1733 header[0] = SHORT_BINBYTES;
1734 header[1] = (unsigned char)size;
1735 len = 2;
1736 }
1737 else if (size <= 0xffffffffL) {
1738 header[0] = BINBYTES;
1739 header[1] = (unsigned char)(size & 0xff);
1740 header[2] = (unsigned char)((size >> 8) & 0xff);
1741 header[3] = (unsigned char)((size >> 16) & 0xff);
1742 header[4] = (unsigned char)((size >> 24) & 0xff);
1743 len = 5;
1744 }
1745 else {
1746 return -1; /* string too large */
1747 }
1748
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001749 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001750 return -1;
1751
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001752 if (_Pickler_Write(self, PyBytes_AS_STRING(obj), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001753 return -1;
1754
1755 if (memo_put(self, obj) < 0)
1756 return -1;
1757
1758 return 0;
1759 }
1760}
1761
1762/* A copy of PyUnicode_EncodeRawUnicodeEscape() that also translates
1763 backslash and newline characters to \uXXXX escapes. */
1764static PyObject *
1765raw_unicode_escape(const Py_UNICODE *s, Py_ssize_t size)
1766{
1767 PyObject *repr, *result;
1768 char *p;
1769 char *q;
1770
1771 static const char *hexdigits = "0123456789abcdef";
1772
1773#ifdef Py_UNICODE_WIDE
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001774 const Py_ssize_t expandsize = 10;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001775#else
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001776 const Py_ssize_t expandsize = 6;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001777#endif
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001778
1779 if (size > PY_SSIZE_T_MAX / expandsize)
1780 return PyErr_NoMemory();
1781
1782 repr = PyByteArray_FromStringAndSize(NULL, expandsize * size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001783 if (repr == NULL)
1784 return NULL;
1785 if (size == 0)
1786 goto done;
1787
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001788 p = q = PyByteArray_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001789 while (size-- > 0) {
1790 Py_UNICODE ch = *s++;
1791#ifdef Py_UNICODE_WIDE
1792 /* Map 32-bit characters to '\Uxxxxxxxx' */
1793 if (ch >= 0x10000) {
1794 *p++ = '\\';
1795 *p++ = 'U';
1796 *p++ = hexdigits[(ch >> 28) & 0xf];
1797 *p++ = hexdigits[(ch >> 24) & 0xf];
1798 *p++ = hexdigits[(ch >> 20) & 0xf];
1799 *p++ = hexdigits[(ch >> 16) & 0xf];
1800 *p++ = hexdigits[(ch >> 12) & 0xf];
1801 *p++ = hexdigits[(ch >> 8) & 0xf];
1802 *p++ = hexdigits[(ch >> 4) & 0xf];
1803 *p++ = hexdigits[ch & 15];
1804 }
1805 else
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001806#else
1807 /* Map UTF-16 surrogate pairs to '\U00xxxxxx' */
1808 if (ch >= 0xD800 && ch < 0xDC00) {
1809 Py_UNICODE ch2;
1810 Py_UCS4 ucs;
1811
1812 ch2 = *s++;
1813 size--;
1814 if (ch2 >= 0xDC00 && ch2 <= 0xDFFF) {
1815 ucs = (((ch & 0x03FF) << 10) | (ch2 & 0x03FF)) + 0x00010000;
1816 *p++ = '\\';
1817 *p++ = 'U';
1818 *p++ = hexdigits[(ucs >> 28) & 0xf];
1819 *p++ = hexdigits[(ucs >> 24) & 0xf];
1820 *p++ = hexdigits[(ucs >> 20) & 0xf];
1821 *p++ = hexdigits[(ucs >> 16) & 0xf];
1822 *p++ = hexdigits[(ucs >> 12) & 0xf];
1823 *p++ = hexdigits[(ucs >> 8) & 0xf];
1824 *p++ = hexdigits[(ucs >> 4) & 0xf];
1825 *p++ = hexdigits[ucs & 0xf];
1826 continue;
1827 }
1828 /* Fall through: isolated surrogates are copied as-is */
1829 s--;
1830 size++;
1831 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001832#endif
1833 /* Map 16-bit characters to '\uxxxx' */
1834 if (ch >= 256 || ch == '\\' || ch == '\n') {
1835 *p++ = '\\';
1836 *p++ = 'u';
1837 *p++ = hexdigits[(ch >> 12) & 0xf];
1838 *p++ = hexdigits[(ch >> 8) & 0xf];
1839 *p++ = hexdigits[(ch >> 4) & 0xf];
1840 *p++ = hexdigits[ch & 15];
1841 }
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001842 /* Copy everything else as-is */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001843 else
1844 *p++ = (char) ch;
1845 }
1846 size = p - q;
1847
1848 done:
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001849 result = PyBytes_FromStringAndSize(PyByteArray_AS_STRING(repr), size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001850 Py_DECREF(repr);
1851 return result;
1852}
1853
1854static int
1855save_unicode(PicklerObject *self, PyObject *obj)
1856{
1857 Py_ssize_t size;
1858 PyObject *encoded = NULL;
1859
1860 if (self->bin) {
1861 char pdata[5];
1862
Victor Stinner485fb562010-04-13 11:07:24 +00001863 encoded = PyUnicode_EncodeUTF8(PyUnicode_AS_UNICODE(obj),
1864 PyUnicode_GET_SIZE(obj),
1865 "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001866 if (encoded == NULL)
1867 goto error;
1868
1869 size = PyBytes_GET_SIZE(encoded);
1870 if (size < 0 || size > 0xffffffffL)
1871 goto error; /* string too large */
1872
1873 pdata[0] = BINUNICODE;
1874 pdata[1] = (unsigned char)(size & 0xff);
1875 pdata[2] = (unsigned char)((size >> 8) & 0xff);
1876 pdata[3] = (unsigned char)((size >> 16) & 0xff);
1877 pdata[4] = (unsigned char)((size >> 24) & 0xff);
1878
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001879 if (_Pickler_Write(self, pdata, 5) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001880 goto error;
1881
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001882 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001883 goto error;
1884 }
1885 else {
1886 const char unicode_op = UNICODE;
1887
1888 encoded = raw_unicode_escape(PyUnicode_AS_UNICODE(obj),
1889 PyUnicode_GET_SIZE(obj));
1890 if (encoded == NULL)
1891 goto error;
1892
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001893 if (_Pickler_Write(self, &unicode_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001894 goto error;
1895
1896 size = PyBytes_GET_SIZE(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001897 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001898 goto error;
1899
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001900 if (_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001901 goto error;
1902 }
1903 if (memo_put(self, obj) < 0)
1904 goto error;
1905
1906 Py_DECREF(encoded);
1907 return 0;
1908
1909 error:
1910 Py_XDECREF(encoded);
1911 return -1;
1912}
1913
1914/* A helper for save_tuple. Push the len elements in tuple t on the stack. */
1915static int
1916store_tuple_elements(PicklerObject *self, PyObject *t, int len)
1917{
1918 int i;
1919
1920 assert(PyTuple_Size(t) == len);
1921
1922 for (i = 0; i < len; i++) {
1923 PyObject *element = PyTuple_GET_ITEM(t, i);
1924
1925 if (element == NULL)
1926 return -1;
1927 if (save(self, element, 0) < 0)
1928 return -1;
1929 }
1930
1931 return 0;
1932}
1933
1934/* Tuples are ubiquitous in the pickle protocols, so many techniques are
1935 * used across protocols to minimize the space needed to pickle them.
1936 * Tuples are also the only builtin immutable type that can be recursive
1937 * (a tuple can be reached from itself), and that requires some subtle
1938 * magic so that it works in all cases. IOW, this is a long routine.
1939 */
1940static int
1941save_tuple(PicklerObject *self, PyObject *obj)
1942{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001943 int len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001944
1945 const char mark_op = MARK;
1946 const char tuple_op = TUPLE;
1947 const char pop_op = POP;
1948 const char pop_mark_op = POP_MARK;
1949 const char len2opcode[] = {EMPTY_TUPLE, TUPLE1, TUPLE2, TUPLE3};
1950
1951 if ((len = PyTuple_Size(obj)) < 0)
1952 return -1;
1953
1954 if (len == 0) {
1955 char pdata[2];
1956
1957 if (self->proto) {
1958 pdata[0] = EMPTY_TUPLE;
1959 len = 1;
1960 }
1961 else {
1962 pdata[0] = MARK;
1963 pdata[1] = TUPLE;
1964 len = 2;
1965 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001966 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001967 return -1;
1968 return 0;
1969 }
1970
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001971 /* The tuple isn't in the memo now. If it shows up there after
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001972 * saving the tuple elements, the tuple must be recursive, in
1973 * which case we'll pop everything we put on the stack, and fetch
1974 * its value from the memo.
1975 */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001976 if (len <= 3 && self->proto >= 2) {
1977 /* Use TUPLE{1,2,3} opcodes. */
1978 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001979 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001980
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001981 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001982 /* pop the len elements */
1983 for (i = 0; i < len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001984 if (_Pickler_Write(self, &pop_op, 1) < 0)
1985 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001986 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001987 if (memo_get(self, obj) < 0)
1988 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001989
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001990 return 0;
1991 }
1992 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001993 if (_Pickler_Write(self, len2opcode + len, 1) < 0)
1994 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001995 }
1996 goto memoize;
1997 }
1998
1999 /* proto < 2 and len > 0, or proto >= 2 and len > 3.
2000 * Generate MARK e1 e2 ... TUPLE
2001 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002002 if (_Pickler_Write(self, &mark_op, 1) < 0)
2003 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002004
2005 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002006 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002007
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002008 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002009 /* pop the stack stuff we pushed */
2010 if (self->bin) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002011 if (_Pickler_Write(self, &pop_mark_op, 1) < 0)
2012 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002013 }
2014 else {
2015 /* Note that we pop one more than len, to remove
2016 * the MARK too.
2017 */
2018 for (i = 0; i <= len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002019 if (_Pickler_Write(self, &pop_op, 1) < 0)
2020 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002021 }
2022 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002023 if (memo_get(self, obj) < 0)
2024 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002025
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002026 return 0;
2027 }
2028 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002029 if (_Pickler_Write(self, &tuple_op, 1) < 0)
2030 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002031 }
2032
2033 memoize:
2034 if (memo_put(self, obj) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002035 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002036
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002037 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002038}
2039
2040/* iter is an iterator giving items, and we batch up chunks of
2041 * MARK item item ... item APPENDS
2042 * opcode sequences. Calling code should have arranged to first create an
2043 * empty list, or list-like object, for the APPENDS to operate on.
2044 * Returns 0 on success, <0 on error.
2045 */
2046static int
2047batch_list(PicklerObject *self, PyObject *iter)
2048{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002049 PyObject *obj = NULL;
2050 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002051 int i, n;
2052
2053 const char mark_op = MARK;
2054 const char append_op = APPEND;
2055 const char appends_op = APPENDS;
2056
2057 assert(iter != NULL);
2058
2059 /* XXX: I think this function could be made faster by avoiding the
2060 iterator interface and fetching objects directly from list using
2061 PyList_GET_ITEM.
2062 */
2063
2064 if (self->proto == 0) {
2065 /* APPENDS isn't available; do one at a time. */
2066 for (;;) {
2067 obj = PyIter_Next(iter);
2068 if (obj == NULL) {
2069 if (PyErr_Occurred())
2070 return -1;
2071 break;
2072 }
2073 i = save(self, obj, 0);
2074 Py_DECREF(obj);
2075 if (i < 0)
2076 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002077 if (_Pickler_Write(self, &append_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002078 return -1;
2079 }
2080 return 0;
2081 }
2082
2083 /* proto > 0: write in batches of BATCHSIZE. */
2084 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002085 /* Get first item */
2086 firstitem = PyIter_Next(iter);
2087 if (firstitem == NULL) {
2088 if (PyErr_Occurred())
2089 goto error;
2090
2091 /* nothing more to add */
2092 break;
2093 }
2094
2095 /* Try to get a second item */
2096 obj = PyIter_Next(iter);
2097 if (obj == NULL) {
2098 if (PyErr_Occurred())
2099 goto error;
2100
2101 /* Only one item to write */
2102 if (save(self, firstitem, 0) < 0)
2103 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002104 if (_Pickler_Write(self, &append_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002105 goto error;
2106 Py_CLEAR(firstitem);
2107 break;
2108 }
2109
2110 /* More than one item to write */
2111
2112 /* Pump out MARK, items, APPENDS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002113 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002114 goto error;
2115
2116 if (save(self, firstitem, 0) < 0)
2117 goto error;
2118 Py_CLEAR(firstitem);
2119 n = 1;
2120
2121 /* Fetch and save up to BATCHSIZE items */
2122 while (obj) {
2123 if (save(self, obj, 0) < 0)
2124 goto error;
2125 Py_CLEAR(obj);
2126 n += 1;
2127
2128 if (n == BATCHSIZE)
2129 break;
2130
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002131 obj = PyIter_Next(iter);
2132 if (obj == NULL) {
2133 if (PyErr_Occurred())
2134 goto error;
2135 break;
2136 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002137 }
2138
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002139 if (_Pickler_Write(self, &appends_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002140 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002141
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002142 } while (n == BATCHSIZE);
2143 return 0;
2144
2145 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002146 Py_XDECREF(firstitem);
2147 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002148 return -1;
2149}
2150
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002151/* This is a variant of batch_list() above, specialized for lists (with no
2152 * support for list subclasses). Like batch_list(), we batch up chunks of
2153 * MARK item item ... item APPENDS
2154 * opcode sequences. Calling code should have arranged to first create an
2155 * empty list, or list-like object, for the APPENDS to operate on.
2156 * Returns 0 on success, -1 on error.
2157 *
2158 * This version is considerably faster than batch_list(), if less general.
2159 *
2160 * Note that this only works for protocols > 0.
2161 */
2162static int
2163batch_list_exact(PicklerObject *self, PyObject *obj)
2164{
2165 PyObject *item = NULL;
2166 int this_batch, total;
2167
2168 const char append_op = APPEND;
2169 const char appends_op = APPENDS;
2170 const char mark_op = MARK;
2171
2172 assert(obj != NULL);
2173 assert(self->proto > 0);
2174 assert(PyList_CheckExact(obj));
2175
2176 if (PyList_GET_SIZE(obj) == 1) {
2177 item = PyList_GET_ITEM(obj, 0);
2178 if (save(self, item, 0) < 0)
2179 return -1;
2180 if (_Pickler_Write(self, &append_op, 1) < 0)
2181 return -1;
2182 return 0;
2183 }
2184
2185 /* Write in batches of BATCHSIZE. */
2186 total = 0;
2187 do {
2188 this_batch = 0;
2189 if (_Pickler_Write(self, &mark_op, 1) < 0)
2190 return -1;
2191 while (total < PyList_GET_SIZE(obj)) {
2192 item = PyList_GET_ITEM(obj, total);
2193 if (save(self, item, 0) < 0)
2194 return -1;
2195 total++;
2196 if (++this_batch == BATCHSIZE)
2197 break;
2198 }
2199 if (_Pickler_Write(self, &appends_op, 1) < 0)
2200 return -1;
2201
2202 } while (total < PyList_GET_SIZE(obj));
2203
2204 return 0;
2205}
2206
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002207static int
2208save_list(PicklerObject *self, PyObject *obj)
2209{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002210 char header[3];
2211 int len;
2212 int status = 0;
2213
2214 if (self->fast && !fast_save_enter(self, obj))
2215 goto error;
2216
2217 /* Create an empty list. */
2218 if (self->bin) {
2219 header[0] = EMPTY_LIST;
2220 len = 1;
2221 }
2222 else {
2223 header[0] = MARK;
2224 header[1] = LIST;
2225 len = 2;
2226 }
2227
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002228 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002229 goto error;
2230
2231 /* Get list length, and bow out early if empty. */
2232 if ((len = PyList_Size(obj)) < 0)
2233 goto error;
2234
2235 if (memo_put(self, obj) < 0)
2236 goto error;
2237
2238 if (len != 0) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002239 /* Materialize the list elements. */
2240 if (PyList_CheckExact(obj) && self->proto > 0) {
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002241 if (Py_EnterRecursiveCall(" while pickling an object"))
2242 goto error;
2243 status = batch_list_exact(self, obj);
2244 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002245 } else {
2246 PyObject *iter = PyObject_GetIter(obj);
2247 if (iter == NULL)
2248 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002249
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002250 if (Py_EnterRecursiveCall(" while pickling an object")) {
2251 Py_DECREF(iter);
2252 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002253 }
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002254 status = batch_list(self, iter);
2255 Py_LeaveRecursiveCall();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002256 Py_DECREF(iter);
2257 }
2258 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002259 if (0) {
2260 error:
2261 status = -1;
2262 }
2263
2264 if (self->fast && !fast_save_leave(self, obj))
2265 status = -1;
2266
2267 return status;
2268}
2269
2270/* iter is an iterator giving (key, value) pairs, and we batch up chunks of
2271 * MARK key value ... key value SETITEMS
2272 * opcode sequences. Calling code should have arranged to first create an
2273 * empty dict, or dict-like object, for the SETITEMS to operate on.
2274 * Returns 0 on success, <0 on error.
2275 *
2276 * This is very much like batch_list(). The difference between saving
2277 * elements directly, and picking apart two-tuples, is so long-winded at
2278 * the C level, though, that attempts to combine these routines were too
2279 * ugly to bear.
2280 */
2281static int
2282batch_dict(PicklerObject *self, PyObject *iter)
2283{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002284 PyObject *obj = NULL;
2285 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002286 int i, n;
2287
2288 const char mark_op = MARK;
2289 const char setitem_op = SETITEM;
2290 const char setitems_op = SETITEMS;
2291
2292 assert(iter != NULL);
2293
2294 if (self->proto == 0) {
2295 /* SETITEMS isn't available; do one at a time. */
2296 for (;;) {
2297 obj = PyIter_Next(iter);
2298 if (obj == NULL) {
2299 if (PyErr_Occurred())
2300 return -1;
2301 break;
2302 }
2303 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2304 PyErr_SetString(PyExc_TypeError, "dict items "
2305 "iterator must return 2-tuples");
2306 return -1;
2307 }
2308 i = save(self, PyTuple_GET_ITEM(obj, 0), 0);
2309 if (i >= 0)
2310 i = save(self, PyTuple_GET_ITEM(obj, 1), 0);
2311 Py_DECREF(obj);
2312 if (i < 0)
2313 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002314 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002315 return -1;
2316 }
2317 return 0;
2318 }
2319
2320 /* proto > 0: write in batches of BATCHSIZE. */
2321 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002322 /* Get first item */
2323 firstitem = PyIter_Next(iter);
2324 if (firstitem == NULL) {
2325 if (PyErr_Occurred())
2326 goto error;
2327
2328 /* nothing more to add */
2329 break;
2330 }
2331 if (!PyTuple_Check(firstitem) || PyTuple_Size(firstitem) != 2) {
2332 PyErr_SetString(PyExc_TypeError, "dict items "
2333 "iterator must return 2-tuples");
2334 goto error;
2335 }
2336
2337 /* Try to get a second item */
2338 obj = PyIter_Next(iter);
2339 if (obj == NULL) {
2340 if (PyErr_Occurred())
2341 goto error;
2342
2343 /* Only one item to write */
2344 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2345 goto error;
2346 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2347 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002348 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002349 goto error;
2350 Py_CLEAR(firstitem);
2351 break;
2352 }
2353
2354 /* More than one item to write */
2355
2356 /* Pump out MARK, items, SETITEMS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002357 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002358 goto error;
2359
2360 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2361 goto error;
2362 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2363 goto error;
2364 Py_CLEAR(firstitem);
2365 n = 1;
2366
2367 /* Fetch and save up to BATCHSIZE items */
2368 while (obj) {
2369 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2370 PyErr_SetString(PyExc_TypeError, "dict items "
2371 "iterator must return 2-tuples");
2372 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002373 }
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002374 if (save(self, PyTuple_GET_ITEM(obj, 0), 0) < 0 ||
2375 save(self, PyTuple_GET_ITEM(obj, 1), 0) < 0)
2376 goto error;
2377 Py_CLEAR(obj);
2378 n += 1;
2379
2380 if (n == BATCHSIZE)
2381 break;
2382
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002383 obj = PyIter_Next(iter);
2384 if (obj == NULL) {
2385 if (PyErr_Occurred())
2386 goto error;
2387 break;
2388 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002389 }
2390
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002391 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002392 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002393
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002394 } while (n == BATCHSIZE);
2395 return 0;
2396
2397 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002398 Py_XDECREF(firstitem);
2399 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002400 return -1;
2401}
2402
Collin Winter5c9b02d2009-05-25 05:43:30 +00002403/* This is a variant of batch_dict() above that specializes for dicts, with no
2404 * support for dict subclasses. Like batch_dict(), we batch up chunks of
2405 * MARK key value ... key value SETITEMS
2406 * opcode sequences. Calling code should have arranged to first create an
2407 * empty dict, or dict-like object, for the SETITEMS to operate on.
2408 * Returns 0 on success, -1 on error.
2409 *
2410 * Note that this currently doesn't work for protocol 0.
2411 */
2412static int
2413batch_dict_exact(PicklerObject *self, PyObject *obj)
2414{
2415 PyObject *key = NULL, *value = NULL;
2416 int i;
2417 Py_ssize_t dict_size, ppos = 0;
2418
Alexandre Vassalottif70b1292009-05-25 18:00:52 +00002419 const char mark_op = MARK;
2420 const char setitem_op = SETITEM;
2421 const char setitems_op = SETITEMS;
Collin Winter5c9b02d2009-05-25 05:43:30 +00002422
2423 assert(obj != NULL);
2424 assert(self->proto > 0);
2425
2426 dict_size = PyDict_Size(obj);
2427
2428 /* Special-case len(d) == 1 to save space. */
2429 if (dict_size == 1) {
2430 PyDict_Next(obj, &ppos, &key, &value);
2431 if (save(self, key, 0) < 0)
2432 return -1;
2433 if (save(self, value, 0) < 0)
2434 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002435 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002436 return -1;
2437 return 0;
2438 }
2439
2440 /* Write in batches of BATCHSIZE. */
2441 do {
2442 i = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002443 if (_Pickler_Write(self, &mark_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002444 return -1;
2445 while (PyDict_Next(obj, &ppos, &key, &value)) {
2446 if (save(self, key, 0) < 0)
2447 return -1;
2448 if (save(self, value, 0) < 0)
2449 return -1;
2450 if (++i == BATCHSIZE)
2451 break;
2452 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002453 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002454 return -1;
2455 if (PyDict_Size(obj) != dict_size) {
2456 PyErr_Format(
2457 PyExc_RuntimeError,
2458 "dictionary changed size during iteration");
2459 return -1;
2460 }
2461
2462 } while (i == BATCHSIZE);
2463 return 0;
2464}
2465
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002466static int
2467save_dict(PicklerObject *self, PyObject *obj)
2468{
2469 PyObject *items, *iter;
2470 char header[3];
2471 int len;
2472 int status = 0;
2473
2474 if (self->fast && !fast_save_enter(self, obj))
2475 goto error;
2476
2477 /* Create an empty dict. */
2478 if (self->bin) {
2479 header[0] = EMPTY_DICT;
2480 len = 1;
2481 }
2482 else {
2483 header[0] = MARK;
2484 header[1] = DICT;
2485 len = 2;
2486 }
2487
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002488 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002489 goto error;
2490
2491 /* Get dict size, and bow out early if empty. */
2492 if ((len = PyDict_Size(obj)) < 0)
2493 goto error;
2494
2495 if (memo_put(self, obj) < 0)
2496 goto error;
2497
2498 if (len != 0) {
2499 /* Save the dict items. */
Collin Winter5c9b02d2009-05-25 05:43:30 +00002500 if (PyDict_CheckExact(obj) && self->proto > 0) {
2501 /* We can take certain shortcuts if we know this is a dict and
2502 not a dict subclass. */
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002503 if (Py_EnterRecursiveCall(" while pickling an object"))
2504 goto error;
2505 status = batch_dict_exact(self, obj);
2506 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002507 } else {
2508 items = PyObject_CallMethod(obj, "items", "()");
2509 if (items == NULL)
2510 goto error;
2511 iter = PyObject_GetIter(items);
2512 Py_DECREF(items);
2513 if (iter == NULL)
2514 goto error;
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002515 if (Py_EnterRecursiveCall(" while pickling an object")) {
2516 Py_DECREF(iter);
2517 goto error;
2518 }
Collin Winter5c9b02d2009-05-25 05:43:30 +00002519 status = batch_dict(self, iter);
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00002520 Py_LeaveRecursiveCall();
Collin Winter5c9b02d2009-05-25 05:43:30 +00002521 Py_DECREF(iter);
2522 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002523 }
2524
2525 if (0) {
2526 error:
2527 status = -1;
2528 }
2529
2530 if (self->fast && !fast_save_leave(self, obj))
2531 status = -1;
2532
2533 return status;
2534}
2535
2536static int
2537save_global(PicklerObject *self, PyObject *obj, PyObject *name)
2538{
2539 static PyObject *name_str = NULL;
2540 PyObject *global_name = NULL;
2541 PyObject *module_name = NULL;
2542 PyObject *module = NULL;
2543 PyObject *cls;
2544 int status = 0;
2545
2546 const char global_op = GLOBAL;
2547
2548 if (name_str == NULL) {
2549 name_str = PyUnicode_InternFromString("__name__");
2550 if (name_str == NULL)
2551 goto error;
2552 }
2553
2554 if (name) {
2555 global_name = name;
2556 Py_INCREF(global_name);
2557 }
2558 else {
2559 global_name = PyObject_GetAttr(obj, name_str);
2560 if (global_name == NULL)
2561 goto error;
2562 }
2563
2564 module_name = whichmodule(obj, global_name);
2565 if (module_name == NULL)
2566 goto error;
2567
2568 /* XXX: Change to use the import C API directly with level=0 to disallow
2569 relative imports.
2570
2571 XXX: PyImport_ImportModuleLevel could be used. However, this bypasses
2572 builtins.__import__. Therefore, _pickle, unlike pickle.py, will ignore
2573 custom import functions (IMHO, this would be a nice security
2574 feature). The import C API would need to be extended to support the
2575 extra parameters of __import__ to fix that. */
2576 module = PyImport_Import(module_name);
2577 if (module == NULL) {
2578 PyErr_Format(PicklingError,
2579 "Can't pickle %R: import of module %R failed",
2580 obj, module_name);
2581 goto error;
2582 }
2583 cls = PyObject_GetAttr(module, global_name);
2584 if (cls == NULL) {
2585 PyErr_Format(PicklingError,
2586 "Can't pickle %R: attribute lookup %S.%S failed",
2587 obj, module_name, global_name);
2588 goto error;
2589 }
2590 if (cls != obj) {
2591 Py_DECREF(cls);
2592 PyErr_Format(PicklingError,
2593 "Can't pickle %R: it's not the same object as %S.%S",
2594 obj, module_name, global_name);
2595 goto error;
2596 }
2597 Py_DECREF(cls);
2598
2599 if (self->proto >= 2) {
2600 /* See whether this is in the extension registry, and if
2601 * so generate an EXT opcode.
2602 */
2603 PyObject *code_obj; /* extension code as Python object */
2604 long code; /* extension code as C value */
2605 char pdata[5];
2606 int n;
2607
2608 PyTuple_SET_ITEM(two_tuple, 0, module_name);
2609 PyTuple_SET_ITEM(two_tuple, 1, global_name);
2610 code_obj = PyDict_GetItem(extension_registry, two_tuple);
2611 /* The object is not registered in the extension registry.
2612 This is the most likely code path. */
2613 if (code_obj == NULL)
2614 goto gen_global;
2615
2616 /* XXX: pickle.py doesn't check neither the type, nor the range
2617 of the value returned by the extension_registry. It should for
2618 consistency. */
2619
2620 /* Verify code_obj has the right type and value. */
2621 if (!PyLong_Check(code_obj)) {
2622 PyErr_Format(PicklingError,
2623 "Can't pickle %R: extension code %R isn't an integer",
2624 obj, code_obj);
2625 goto error;
2626 }
2627 code = PyLong_AS_LONG(code_obj);
2628 if (code <= 0 || code > 0x7fffffffL) {
2629 PyErr_Format(PicklingError,
2630 "Can't pickle %R: extension code %ld is out of range",
2631 obj, code);
2632 goto error;
2633 }
2634
2635 /* Generate an EXT opcode. */
2636 if (code <= 0xff) {
2637 pdata[0] = EXT1;
2638 pdata[1] = (unsigned char)code;
2639 n = 2;
2640 }
2641 else if (code <= 0xffff) {
2642 pdata[0] = EXT2;
2643 pdata[1] = (unsigned char)(code & 0xff);
2644 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2645 n = 3;
2646 }
2647 else {
2648 pdata[0] = EXT4;
2649 pdata[1] = (unsigned char)(code & 0xff);
2650 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2651 pdata[3] = (unsigned char)((code >> 16) & 0xff);
2652 pdata[4] = (unsigned char)((code >> 24) & 0xff);
2653 n = 5;
2654 }
2655
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002656 if (_Pickler_Write(self, pdata, n) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002657 goto error;
2658 }
2659 else {
2660 /* Generate a normal global opcode if we are using a pickle
2661 protocol <= 2, or if the object is not registered in the
2662 extension registry. */
2663 PyObject *encoded;
2664 PyObject *(*unicode_encoder)(PyObject *);
2665
2666 gen_global:
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002667 if (_Pickler_Write(self, &global_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002668 goto error;
2669
2670 /* Since Python 3.0 now supports non-ASCII identifiers, we encode both
2671 the module name and the global name using UTF-8. We do so only when
2672 we are using the pickle protocol newer than version 3. This is to
2673 ensure compatibility with older Unpickler running on Python 2.x. */
2674 if (self->proto >= 3) {
2675 unicode_encoder = PyUnicode_AsUTF8String;
2676 }
2677 else {
2678 unicode_encoder = PyUnicode_AsASCIIString;
2679 }
2680
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00002681 /* For protocol < 3 and if the user didn't request against doing so,
2682 we convert module names to the old 2.x module names. */
2683 if (self->fix_imports) {
2684 PyObject *key;
2685 PyObject *item;
2686
2687 key = PyTuple_Pack(2, module_name, global_name);
2688 if (key == NULL)
2689 goto error;
2690 item = PyDict_GetItemWithError(name_mapping_3to2, key);
2691 Py_DECREF(key);
2692 if (item) {
2693 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
2694 PyErr_Format(PyExc_RuntimeError,
2695 "_compat_pickle.REVERSE_NAME_MAPPING values "
2696 "should be 2-tuples, not %.200s",
2697 Py_TYPE(item)->tp_name);
2698 goto error;
2699 }
2700 Py_CLEAR(module_name);
2701 Py_CLEAR(global_name);
2702 module_name = PyTuple_GET_ITEM(item, 0);
2703 global_name = PyTuple_GET_ITEM(item, 1);
2704 if (!PyUnicode_Check(module_name) ||
2705 !PyUnicode_Check(global_name)) {
2706 PyErr_Format(PyExc_RuntimeError,
2707 "_compat_pickle.REVERSE_NAME_MAPPING values "
2708 "should be pairs of str, not (%.200s, %.200s)",
2709 Py_TYPE(module_name)->tp_name,
2710 Py_TYPE(global_name)->tp_name);
2711 goto error;
2712 }
2713 Py_INCREF(module_name);
2714 Py_INCREF(global_name);
2715 }
2716 else if (PyErr_Occurred()) {
2717 goto error;
2718 }
2719
2720 item = PyDict_GetItemWithError(import_mapping_3to2, module_name);
2721 if (item) {
2722 if (!PyUnicode_Check(item)) {
2723 PyErr_Format(PyExc_RuntimeError,
2724 "_compat_pickle.REVERSE_IMPORT_MAPPING values "
2725 "should be strings, not %.200s",
2726 Py_TYPE(item)->tp_name);
2727 goto error;
2728 }
2729 Py_CLEAR(module_name);
2730 module_name = item;
2731 Py_INCREF(module_name);
2732 }
2733 else if (PyErr_Occurred()) {
2734 goto error;
2735 }
2736 }
2737
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002738 /* Save the name of the module. */
2739 encoded = unicode_encoder(module_name);
2740 if (encoded == NULL) {
2741 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2742 PyErr_Format(PicklingError,
2743 "can't pickle module identifier '%S' using "
2744 "pickle protocol %i", module_name, self->proto);
2745 goto error;
2746 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002747 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002748 PyBytes_GET_SIZE(encoded)) < 0) {
2749 Py_DECREF(encoded);
2750 goto error;
2751 }
2752 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002753 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002754 goto error;
2755
2756 /* Save the name of the module. */
2757 encoded = unicode_encoder(global_name);
2758 if (encoded == NULL) {
2759 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2760 PyErr_Format(PicklingError,
2761 "can't pickle global identifier '%S' using "
2762 "pickle protocol %i", global_name, self->proto);
2763 goto error;
2764 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002765 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002766 PyBytes_GET_SIZE(encoded)) < 0) {
2767 Py_DECREF(encoded);
2768 goto error;
2769 }
2770 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002771 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002772 goto error;
2773
2774 /* Memoize the object. */
2775 if (memo_put(self, obj) < 0)
2776 goto error;
2777 }
2778
2779 if (0) {
2780 error:
2781 status = -1;
2782 }
2783 Py_XDECREF(module_name);
2784 Py_XDECREF(global_name);
2785 Py_XDECREF(module);
2786
2787 return status;
2788}
2789
2790static int
2791save_pers(PicklerObject *self, PyObject *obj, PyObject *func)
2792{
2793 PyObject *pid = NULL;
2794 int status = 0;
2795
2796 const char persid_op = PERSID;
2797 const char binpersid_op = BINPERSID;
2798
2799 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002800 pid = _Pickler_FastCall(self, func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002801 if (pid == NULL)
2802 return -1;
2803
2804 if (pid != Py_None) {
2805 if (self->bin) {
2806 if (save(self, pid, 1) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002807 _Pickler_Write(self, &binpersid_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002808 goto error;
2809 }
2810 else {
2811 PyObject *pid_str = NULL;
2812 char *pid_ascii_bytes;
2813 Py_ssize_t size;
2814
2815 pid_str = PyObject_Str(pid);
2816 if (pid_str == NULL)
2817 goto error;
2818
2819 /* XXX: Should it check whether the persistent id only contains
2820 ASCII characters? And what if the pid contains embedded
2821 newlines? */
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00002822 pid_ascii_bytes = _PyUnicode_AsStringAndSize(pid_str, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002823 Py_DECREF(pid_str);
2824 if (pid_ascii_bytes == NULL)
2825 goto error;
2826
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002827 if (_Pickler_Write(self, &persid_op, 1) < 0 ||
2828 _Pickler_Write(self, pid_ascii_bytes, size) < 0 ||
2829 _Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002830 goto error;
2831 }
2832 status = 1;
2833 }
2834
2835 if (0) {
2836 error:
2837 status = -1;
2838 }
2839 Py_XDECREF(pid);
2840
2841 return status;
2842}
2843
2844/* We're saving obj, and args is the 2-thru-5 tuple returned by the
2845 * appropriate __reduce__ method for obj.
2846 */
2847static int
2848save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
2849{
2850 PyObject *callable;
2851 PyObject *argtup;
2852 PyObject *state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002853 PyObject *listitems = Py_None;
2854 PyObject *dictitems = Py_None;
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002855 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002856
2857 int use_newobj = self->proto >= 2;
2858
2859 const char reduce_op = REDUCE;
2860 const char build_op = BUILD;
2861 const char newobj_op = NEWOBJ;
2862
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002863 size = PyTuple_Size(args);
2864 if (size < 2 || size > 5) {
2865 PyErr_SetString(PicklingError, "tuple returned by "
2866 "__reduce__ must contain 2 through 5 elements");
2867 return -1;
2868 }
2869
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002870 if (!PyArg_UnpackTuple(args, "save_reduce", 2, 5,
2871 &callable, &argtup, &state, &listitems, &dictitems))
2872 return -1;
2873
2874 if (!PyCallable_Check(callable)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002875 PyErr_SetString(PicklingError, "first item of the tuple "
2876 "returned by __reduce__ must be callable");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002877 return -1;
2878 }
2879 if (!PyTuple_Check(argtup)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002880 PyErr_SetString(PicklingError, "second item of the tuple "
2881 "returned by __reduce__ must be a tuple");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002882 return -1;
2883 }
2884
2885 if (state == Py_None)
2886 state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002887
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002888 if (listitems == Py_None)
2889 listitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002890 else if (!PyIter_Check(listitems)) {
2891 PyErr_Format(PicklingError, "Fourth element of tuple"
2892 "returned by __reduce__ must be an iterator, not %s",
2893 Py_TYPE(listitems)->tp_name);
2894 return -1;
2895 }
2896
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002897 if (dictitems == Py_None)
2898 dictitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002899 else if (!PyIter_Check(dictitems)) {
2900 PyErr_Format(PicklingError, "Fifth element of tuple"
2901 "returned by __reduce__ must be an iterator, not %s",
2902 Py_TYPE(dictitems)->tp_name);
2903 return -1;
2904 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002905
2906 /* Protocol 2 special case: if callable's name is __newobj__, use
2907 NEWOBJ. */
2908 if (use_newobj) {
Antoine Pitrouff150f22010-10-22 21:41:05 +00002909 static PyObject *newobj_str = NULL;
2910 PyObject *name_str;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002911
2912 if (newobj_str == NULL) {
2913 newobj_str = PyUnicode_InternFromString("__newobj__");
Antoine Pitrouff150f22010-10-22 21:41:05 +00002914 if (newobj_str == NULL)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002915 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002916 }
2917
Antoine Pitrouff150f22010-10-22 21:41:05 +00002918 name_str = PyObject_GetAttrString(callable, "__name__");
2919 if (name_str == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002920 if (PyErr_ExceptionMatches(PyExc_AttributeError))
2921 PyErr_Clear();
2922 else
2923 return -1;
2924 use_newobj = 0;
2925 }
2926 else {
Antoine Pitrouff150f22010-10-22 21:41:05 +00002927 use_newobj = PyUnicode_Check(name_str) &&
2928 PyUnicode_Compare(name_str, newobj_str) == 0;
2929 Py_DECREF(name_str);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002930 }
2931 }
2932 if (use_newobj) {
2933 PyObject *cls;
2934 PyObject *newargtup;
2935 PyObject *obj_class;
2936 int p;
2937
2938 /* Sanity checks. */
2939 if (Py_SIZE(argtup) < 1) {
2940 PyErr_SetString(PicklingError, "__newobj__ arglist is empty");
2941 return -1;
2942 }
2943
2944 cls = PyTuple_GET_ITEM(argtup, 0);
Antoine Pitrouff150f22010-10-22 21:41:05 +00002945 if (!PyObject_HasAttrString(cls, "__new__")) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002946 PyErr_SetString(PicklingError, "args[0] from "
Antoine Pitrouff150f22010-10-22 21:41:05 +00002947 "__newobj__ args has no __new__");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002948 return -1;
2949 }
2950
2951 if (obj != NULL) {
Antoine Pitrouff150f22010-10-22 21:41:05 +00002952 obj_class = PyObject_GetAttrString(obj, "__class__");
2953 if (obj_class == NULL) {
2954 if (PyErr_ExceptionMatches(PyExc_AttributeError))
2955 PyErr_Clear();
2956 else
2957 return -1;
2958 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002959 p = obj_class != cls; /* true iff a problem */
2960 Py_DECREF(obj_class);
2961 if (p) {
2962 PyErr_SetString(PicklingError, "args[0] from "
2963 "__newobj__ args has the wrong class");
2964 return -1;
2965 }
2966 }
2967 /* XXX: These calls save() are prone to infinite recursion. Imagine
2968 what happen if the value returned by the __reduce__() method of
2969 some extension type contains another object of the same type. Ouch!
2970
2971 Here is a quick example, that I ran into, to illustrate what I
2972 mean:
2973
2974 >>> import pickle, copyreg
2975 >>> copyreg.dispatch_table.pop(complex)
2976 >>> pickle.dumps(1+2j)
2977 Traceback (most recent call last):
2978 ...
2979 RuntimeError: maximum recursion depth exceeded
2980
2981 Removing the complex class from copyreg.dispatch_table made the
2982 __reduce_ex__() method emit another complex object:
2983
2984 >>> (1+1j).__reduce_ex__(2)
2985 (<function __newobj__ at 0xb7b71c3c>,
2986 (<class 'complex'>, (1+1j)), None, None, None)
2987
2988 Thus when save() was called on newargstup (the 2nd item) recursion
2989 ensued. Of course, the bug was in the complex class which had a
2990 broken __getnewargs__() that emitted another complex object. But,
2991 the point, here, is it is quite easy to end up with a broken reduce
2992 function. */
2993
2994 /* Save the class and its __new__ arguments. */
2995 if (save(self, cls, 0) < 0)
2996 return -1;
2997
2998 newargtup = PyTuple_GetSlice(argtup, 1, Py_SIZE(argtup));
2999 if (newargtup == NULL)
3000 return -1;
3001
3002 p = save(self, newargtup, 0);
3003 Py_DECREF(newargtup);
3004 if (p < 0)
3005 return -1;
3006
3007 /* Add NEWOBJ opcode. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003008 if (_Pickler_Write(self, &newobj_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003009 return -1;
3010 }
3011 else { /* Not using NEWOBJ. */
3012 if (save(self, callable, 0) < 0 ||
3013 save(self, argtup, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003014 _Pickler_Write(self, &reduce_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003015 return -1;
3016 }
3017
3018 /* obj can be NULL when save_reduce() is used directly. A NULL obj means
3019 the caller do not want to memoize the object. Not particularly useful,
3020 but that is to mimic the behavior save_reduce() in pickle.py when
3021 obj is None. */
3022 if (obj && memo_put(self, obj) < 0)
3023 return -1;
3024
3025 if (listitems && batch_list(self, listitems) < 0)
3026 return -1;
3027
3028 if (dictitems && batch_dict(self, dictitems) < 0)
3029 return -1;
3030
3031 if (state) {
3032 if (save(self, state, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003033 _Pickler_Write(self, &build_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003034 return -1;
3035 }
3036
3037 return 0;
3038}
3039
3040static int
3041save(PicklerObject *self, PyObject *obj, int pers_save)
3042{
3043 PyTypeObject *type;
3044 PyObject *reduce_func = NULL;
3045 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003046 int status = 0;
3047
Antoine Pitroue6d4c5b2011-01-23 17:12:25 +00003048 if (Py_EnterRecursiveCall(" while pickling an object"))
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003049 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003050
3051 /* The extra pers_save argument is necessary to avoid calling save_pers()
3052 on its returned object. */
3053 if (!pers_save && self->pers_func) {
3054 /* save_pers() returns:
3055 -1 to signal an error;
3056 0 if it did nothing successfully;
3057 1 if a persistent id was saved.
3058 */
3059 if ((status = save_pers(self, obj, self->pers_func)) != 0)
3060 goto done;
3061 }
3062
3063 type = Py_TYPE(obj);
3064
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003065 /* The old cPickle had an optimization that used switch-case statement
3066 dispatching on the first letter of the type name. This has was removed
3067 since benchmarks shown that this optimization was actually slowing
3068 things down. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003069
3070 /* Atom types; these aren't memoized, so don't check the memo. */
3071
3072 if (obj == Py_None) {
3073 status = save_none(self, obj);
3074 goto done;
3075 }
3076 else if (obj == Py_False || obj == Py_True) {
3077 status = save_bool(self, obj);
3078 goto done;
3079 }
3080 else if (type == &PyLong_Type) {
3081 status = save_long(self, obj);
3082 goto done;
3083 }
3084 else if (type == &PyFloat_Type) {
3085 status = save_float(self, obj);
3086 goto done;
3087 }
3088
3089 /* Check the memo to see if it has the object. If so, generate
3090 a GET (or BINGET) opcode, instead of pickling the object
3091 once again. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003092 if (PyMemoTable_Get(self->memo, obj)) {
3093 if (memo_get(self, obj) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003094 goto error;
3095 goto done;
3096 }
3097
3098 if (type == &PyBytes_Type) {
3099 status = save_bytes(self, obj);
3100 goto done;
3101 }
3102 else if (type == &PyUnicode_Type) {
3103 status = save_unicode(self, obj);
3104 goto done;
3105 }
3106 else if (type == &PyDict_Type) {
3107 status = save_dict(self, obj);
3108 goto done;
3109 }
3110 else if (type == &PyList_Type) {
3111 status = save_list(self, obj);
3112 goto done;
3113 }
3114 else if (type == &PyTuple_Type) {
3115 status = save_tuple(self, obj);
3116 goto done;
3117 }
3118 else if (type == &PyType_Type) {
3119 status = save_global(self, obj, NULL);
3120 goto done;
3121 }
3122 else if (type == &PyFunction_Type) {
3123 status = save_global(self, obj, NULL);
3124 if (status < 0 && PyErr_ExceptionMatches(PickleError)) {
3125 /* fall back to reduce */
3126 PyErr_Clear();
3127 }
3128 else {
3129 goto done;
3130 }
3131 }
3132 else if (type == &PyCFunction_Type) {
3133 status = save_global(self, obj, NULL);
3134 goto done;
3135 }
3136 else if (PyType_IsSubtype(type, &PyType_Type)) {
3137 status = save_global(self, obj, NULL);
3138 goto done;
3139 }
3140
3141 /* XXX: This part needs some unit tests. */
3142
3143 /* Get a reduction callable, and call it. This may come from
3144 * copyreg.dispatch_table, the object's __reduce_ex__ method,
3145 * or the object's __reduce__ method.
3146 */
3147 reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
3148 if (reduce_func != NULL) {
3149 /* Here, the reference count of the reduce_func object returned by
3150 PyDict_GetItem needs to be increased to be consistent with the one
3151 returned by PyObject_GetAttr. This is allow us to blindly DECREF
3152 reduce_func at the end of the save() routine.
3153 */
3154 Py_INCREF(reduce_func);
3155 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003156 reduce_value = _Pickler_FastCall(self, reduce_func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003157 }
3158 else {
3159 static PyObject *reduce_str = NULL;
3160 static PyObject *reduce_ex_str = NULL;
3161
3162 /* Cache the name of the reduce methods. */
3163 if (reduce_str == NULL) {
3164 reduce_str = PyUnicode_InternFromString("__reduce__");
3165 if (reduce_str == NULL)
3166 goto error;
3167 reduce_ex_str = PyUnicode_InternFromString("__reduce_ex__");
3168 if (reduce_ex_str == NULL)
3169 goto error;
3170 }
3171
3172 /* XXX: If the __reduce__ method is defined, __reduce_ex__ is
3173 automatically defined as __reduce__. While this is convenient, this
3174 make it impossible to know which method was actually called. Of
3175 course, this is not a big deal. But still, it would be nice to let
3176 the user know which method was called when something go
3177 wrong. Incidentally, this means if __reduce_ex__ is not defined, we
3178 don't actually have to check for a __reduce__ method. */
3179
3180 /* Check for a __reduce_ex__ method. */
3181 reduce_func = PyObject_GetAttr(obj, reduce_ex_str);
3182 if (reduce_func != NULL) {
3183 PyObject *proto;
3184 proto = PyLong_FromLong(self->proto);
3185 if (proto != NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003186 reduce_value = _Pickler_FastCall(self, reduce_func, proto);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003187 }
3188 }
3189 else {
3190 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3191 PyErr_Clear();
3192 else
3193 goto error;
3194 /* Check for a __reduce__ method. */
3195 reduce_func = PyObject_GetAttr(obj, reduce_str);
3196 if (reduce_func != NULL) {
3197 reduce_value = PyObject_Call(reduce_func, empty_tuple, NULL);
3198 }
3199 else {
3200 PyErr_Format(PicklingError, "can't pickle '%.200s' object: %R",
3201 type->tp_name, obj);
3202 goto error;
3203 }
3204 }
3205 }
3206
3207 if (reduce_value == NULL)
3208 goto error;
3209
3210 if (PyUnicode_Check(reduce_value)) {
3211 status = save_global(self, obj, reduce_value);
3212 goto done;
3213 }
3214
3215 if (!PyTuple_Check(reduce_value)) {
3216 PyErr_SetString(PicklingError,
3217 "__reduce__ must return a string or tuple");
3218 goto error;
3219 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003220
3221 status = save_reduce(self, reduce_value, obj);
3222
3223 if (0) {
3224 error:
3225 status = -1;
3226 }
3227 done:
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003228 Py_LeaveRecursiveCall();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003229 Py_XDECREF(reduce_func);
3230 Py_XDECREF(reduce_value);
3231
3232 return status;
3233}
3234
3235static int
3236dump(PicklerObject *self, PyObject *obj)
3237{
3238 const char stop_op = STOP;
3239
3240 if (self->proto >= 2) {
3241 char header[2];
3242
3243 header[0] = PROTO;
3244 assert(self->proto >= 0 && self->proto < 256);
3245 header[1] = (unsigned char)self->proto;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003246 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003247 return -1;
3248 }
3249
3250 if (save(self, obj, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003251 _Pickler_Write(self, &stop_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003252 return -1;
3253
3254 return 0;
3255}
3256
3257PyDoc_STRVAR(Pickler_clear_memo_doc,
3258"clear_memo() -> None. Clears the pickler's \"memo\"."
3259"\n"
3260"The memo is the data structure that remembers which objects the\n"
3261"pickler has already seen, so that shared or recursive objects are\n"
3262"pickled by reference and not by value. This method is useful when\n"
3263"re-using picklers.");
3264
3265static PyObject *
3266Pickler_clear_memo(PicklerObject *self)
3267{
3268 if (self->memo)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003269 PyMemoTable_Clear(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003270
3271 Py_RETURN_NONE;
3272}
3273
3274PyDoc_STRVAR(Pickler_dump_doc,
3275"dump(obj) -> None. Write a pickled representation of obj to the open file.");
3276
3277static PyObject *
3278Pickler_dump(PicklerObject *self, PyObject *args)
3279{
3280 PyObject *obj;
3281
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003282 /* Check whether the Pickler was initialized correctly (issue3664).
3283 Developers often forget to call __init__() in their subclasses, which
3284 would trigger a segfault without this check. */
3285 if (self->write == NULL) {
3286 PyErr_Format(PicklingError,
3287 "Pickler.__init__() was not called by %s.__init__()",
3288 Py_TYPE(self)->tp_name);
3289 return NULL;
3290 }
3291
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003292 if (!PyArg_ParseTuple(args, "O:dump", &obj))
3293 return NULL;
3294
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003295 if (_Pickler_ClearBuffer(self) < 0)
3296 return NULL;
3297
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003298 if (dump(self, obj) < 0)
3299 return NULL;
3300
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003301 if (_Pickler_FlushToFile(self) < 0)
3302 return NULL;
3303
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003304 Py_RETURN_NONE;
3305}
3306
3307static struct PyMethodDef Pickler_methods[] = {
3308 {"dump", (PyCFunction)Pickler_dump, METH_VARARGS,
3309 Pickler_dump_doc},
3310 {"clear_memo", (PyCFunction)Pickler_clear_memo, METH_NOARGS,
3311 Pickler_clear_memo_doc},
3312 {NULL, NULL} /* sentinel */
3313};
3314
3315static void
3316Pickler_dealloc(PicklerObject *self)
3317{
3318 PyObject_GC_UnTrack(self);
3319
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003320 Py_XDECREF(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003321 Py_XDECREF(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003322 Py_XDECREF(self->pers_func);
3323 Py_XDECREF(self->arg);
3324 Py_XDECREF(self->fast_memo);
3325
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003326 PyMemoTable_Del(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003327
3328 Py_TYPE(self)->tp_free((PyObject *)self);
3329}
3330
3331static int
3332Pickler_traverse(PicklerObject *self, visitproc visit, void *arg)
3333{
3334 Py_VISIT(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003335 Py_VISIT(self->pers_func);
3336 Py_VISIT(self->arg);
3337 Py_VISIT(self->fast_memo);
3338 return 0;
3339}
3340
3341static int
3342Pickler_clear(PicklerObject *self)
3343{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003344 Py_CLEAR(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003345 Py_CLEAR(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003346 Py_CLEAR(self->pers_func);
3347 Py_CLEAR(self->arg);
3348 Py_CLEAR(self->fast_memo);
3349
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003350 if (self->memo != NULL) {
3351 PyMemoTable *memo = self->memo;
3352 self->memo = NULL;
3353 PyMemoTable_Del(memo);
3354 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003355 return 0;
3356}
3357
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003358
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003359PyDoc_STRVAR(Pickler_doc,
3360"Pickler(file, protocol=None)"
3361"\n"
3362"This takes a binary file for writing a pickle data stream.\n"
3363"\n"
3364"The optional protocol argument tells the pickler to use the\n"
3365"given protocol; supported protocols are 0, 1, 2, 3. The default\n"
3366"protocol is 3; a backward-incompatible protocol designed for\n"
3367"Python 3.0.\n"
3368"\n"
3369"Specifying a negative protocol version selects the highest\n"
3370"protocol version supported. The higher the protocol used, the\n"
3371"more recent the version of Python needed to read the pickle\n"
3372"produced.\n"
3373"\n"
3374"The file argument must have a write() method that accepts a single\n"
3375"bytes argument. It can thus be a file object opened for binary\n"
3376"writing, a io.BytesIO instance, or any other custom object that\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003377"meets this interface.\n"
3378"\n"
3379"If fix_imports is True and protocol is less than 3, pickle will try to\n"
3380"map the new Python 3.x names to the old module names used in Python\n"
3381"2.x, so that the pickle data stream is readable with Python 2.x.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003382
3383static int
3384Pickler_init(PicklerObject *self, PyObject *args, PyObject *kwds)
3385{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003386 static char *kwlist[] = {"file", "protocol", "fix_imports", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003387 PyObject *file;
3388 PyObject *proto_obj = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003389 PyObject *fix_imports = Py_True;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003390
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003391 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:Pickler",
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003392 kwlist, &file, &proto_obj, &fix_imports))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003393 return -1;
3394
3395 /* In case of multiple __init__() calls, clear previous content. */
3396 if (self->write != NULL)
3397 (void)Pickler_clear(self);
3398
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003399 if (_Pickler_SetProtocol(self, proto_obj, fix_imports) < 0)
3400 return -1;
3401
3402 if (_Pickler_SetOutputStream(self, file) < 0)
3403 return -1;
3404
3405 /* memo and output_buffer may have already been created in _Pickler_New */
3406 if (self->memo == NULL) {
3407 self->memo = PyMemoTable_New();
3408 if (self->memo == NULL)
3409 return -1;
3410 }
3411 self->output_len = 0;
3412 if (self->output_buffer == NULL) {
3413 self->max_output_len = WRITE_BUF_SIZE;
3414 self->output_buffer = PyBytes_FromStringAndSize(NULL,
3415 self->max_output_len);
3416 if (self->output_buffer == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003417 return -1;
3418 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003419
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003420 self->arg = NULL;
3421 self->fast = 0;
3422 self->fast_nesting = 0;
3423 self->fast_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003424 self->pers_func = NULL;
3425 if (PyObject_HasAttrString((PyObject *)self, "persistent_id")) {
3426 self->pers_func = PyObject_GetAttrString((PyObject *)self,
3427 "persistent_id");
3428 if (self->pers_func == NULL)
3429 return -1;
3430 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003431 return 0;
3432}
3433
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003434/* Define a proxy object for the Pickler's internal memo object. This is to
3435 * avoid breaking code like:
3436 * pickler.memo.clear()
3437 * and
3438 * pickler.memo = saved_memo
3439 * Is this a good idea? Not really, but we don't want to break code that uses
3440 * it. Note that we don't implement the entire mapping API here. This is
3441 * intentional, as these should be treated as black-box implementation details.
3442 */
3443
3444typedef struct {
3445 PyObject_HEAD
3446 PicklerObject *pickler; /* Pickler whose memo table we're proxying. */
3447} PicklerMemoProxyObject;
3448
3449PyDoc_STRVAR(pmp_clear_doc,
3450"memo.clear() -> None. Remove all items from memo.");
3451
3452static PyObject *
3453pmp_clear(PicklerMemoProxyObject *self)
3454{
3455 if (self->pickler->memo)
3456 PyMemoTable_Clear(self->pickler->memo);
3457 Py_RETURN_NONE;
3458}
3459
3460PyDoc_STRVAR(pmp_copy_doc,
3461"memo.copy() -> new_memo. Copy the memo to a new object.");
3462
3463static PyObject *
3464pmp_copy(PicklerMemoProxyObject *self)
3465{
3466 Py_ssize_t i;
3467 PyMemoTable *memo;
3468 PyObject *new_memo = PyDict_New();
3469 if (new_memo == NULL)
3470 return NULL;
3471
3472 memo = self->pickler->memo;
3473 for (i = 0; i < memo->mt_allocated; ++i) {
3474 PyMemoEntry entry = memo->mt_table[i];
3475 if (entry.me_key != NULL) {
3476 int status;
3477 PyObject *key, *value;
3478
3479 key = PyLong_FromVoidPtr(entry.me_key);
3480 value = Py_BuildValue("lO", entry.me_value, entry.me_key);
3481
3482 if (key == NULL || value == NULL) {
3483 Py_XDECREF(key);
3484 Py_XDECREF(value);
3485 goto error;
3486 }
3487 status = PyDict_SetItem(new_memo, key, value);
3488 Py_DECREF(key);
3489 Py_DECREF(value);
3490 if (status < 0)
3491 goto error;
3492 }
3493 }
3494 return new_memo;
3495
3496 error:
3497 Py_XDECREF(new_memo);
3498 return NULL;
3499}
3500
3501PyDoc_STRVAR(pmp_reduce_doc,
3502"memo.__reduce__(). Pickling support.");
3503
3504static PyObject *
3505pmp_reduce(PicklerMemoProxyObject *self, PyObject *args)
3506{
3507 PyObject *reduce_value, *dict_args;
3508 PyObject *contents = pmp_copy(self);
3509 if (contents == NULL)
3510 return NULL;
3511
3512 reduce_value = PyTuple_New(2);
3513 if (reduce_value == NULL) {
3514 Py_DECREF(contents);
3515 return NULL;
3516 }
3517 dict_args = PyTuple_New(1);
3518 if (dict_args == NULL) {
3519 Py_DECREF(contents);
3520 Py_DECREF(reduce_value);
3521 return NULL;
3522 }
3523 PyTuple_SET_ITEM(dict_args, 0, contents);
3524 Py_INCREF((PyObject *)&PyDict_Type);
3525 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
3526 PyTuple_SET_ITEM(reduce_value, 1, dict_args);
3527 return reduce_value;
3528}
3529
3530static PyMethodDef picklerproxy_methods[] = {
3531 {"clear", (PyCFunction)pmp_clear, METH_NOARGS, pmp_clear_doc},
3532 {"copy", (PyCFunction)pmp_copy, METH_NOARGS, pmp_copy_doc},
3533 {"__reduce__", (PyCFunction)pmp_reduce, METH_VARARGS, pmp_reduce_doc},
3534 {NULL, NULL} /* sentinel */
3535};
3536
3537static void
3538PicklerMemoProxy_dealloc(PicklerMemoProxyObject *self)
3539{
3540 PyObject_GC_UnTrack(self);
3541 Py_XDECREF(self->pickler);
3542 PyObject_GC_Del((PyObject *)self);
3543}
3544
3545static int
3546PicklerMemoProxy_traverse(PicklerMemoProxyObject *self,
3547 visitproc visit, void *arg)
3548{
3549 Py_VISIT(self->pickler);
3550 return 0;
3551}
3552
3553static int
3554PicklerMemoProxy_clear(PicklerMemoProxyObject *self)
3555{
3556 Py_CLEAR(self->pickler);
3557 return 0;
3558}
3559
3560static PyTypeObject PicklerMemoProxyType = {
3561 PyVarObject_HEAD_INIT(NULL, 0)
3562 "_pickle.PicklerMemoProxy", /*tp_name*/
3563 sizeof(PicklerMemoProxyObject), /*tp_basicsize*/
3564 0,
3565 (destructor)PicklerMemoProxy_dealloc, /* tp_dealloc */
3566 0, /* tp_print */
3567 0, /* tp_getattr */
3568 0, /* tp_setattr */
3569 0, /* tp_compare */
3570 0, /* tp_repr */
3571 0, /* tp_as_number */
3572 0, /* tp_as_sequence */
3573 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00003574 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003575 0, /* tp_call */
3576 0, /* tp_str */
3577 PyObject_GenericGetAttr, /* tp_getattro */
3578 PyObject_GenericSetAttr, /* tp_setattro */
3579 0, /* tp_as_buffer */
3580 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3581 0, /* tp_doc */
3582 (traverseproc)PicklerMemoProxy_traverse, /* tp_traverse */
3583 (inquiry)PicklerMemoProxy_clear, /* tp_clear */
3584 0, /* tp_richcompare */
3585 0, /* tp_weaklistoffset */
3586 0, /* tp_iter */
3587 0, /* tp_iternext */
3588 picklerproxy_methods, /* tp_methods */
3589};
3590
3591static PyObject *
3592PicklerMemoProxy_New(PicklerObject *pickler)
3593{
3594 PicklerMemoProxyObject *self;
3595
3596 self = PyObject_GC_New(PicklerMemoProxyObject, &PicklerMemoProxyType);
3597 if (self == NULL)
3598 return NULL;
3599 Py_INCREF(pickler);
3600 self->pickler = pickler;
3601 PyObject_GC_Track(self);
3602 return (PyObject *)self;
3603}
3604
3605/*****************************************************************************/
3606
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003607static PyObject *
3608Pickler_get_memo(PicklerObject *self)
3609{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003610 return PicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003611}
3612
3613static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003614Pickler_set_memo(PicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003615{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003616 PyMemoTable *new_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003617
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003618 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003619 PyErr_SetString(PyExc_TypeError,
3620 "attribute deletion is not supported");
3621 return -1;
3622 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003623
3624 if (Py_TYPE(obj) == &PicklerMemoProxyType) {
3625 PicklerObject *pickler =
3626 ((PicklerMemoProxyObject *)obj)->pickler;
3627
3628 new_memo = PyMemoTable_Copy(pickler->memo);
3629 if (new_memo == NULL)
3630 return -1;
3631 }
3632 else if (PyDict_Check(obj)) {
3633 Py_ssize_t i = 0;
3634 PyObject *key, *value;
3635
3636 new_memo = PyMemoTable_New();
3637 if (new_memo == NULL)
3638 return -1;
3639
3640 while (PyDict_Next(obj, &i, &key, &value)) {
3641 long memo_id;
3642 PyObject *memo_obj;
3643
3644 if (!PyTuple_Check(value) || Py_SIZE(value) != 2) {
3645 PyErr_SetString(PyExc_TypeError,
3646 "'memo' values must be 2-item tuples");
3647 goto error;
3648 }
3649 memo_id = PyLong_AsLong(PyTuple_GET_ITEM(value, 0));
3650 if (memo_id == -1 && PyErr_Occurred())
3651 goto error;
3652 memo_obj = PyTuple_GET_ITEM(value, 1);
3653 if (PyMemoTable_Set(new_memo, memo_obj, memo_id) < 0)
3654 goto error;
3655 }
3656 }
3657 else {
3658 PyErr_Format(PyExc_TypeError,
3659 "'memo' attribute must be an PicklerMemoProxy object"
3660 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003661 return -1;
3662 }
3663
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003664 PyMemoTable_Del(self->memo);
3665 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003666
3667 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003668
3669 error:
3670 if (new_memo)
3671 PyMemoTable_Del(new_memo);
3672 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003673}
3674
3675static PyObject *
3676Pickler_get_persid(PicklerObject *self)
3677{
3678 if (self->pers_func == NULL)
3679 PyErr_SetString(PyExc_AttributeError, "persistent_id");
3680 else
3681 Py_INCREF(self->pers_func);
3682 return self->pers_func;
3683}
3684
3685static int
3686Pickler_set_persid(PicklerObject *self, PyObject *value)
3687{
3688 PyObject *tmp;
3689
3690 if (value == NULL) {
3691 PyErr_SetString(PyExc_TypeError,
3692 "attribute deletion is not supported");
3693 return -1;
3694 }
3695 if (!PyCallable_Check(value)) {
3696 PyErr_SetString(PyExc_TypeError,
3697 "persistent_id must be a callable taking one argument");
3698 return -1;
3699 }
3700
3701 tmp = self->pers_func;
3702 Py_INCREF(value);
3703 self->pers_func = value;
3704 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
3705
3706 return 0;
3707}
3708
3709static PyMemberDef Pickler_members[] = {
3710 {"bin", T_INT, offsetof(PicklerObject, bin)},
3711 {"fast", T_INT, offsetof(PicklerObject, fast)},
3712 {NULL}
3713};
3714
3715static PyGetSetDef Pickler_getsets[] = {
3716 {"memo", (getter)Pickler_get_memo,
3717 (setter)Pickler_set_memo},
3718 {"persistent_id", (getter)Pickler_get_persid,
3719 (setter)Pickler_set_persid},
3720 {NULL}
3721};
3722
3723static PyTypeObject Pickler_Type = {
3724 PyVarObject_HEAD_INIT(NULL, 0)
3725 "_pickle.Pickler" , /*tp_name*/
3726 sizeof(PicklerObject), /*tp_basicsize*/
3727 0, /*tp_itemsize*/
3728 (destructor)Pickler_dealloc, /*tp_dealloc*/
3729 0, /*tp_print*/
3730 0, /*tp_getattr*/
3731 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00003732 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003733 0, /*tp_repr*/
3734 0, /*tp_as_number*/
3735 0, /*tp_as_sequence*/
3736 0, /*tp_as_mapping*/
3737 0, /*tp_hash*/
3738 0, /*tp_call*/
3739 0, /*tp_str*/
3740 0, /*tp_getattro*/
3741 0, /*tp_setattro*/
3742 0, /*tp_as_buffer*/
3743 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3744 Pickler_doc, /*tp_doc*/
3745 (traverseproc)Pickler_traverse, /*tp_traverse*/
3746 (inquiry)Pickler_clear, /*tp_clear*/
3747 0, /*tp_richcompare*/
3748 0, /*tp_weaklistoffset*/
3749 0, /*tp_iter*/
3750 0, /*tp_iternext*/
3751 Pickler_methods, /*tp_methods*/
3752 Pickler_members, /*tp_members*/
3753 Pickler_getsets, /*tp_getset*/
3754 0, /*tp_base*/
3755 0, /*tp_dict*/
3756 0, /*tp_descr_get*/
3757 0, /*tp_descr_set*/
3758 0, /*tp_dictoffset*/
3759 (initproc)Pickler_init, /*tp_init*/
3760 PyType_GenericAlloc, /*tp_alloc*/
3761 PyType_GenericNew, /*tp_new*/
3762 PyObject_GC_Del, /*tp_free*/
3763 0, /*tp_is_gc*/
3764};
3765
3766/* Temporary helper for calling self.find_class().
3767
3768 XXX: It would be nice to able to avoid Python function call overhead, by
3769 using directly the C version of find_class(), when find_class() is not
3770 overridden by a subclass. Although, this could become rather hackish. A
3771 simpler optimization would be to call the C function when self is not a
3772 subclass instance. */
3773static PyObject *
3774find_class(UnpicklerObject *self, PyObject *module_name, PyObject *global_name)
3775{
3776 return PyObject_CallMethod((PyObject *)self, "find_class", "OO",
3777 module_name, global_name);
3778}
3779
3780static int
3781marker(UnpicklerObject *self)
3782{
3783 if (self->num_marks < 1) {
3784 PyErr_SetString(UnpicklingError, "could not find MARK");
3785 return -1;
3786 }
3787
3788 return self->marks[--self->num_marks];
3789}
3790
3791static int
3792load_none(UnpicklerObject *self)
3793{
3794 PDATA_APPEND(self->stack, Py_None, -1);
3795 return 0;
3796}
3797
3798static int
3799bad_readline(void)
3800{
3801 PyErr_SetString(UnpicklingError, "pickle data was truncated");
3802 return -1;
3803}
3804
3805static int
3806load_int(UnpicklerObject *self)
3807{
3808 PyObject *value;
3809 char *endptr, *s;
3810 Py_ssize_t len;
3811 long x;
3812
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003813 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003814 return -1;
3815 if (len < 2)
3816 return bad_readline();
3817
3818 errno = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003819 /* XXX: Should the base argument of strtol() be explicitly set to 10?
3820 XXX(avassalotti): Should this uses PyOS_strtol()? */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003821 x = strtol(s, &endptr, 0);
3822
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003823 if (errno || (*endptr != '\n' && *endptr != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003824 /* Hm, maybe we've got something long. Let's try reading
3825 * it as a Python long object. */
3826 errno = 0;
3827 /* XXX: Same thing about the base here. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003828 value = PyLong_FromString(s, NULL, 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003829 if (value == NULL) {
3830 PyErr_SetString(PyExc_ValueError,
3831 "could not convert string to int");
3832 return -1;
3833 }
3834 }
3835 else {
3836 if (len == 3 && (x == 0 || x == 1)) {
3837 if ((value = PyBool_FromLong(x)) == NULL)
3838 return -1;
3839 }
3840 else {
3841 if ((value = PyLong_FromLong(x)) == NULL)
3842 return -1;
3843 }
3844 }
3845
3846 PDATA_PUSH(self->stack, value, -1);
3847 return 0;
3848}
3849
3850static int
3851load_bool(UnpicklerObject *self, PyObject *boolean)
3852{
3853 assert(boolean == Py_True || boolean == Py_False);
3854 PDATA_APPEND(self->stack, boolean, -1);
3855 return 0;
3856}
3857
3858/* s contains x bytes of a little-endian integer. Return its value as a
3859 * C int. Obscure: when x is 1 or 2, this is an unsigned little-endian
3860 * int, but when x is 4 it's a signed one. This is an historical source
3861 * of x-platform bugs.
3862 */
3863static long
3864calc_binint(char *bytes, int size)
3865{
3866 unsigned char *s = (unsigned char *)bytes;
3867 int i = size;
3868 long x = 0;
3869
3870 for (i = 0; i < size; i++) {
3871 x |= (long)s[i] << (i * 8);
3872 }
3873
3874 /* Unlike BININT1 and BININT2, BININT (more accurately BININT4)
3875 * is signed, so on a box with longs bigger than 4 bytes we need
3876 * to extend a BININT's sign bit to the full width.
3877 */
3878 if (SIZEOF_LONG > 4 && size == 4) {
3879 x |= -(x & (1L << 31));
3880 }
3881
3882 return x;
3883}
3884
3885static int
3886load_binintx(UnpicklerObject *self, char *s, int size)
3887{
3888 PyObject *value;
3889 long x;
3890
3891 x = calc_binint(s, size);
3892
3893 if ((value = PyLong_FromLong(x)) == NULL)
3894 return -1;
3895
3896 PDATA_PUSH(self->stack, value, -1);
3897 return 0;
3898}
3899
3900static int
3901load_binint(UnpicklerObject *self)
3902{
3903 char *s;
3904
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003905 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003906 return -1;
3907
3908 return load_binintx(self, s, 4);
3909}
3910
3911static int
3912load_binint1(UnpicklerObject *self)
3913{
3914 char *s;
3915
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003916 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003917 return -1;
3918
3919 return load_binintx(self, s, 1);
3920}
3921
3922static int
3923load_binint2(UnpicklerObject *self)
3924{
3925 char *s;
3926
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003927 if (_Unpickler_Read(self, &s, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003928 return -1;
3929
3930 return load_binintx(self, s, 2);
3931}
3932
3933static int
3934load_long(UnpicklerObject *self)
3935{
3936 PyObject *value;
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00003937 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003938 Py_ssize_t len;
3939
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003940 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003941 return -1;
3942 if (len < 2)
3943 return bad_readline();
3944
Mark Dickinson8dd05142009-01-20 20:43:58 +00003945 /* s[len-2] will usually be 'L' (and s[len-1] is '\n'); we need to remove
3946 the 'L' before calling PyLong_FromString. In order to maintain
3947 compatibility with Python 3.0.0, we don't actually *require*
3948 the 'L' to be present. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003949 if (s[len-2] == 'L')
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00003950 s[len-2] = '\0';
Alexandre Vassalottie4bccb72009-01-24 01:47:57 +00003951 /* XXX: Should the base argument explicitly set to 10? */
3952 value = PyLong_FromString(s, NULL, 0);
Mark Dickinson8dd05142009-01-20 20:43:58 +00003953 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003954 return -1;
3955
3956 PDATA_PUSH(self->stack, value, -1);
3957 return 0;
3958}
3959
3960/* 'size' bytes contain the # of bytes of little-endian 256's-complement
3961 * data following.
3962 */
3963static int
3964load_counted_long(UnpicklerObject *self, int size)
3965{
3966 PyObject *value;
3967 char *nbytes;
3968 char *pdata;
3969
3970 assert(size == 1 || size == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003971 if (_Unpickler_Read(self, &nbytes, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003972 return -1;
3973
3974 size = calc_binint(nbytes, size);
3975 if (size < 0) {
3976 /* Corrupt or hostile pickle -- we never write one like this */
3977 PyErr_SetString(UnpicklingError,
3978 "LONG pickle has negative byte count");
3979 return -1;
3980 }
3981
3982 if (size == 0)
3983 value = PyLong_FromLong(0L);
3984 else {
3985 /* Read the raw little-endian bytes and convert. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003986 if (_Unpickler_Read(self, &pdata, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003987 return -1;
3988 value = _PyLong_FromByteArray((unsigned char *)pdata, (size_t)size,
3989 1 /* little endian */ , 1 /* signed */ );
3990 }
3991 if (value == NULL)
3992 return -1;
3993 PDATA_PUSH(self->stack, value, -1);
3994 return 0;
3995}
3996
3997static int
3998load_float(UnpicklerObject *self)
3999{
4000 PyObject *value;
4001 char *endptr, *s;
4002 Py_ssize_t len;
4003 double d;
4004
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004005 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004006 return -1;
4007 if (len < 2)
4008 return bad_readline();
4009
4010 errno = 0;
Mark Dickinson725bfd82009-05-03 20:33:40 +00004011 d = PyOS_string_to_double(s, &endptr, PyExc_OverflowError);
4012 if (d == -1.0 && PyErr_Occurred())
4013 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004014 if ((endptr[0] != '\n') && (endptr[0] != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004015 PyErr_SetString(PyExc_ValueError, "could not convert string to float");
4016 return -1;
4017 }
Mark Dickinson725bfd82009-05-03 20:33:40 +00004018 value = PyFloat_FromDouble(d);
4019 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004020 return -1;
4021
4022 PDATA_PUSH(self->stack, value, -1);
4023 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004024}
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004025
4026static int
4027load_binfloat(UnpicklerObject *self)
4028{
4029 PyObject *value;
4030 double x;
4031 char *s;
4032
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004033 if (_Unpickler_Read(self, &s, 8) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004034 return -1;
4035
4036 x = _PyFloat_Unpack8((unsigned char *)s, 0);
4037 if (x == -1.0 && PyErr_Occurred())
4038 return -1;
4039
4040 if ((value = PyFloat_FromDouble(x)) == NULL)
4041 return -1;
4042
4043 PDATA_PUSH(self->stack, value, -1);
4044 return 0;
4045}
4046
4047static int
4048load_string(UnpicklerObject *self)
4049{
4050 PyObject *bytes;
4051 PyObject *str = NULL;
4052 Py_ssize_t len;
4053 char *s, *p;
4054
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004055 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004056 return -1;
4057 if (len < 3)
4058 return bad_readline();
4059 if ((s = strdup(s)) == NULL) {
4060 PyErr_NoMemory();
4061 return -1;
4062 }
4063
4064 /* Strip outermost quotes */
4065 while (s[len - 1] <= ' ')
4066 len--;
4067 if (s[0] == '"' && s[len - 1] == '"') {
4068 s[len - 1] = '\0';
4069 p = s + 1;
4070 len -= 2;
4071 }
4072 else if (s[0] == '\'' && s[len - 1] == '\'') {
4073 s[len - 1] = '\0';
4074 p = s + 1;
4075 len -= 2;
4076 }
4077 else {
4078 free(s);
4079 PyErr_SetString(PyExc_ValueError, "insecure string pickle");
4080 return -1;
4081 }
4082
4083 /* Use the PyBytes API to decode the string, since that is what is used
4084 to encode, and then coerce the result to Unicode. */
4085 bytes = PyBytes_DecodeEscape(p, len, NULL, 0, NULL);
4086 free(s);
4087 if (bytes == NULL)
4088 return -1;
4089 str = PyUnicode_FromEncodedObject(bytes, self->encoding, self->errors);
4090 Py_DECREF(bytes);
4091 if (str == NULL)
4092 return -1;
4093
4094 PDATA_PUSH(self->stack, str, -1);
4095 return 0;
4096}
4097
4098static int
4099load_binbytes(UnpicklerObject *self)
4100{
4101 PyObject *bytes;
4102 long x;
4103 char *s;
4104
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004105 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004106 return -1;
4107
4108 x = calc_binint(s, 4);
4109 if (x < 0) {
4110 PyErr_SetString(UnpicklingError,
4111 "BINBYTES pickle has negative byte count");
4112 return -1;
4113 }
4114
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004115 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004116 return -1;
4117 bytes = PyBytes_FromStringAndSize(s, x);
4118 if (bytes == NULL)
4119 return -1;
4120
4121 PDATA_PUSH(self->stack, bytes, -1);
4122 return 0;
4123}
4124
4125static int
4126load_short_binbytes(UnpicklerObject *self)
4127{
4128 PyObject *bytes;
4129 unsigned char x;
4130 char *s;
4131
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004132 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004133 return -1;
4134
4135 x = (unsigned char)s[0];
4136
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004137 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004138 return -1;
4139
4140 bytes = PyBytes_FromStringAndSize(s, x);
4141 if (bytes == NULL)
4142 return -1;
4143
4144 PDATA_PUSH(self->stack, bytes, -1);
4145 return 0;
4146}
4147
4148static int
4149load_binstring(UnpicklerObject *self)
4150{
4151 PyObject *str;
4152 long x;
4153 char *s;
4154
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004155 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004156 return -1;
4157
4158 x = calc_binint(s, 4);
4159 if (x < 0) {
4160 PyErr_SetString(UnpicklingError,
4161 "BINSTRING pickle has negative byte count");
4162 return -1;
4163 }
4164
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004165 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004166 return -1;
4167
4168 /* Convert Python 2.x strings to unicode. */
4169 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4170 if (str == NULL)
4171 return -1;
4172
4173 PDATA_PUSH(self->stack, str, -1);
4174 return 0;
4175}
4176
4177static int
4178load_short_binstring(UnpicklerObject *self)
4179{
4180 PyObject *str;
4181 unsigned char x;
4182 char *s;
4183
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004184 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004185 return -1;
4186
4187 x = (unsigned char)s[0];
4188
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004189 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004190 return -1;
4191
4192 /* Convert Python 2.x strings to unicode. */
4193 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4194 if (str == NULL)
4195 return -1;
4196
4197 PDATA_PUSH(self->stack, str, -1);
4198 return 0;
4199}
4200
4201static int
4202load_unicode(UnpicklerObject *self)
4203{
4204 PyObject *str;
4205 Py_ssize_t len;
4206 char *s;
4207
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004208 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004209 return -1;
4210 if (len < 1)
4211 return bad_readline();
4212
4213 str = PyUnicode_DecodeRawUnicodeEscape(s, len - 1, NULL);
4214 if (str == NULL)
4215 return -1;
4216
4217 PDATA_PUSH(self->stack, str, -1);
4218 return 0;
4219}
4220
4221static int
4222load_binunicode(UnpicklerObject *self)
4223{
4224 PyObject *str;
4225 long size;
4226 char *s;
4227
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004228 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004229 return -1;
4230
4231 size = calc_binint(s, 4);
4232 if (size < 0) {
4233 PyErr_SetString(UnpicklingError,
4234 "BINUNICODE pickle has negative byte count");
4235 return -1;
4236 }
4237
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004238 if (_Unpickler_Read(self, &s, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004239 return -1;
4240
Victor Stinner485fb562010-04-13 11:07:24 +00004241 str = PyUnicode_DecodeUTF8(s, size, "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004242 if (str == NULL)
4243 return -1;
4244
4245 PDATA_PUSH(self->stack, str, -1);
4246 return 0;
4247}
4248
4249static int
4250load_tuple(UnpicklerObject *self)
4251{
4252 PyObject *tuple;
4253 int i;
4254
4255 if ((i = marker(self)) < 0)
4256 return -1;
4257
4258 tuple = Pdata_poptuple(self->stack, i);
4259 if (tuple == NULL)
4260 return -1;
4261 PDATA_PUSH(self->stack, tuple, -1);
4262 return 0;
4263}
4264
4265static int
4266load_counted_tuple(UnpicklerObject *self, int len)
4267{
4268 PyObject *tuple;
4269
4270 tuple = PyTuple_New(len);
4271 if (tuple == NULL)
4272 return -1;
4273
4274 while (--len >= 0) {
4275 PyObject *item;
4276
4277 PDATA_POP(self->stack, item);
4278 if (item == NULL)
4279 return -1;
4280 PyTuple_SET_ITEM(tuple, len, item);
4281 }
4282 PDATA_PUSH(self->stack, tuple, -1);
4283 return 0;
4284}
4285
4286static int
4287load_empty_list(UnpicklerObject *self)
4288{
4289 PyObject *list;
4290
4291 if ((list = PyList_New(0)) == NULL)
4292 return -1;
4293 PDATA_PUSH(self->stack, list, -1);
4294 return 0;
4295}
4296
4297static int
4298load_empty_dict(UnpicklerObject *self)
4299{
4300 PyObject *dict;
4301
4302 if ((dict = PyDict_New()) == NULL)
4303 return -1;
4304 PDATA_PUSH(self->stack, dict, -1);
4305 return 0;
4306}
4307
4308static int
4309load_list(UnpicklerObject *self)
4310{
4311 PyObject *list;
4312 int i;
4313
4314 if ((i = marker(self)) < 0)
4315 return -1;
4316
4317 list = Pdata_poplist(self->stack, i);
4318 if (list == NULL)
4319 return -1;
4320 PDATA_PUSH(self->stack, list, -1);
4321 return 0;
4322}
4323
4324static int
4325load_dict(UnpicklerObject *self)
4326{
4327 PyObject *dict, *key, *value;
4328 int i, j, k;
4329
4330 if ((i = marker(self)) < 0)
4331 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004332 j = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004333
4334 if ((dict = PyDict_New()) == NULL)
4335 return -1;
4336
4337 for (k = i + 1; k < j; k += 2) {
4338 key = self->stack->data[k - 1];
4339 value = self->stack->data[k];
4340 if (PyDict_SetItem(dict, key, value) < 0) {
4341 Py_DECREF(dict);
4342 return -1;
4343 }
4344 }
4345 Pdata_clear(self->stack, i);
4346 PDATA_PUSH(self->stack, dict, -1);
4347 return 0;
4348}
4349
4350static PyObject *
4351instantiate(PyObject *cls, PyObject *args)
4352{
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004353 PyObject *result = NULL;
4354 /* Caller must assure args are a tuple. Normally, args come from
4355 Pdata_poptuple which packs objects from the top of the stack
4356 into a newly created tuple. */
4357 assert(PyTuple_Check(args));
4358 if (Py_SIZE(args) > 0 || !PyType_Check(cls) ||
4359 PyObject_HasAttrString(cls, "__getinitargs__")) {
4360 result = PyObject_CallObject(cls, args);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004361 }
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004362 else {
4363 result = PyObject_CallMethod(cls, "__new__", "O", cls);
4364 }
4365 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004366}
4367
4368static int
4369load_obj(UnpicklerObject *self)
4370{
4371 PyObject *cls, *args, *obj = NULL;
4372 int i;
4373
4374 if ((i = marker(self)) < 0)
4375 return -1;
4376
4377 args = Pdata_poptuple(self->stack, i + 1);
4378 if (args == NULL)
4379 return -1;
4380
4381 PDATA_POP(self->stack, cls);
4382 if (cls) {
4383 obj = instantiate(cls, args);
4384 Py_DECREF(cls);
4385 }
4386 Py_DECREF(args);
4387 if (obj == NULL)
4388 return -1;
4389
4390 PDATA_PUSH(self->stack, obj, -1);
4391 return 0;
4392}
4393
4394static int
4395load_inst(UnpicklerObject *self)
4396{
4397 PyObject *cls = NULL;
4398 PyObject *args = NULL;
4399 PyObject *obj = NULL;
4400 PyObject *module_name;
4401 PyObject *class_name;
4402 Py_ssize_t len;
4403 int i;
4404 char *s;
4405
4406 if ((i = marker(self)) < 0)
4407 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004408 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004409 return -1;
4410 if (len < 2)
4411 return bad_readline();
4412
4413 /* Here it is safe to use PyUnicode_DecodeASCII(), even though non-ASCII
4414 identifiers are permitted in Python 3.0, since the INST opcode is only
4415 supported by older protocols on Python 2.x. */
4416 module_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
4417 if (module_name == NULL)
4418 return -1;
4419
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004420 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004421 if (len < 2)
4422 return bad_readline();
4423 class_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004424 if (class_name != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004425 cls = find_class(self, module_name, class_name);
4426 Py_DECREF(class_name);
4427 }
4428 }
4429 Py_DECREF(module_name);
4430
4431 if (cls == NULL)
4432 return -1;
4433
4434 if ((args = Pdata_poptuple(self->stack, i)) != NULL) {
4435 obj = instantiate(cls, args);
4436 Py_DECREF(args);
4437 }
4438 Py_DECREF(cls);
4439
4440 if (obj == NULL)
4441 return -1;
4442
4443 PDATA_PUSH(self->stack, obj, -1);
4444 return 0;
4445}
4446
4447static int
4448load_newobj(UnpicklerObject *self)
4449{
4450 PyObject *args = NULL;
4451 PyObject *clsraw = NULL;
4452 PyTypeObject *cls; /* clsraw cast to its true type */
4453 PyObject *obj;
4454
4455 /* Stack is ... cls argtuple, and we want to call
4456 * cls.__new__(cls, *argtuple).
4457 */
4458 PDATA_POP(self->stack, args);
4459 if (args == NULL)
4460 goto error;
4461 if (!PyTuple_Check(args)) {
4462 PyErr_SetString(UnpicklingError, "NEWOBJ expected an arg " "tuple.");
4463 goto error;
4464 }
4465
4466 PDATA_POP(self->stack, clsraw);
4467 cls = (PyTypeObject *)clsraw;
4468 if (cls == NULL)
4469 goto error;
4470 if (!PyType_Check(cls)) {
4471 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4472 "isn't a type object");
4473 goto error;
4474 }
4475 if (cls->tp_new == NULL) {
4476 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4477 "has NULL tp_new");
4478 goto error;
4479 }
4480
4481 /* Call __new__. */
4482 obj = cls->tp_new(cls, args, NULL);
4483 if (obj == NULL)
4484 goto error;
4485
4486 Py_DECREF(args);
4487 Py_DECREF(clsraw);
4488 PDATA_PUSH(self->stack, obj, -1);
4489 return 0;
4490
4491 error:
4492 Py_XDECREF(args);
4493 Py_XDECREF(clsraw);
4494 return -1;
4495}
4496
4497static int
4498load_global(UnpicklerObject *self)
4499{
4500 PyObject *global = NULL;
4501 PyObject *module_name;
4502 PyObject *global_name;
4503 Py_ssize_t len;
4504 char *s;
4505
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004506 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004507 return -1;
4508 if (len < 2)
4509 return bad_readline();
4510 module_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4511 if (!module_name)
4512 return -1;
4513
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004514 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004515 if (len < 2) {
4516 Py_DECREF(module_name);
4517 return bad_readline();
4518 }
4519 global_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4520 if (global_name) {
4521 global = find_class(self, module_name, global_name);
4522 Py_DECREF(global_name);
4523 }
4524 }
4525 Py_DECREF(module_name);
4526
4527 if (global == NULL)
4528 return -1;
4529 PDATA_PUSH(self->stack, global, -1);
4530 return 0;
4531}
4532
4533static int
4534load_persid(UnpicklerObject *self)
4535{
4536 PyObject *pid;
4537 Py_ssize_t len;
4538 char *s;
4539
4540 if (self->pers_func) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004541 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004542 return -1;
4543 if (len < 2)
4544 return bad_readline();
4545
4546 pid = PyBytes_FromStringAndSize(s, len - 1);
4547 if (pid == NULL)
4548 return -1;
4549
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004550 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004551 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004552 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004553 if (pid == NULL)
4554 return -1;
4555
4556 PDATA_PUSH(self->stack, pid, -1);
4557 return 0;
4558 }
4559 else {
4560 PyErr_SetString(UnpicklingError,
4561 "A load persistent id instruction was encountered,\n"
4562 "but no persistent_load function was specified.");
4563 return -1;
4564 }
4565}
4566
4567static int
4568load_binpersid(UnpicklerObject *self)
4569{
4570 PyObject *pid;
4571
4572 if (self->pers_func) {
4573 PDATA_POP(self->stack, pid);
4574 if (pid == NULL)
4575 return -1;
4576
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004577 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004578 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004579 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004580 if (pid == NULL)
4581 return -1;
4582
4583 PDATA_PUSH(self->stack, pid, -1);
4584 return 0;
4585 }
4586 else {
4587 PyErr_SetString(UnpicklingError,
4588 "A load persistent id instruction was encountered,\n"
4589 "but no persistent_load function was specified.");
4590 return -1;
4591 }
4592}
4593
4594static int
4595load_pop(UnpicklerObject *self)
4596{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004597 int len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004598
4599 /* Note that we split the (pickle.py) stack into two stacks,
4600 * an object stack and a mark stack. We have to be clever and
4601 * pop the right one. We do this by looking at the top of the
Collin Winter8ca69de2009-05-26 16:53:41 +00004602 * mark stack first, and only signalling a stack underflow if
4603 * the object stack is empty and the mark stack doesn't match
4604 * our expectations.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004605 */
Collin Winter8ca69de2009-05-26 16:53:41 +00004606 if (self->num_marks > 0 && self->marks[self->num_marks - 1] == len) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004607 self->num_marks--;
Antoine Pitrou01a15ea2010-01-07 17:57:31 +00004608 } else if (len > 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004609 len--;
4610 Py_DECREF(self->stack->data[len]);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004611 Py_SIZE(self->stack) = len;
Collin Winter8ca69de2009-05-26 16:53:41 +00004612 } else {
4613 return stack_underflow();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004614 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004615 return 0;
4616}
4617
4618static int
4619load_pop_mark(UnpicklerObject *self)
4620{
4621 int i;
4622
4623 if ((i = marker(self)) < 0)
4624 return -1;
4625
4626 Pdata_clear(self->stack, i);
4627
4628 return 0;
4629}
4630
4631static int
4632load_dup(UnpicklerObject *self)
4633{
4634 PyObject *last;
4635 int len;
4636
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004637 if ((len = Py_SIZE(self->stack)) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004638 return stack_underflow();
4639 last = self->stack->data[len - 1];
4640 PDATA_APPEND(self->stack, last, -1);
4641 return 0;
4642}
4643
4644static int
4645load_get(UnpicklerObject *self)
4646{
4647 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004648 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004649 Py_ssize_t len;
4650 char *s;
4651
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004652 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004653 return -1;
4654 if (len < 2)
4655 return bad_readline();
4656
4657 key = PyLong_FromString(s, NULL, 10);
4658 if (key == NULL)
4659 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004660 idx = PyLong_AsSsize_t(key);
4661 if (idx == -1 && PyErr_Occurred()) {
4662 Py_DECREF(key);
4663 return -1;
4664 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004665
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004666 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004667 if (value == NULL) {
4668 if (!PyErr_Occurred())
4669 PyErr_SetObject(PyExc_KeyError, key);
4670 Py_DECREF(key);
4671 return -1;
4672 }
4673 Py_DECREF(key);
4674
4675 PDATA_APPEND(self->stack, value, -1);
4676 return 0;
4677}
4678
4679static int
4680load_binget(UnpicklerObject *self)
4681{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004682 PyObject *value;
4683 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004684 char *s;
4685
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004686 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004687 return -1;
4688
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004689 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004690
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004691 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004692 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004693 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004694 if (!PyErr_Occurred())
4695 PyErr_SetObject(PyExc_KeyError, key);
4696 Py_DECREF(key);
4697 return -1;
4698 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004699
4700 PDATA_APPEND(self->stack, value, -1);
4701 return 0;
4702}
4703
4704static int
4705load_long_binget(UnpicklerObject *self)
4706{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004707 PyObject *value;
4708 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004709 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004710
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004711 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004712 return -1;
4713
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004714 idx = (long)Py_CHARMASK(s[0]);
4715 idx |= (long)Py_CHARMASK(s[1]) << 8;
4716 idx |= (long)Py_CHARMASK(s[2]) << 16;
4717 idx |= (long)Py_CHARMASK(s[3]) << 24;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004718
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004719 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004720 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004721 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004722 if (!PyErr_Occurred())
4723 PyErr_SetObject(PyExc_KeyError, key);
4724 Py_DECREF(key);
4725 return -1;
4726 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004727
4728 PDATA_APPEND(self->stack, value, -1);
4729 return 0;
4730}
4731
4732/* Push an object from the extension registry (EXT[124]). nbytes is
4733 * the number of bytes following the opcode, holding the index (code) value.
4734 */
4735static int
4736load_extension(UnpicklerObject *self, int nbytes)
4737{
4738 char *codebytes; /* the nbytes bytes after the opcode */
4739 long code; /* calc_binint returns long */
4740 PyObject *py_code; /* code as a Python int */
4741 PyObject *obj; /* the object to push */
4742 PyObject *pair; /* (module_name, class_name) */
4743 PyObject *module_name, *class_name;
4744
4745 assert(nbytes == 1 || nbytes == 2 || nbytes == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004746 if (_Unpickler_Read(self, &codebytes, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004747 return -1;
4748 code = calc_binint(codebytes, nbytes);
4749 if (code <= 0) { /* note that 0 is forbidden */
4750 /* Corrupt or hostile pickle. */
4751 PyErr_SetString(UnpicklingError, "EXT specifies code <= 0");
4752 return -1;
4753 }
4754
4755 /* Look for the code in the cache. */
4756 py_code = PyLong_FromLong(code);
4757 if (py_code == NULL)
4758 return -1;
4759 obj = PyDict_GetItem(extension_cache, py_code);
4760 if (obj != NULL) {
4761 /* Bingo. */
4762 Py_DECREF(py_code);
4763 PDATA_APPEND(self->stack, obj, -1);
4764 return 0;
4765 }
4766
4767 /* Look up the (module_name, class_name) pair. */
4768 pair = PyDict_GetItem(inverted_registry, py_code);
4769 if (pair == NULL) {
4770 Py_DECREF(py_code);
4771 PyErr_Format(PyExc_ValueError, "unregistered extension "
4772 "code %ld", code);
4773 return -1;
4774 }
4775 /* Since the extension registry is manipulable via Python code,
4776 * confirm that pair is really a 2-tuple of strings.
4777 */
4778 if (!PyTuple_Check(pair) || PyTuple_Size(pair) != 2 ||
4779 !PyUnicode_Check(module_name = PyTuple_GET_ITEM(pair, 0)) ||
4780 !PyUnicode_Check(class_name = PyTuple_GET_ITEM(pair, 1))) {
4781 Py_DECREF(py_code);
4782 PyErr_Format(PyExc_ValueError, "_inverted_registry[%ld] "
4783 "isn't a 2-tuple of strings", code);
4784 return -1;
4785 }
4786 /* Load the object. */
4787 obj = find_class(self, module_name, class_name);
4788 if (obj == NULL) {
4789 Py_DECREF(py_code);
4790 return -1;
4791 }
4792 /* Cache code -> obj. */
4793 code = PyDict_SetItem(extension_cache, py_code, obj);
4794 Py_DECREF(py_code);
4795 if (code < 0) {
4796 Py_DECREF(obj);
4797 return -1;
4798 }
4799 PDATA_PUSH(self->stack, obj, -1);
4800 return 0;
4801}
4802
4803static int
4804load_put(UnpicklerObject *self)
4805{
4806 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004807 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004808 Py_ssize_t len;
4809 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004810
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004811 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004812 return -1;
4813 if (len < 2)
4814 return bad_readline();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004815 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004816 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004817 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004818
4819 key = PyLong_FromString(s, NULL, 10);
4820 if (key == NULL)
4821 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004822 idx = PyLong_AsSsize_t(key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004823 Py_DECREF(key);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004824 if (idx == -1 && PyErr_Occurred())
4825 return -1;
4826
4827 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004828}
4829
4830static int
4831load_binput(UnpicklerObject *self)
4832{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004833 PyObject *value;
4834 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004835 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004836
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004837 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004838 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004839
4840 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004841 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004842 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004843
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004844 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004845
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004846 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004847}
4848
4849static int
4850load_long_binput(UnpicklerObject *self)
4851{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004852 PyObject *value;
4853 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004854 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004855
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004856 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004857 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004858
4859 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004860 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004861 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004862
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004863 idx = (long)Py_CHARMASK(s[0]);
4864 idx |= (long)Py_CHARMASK(s[1]) << 8;
4865 idx |= (long)Py_CHARMASK(s[2]) << 16;
4866 idx |= (long)Py_CHARMASK(s[3]) << 24;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004867
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004868 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004869}
4870
4871static int
4872do_append(UnpicklerObject *self, int x)
4873{
4874 PyObject *value;
4875 PyObject *list;
4876 int len, i;
4877
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004878 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004879 if (x > len || x <= 0)
4880 return stack_underflow();
4881 if (len == x) /* nothing to do */
4882 return 0;
4883
4884 list = self->stack->data[x - 1];
4885
4886 if (PyList_Check(list)) {
4887 PyObject *slice;
4888 Py_ssize_t list_len;
4889
4890 slice = Pdata_poplist(self->stack, x);
4891 if (!slice)
4892 return -1;
4893 list_len = PyList_GET_SIZE(list);
4894 i = PyList_SetSlice(list, list_len, list_len, slice);
4895 Py_DECREF(slice);
4896 return i;
4897 }
4898 else {
4899 PyObject *append_func;
4900
4901 append_func = PyObject_GetAttrString(list, "append");
4902 if (append_func == NULL)
4903 return -1;
4904 for (i = x; i < len; i++) {
4905 PyObject *result;
4906
4907 value = self->stack->data[i];
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004908 result = _Unpickler_FastCall(self, append_func, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004909 if (result == NULL) {
4910 Pdata_clear(self->stack, i + 1);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004911 Py_SIZE(self->stack) = x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004912 return -1;
4913 }
4914 Py_DECREF(result);
4915 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004916 Py_SIZE(self->stack) = x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004917 }
4918
4919 return 0;
4920}
4921
4922static int
4923load_append(UnpicklerObject *self)
4924{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004925 return do_append(self, Py_SIZE(self->stack) - 1);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004926}
4927
4928static int
4929load_appends(UnpicklerObject *self)
4930{
4931 return do_append(self, marker(self));
4932}
4933
4934static int
4935do_setitems(UnpicklerObject *self, int x)
4936{
4937 PyObject *value, *key;
4938 PyObject *dict;
4939 int len, i;
4940 int status = 0;
4941
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004942 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004943 if (x > len || x <= 0)
4944 return stack_underflow();
4945 if (len == x) /* nothing to do */
4946 return 0;
4947 if ((len - x) % 2 != 0) {
4948 /* Currupt or hostile pickle -- we never write one like this. */
4949 PyErr_SetString(UnpicklingError, "odd number of items for SETITEMS");
4950 return -1;
4951 }
4952
4953 /* Here, dict does not actually need to be a PyDict; it could be anything
4954 that supports the __setitem__ attribute. */
4955 dict = self->stack->data[x - 1];
4956
4957 for (i = x + 1; i < len; i += 2) {
4958 key = self->stack->data[i - 1];
4959 value = self->stack->data[i];
4960 if (PyObject_SetItem(dict, key, value) < 0) {
4961 status = -1;
4962 break;
4963 }
4964 }
4965
4966 Pdata_clear(self->stack, x);
4967 return status;
4968}
4969
4970static int
4971load_setitem(UnpicklerObject *self)
4972{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004973 return do_setitems(self, Py_SIZE(self->stack) - 2);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004974}
4975
4976static int
4977load_setitems(UnpicklerObject *self)
4978{
4979 return do_setitems(self, marker(self));
4980}
4981
4982static int
4983load_build(UnpicklerObject *self)
4984{
4985 PyObject *state, *inst, *slotstate;
4986 PyObject *setstate;
4987 int status = 0;
4988
4989 /* Stack is ... instance, state. We want to leave instance at
4990 * the stack top, possibly mutated via instance.__setstate__(state).
4991 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004992 if (Py_SIZE(self->stack) < 2)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004993 return stack_underflow();
4994
4995 PDATA_POP(self->stack, state);
4996 if (state == NULL)
4997 return -1;
4998
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004999 inst = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005000
5001 setstate = PyObject_GetAttrString(inst, "__setstate__");
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005002 if (setstate == NULL) {
5003 if (PyErr_ExceptionMatches(PyExc_AttributeError))
5004 PyErr_Clear();
Antoine Pitroud79dc622008-09-05 00:03:33 +00005005 else {
5006 Py_DECREF(state);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005007 return -1;
Antoine Pitroud79dc622008-09-05 00:03:33 +00005008 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005009 }
5010 else {
5011 PyObject *result;
5012
5013 /* The explicit __setstate__ is responsible for everything. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005014 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Antoine Pitroud79dc622008-09-05 00:03:33 +00005015 reference to state first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005016 result = _Unpickler_FastCall(self, setstate, state);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005017 Py_DECREF(setstate);
5018 if (result == NULL)
5019 return -1;
5020 Py_DECREF(result);
5021 return 0;
5022 }
5023
5024 /* A default __setstate__. First see whether state embeds a
5025 * slot state dict too (a proto 2 addition).
5026 */
5027 if (PyTuple_Check(state) && Py_SIZE(state) == 2) {
5028 PyObject *tmp = state;
5029
5030 state = PyTuple_GET_ITEM(tmp, 0);
5031 slotstate = PyTuple_GET_ITEM(tmp, 1);
5032 Py_INCREF(state);
5033 Py_INCREF(slotstate);
5034 Py_DECREF(tmp);
5035 }
5036 else
5037 slotstate = NULL;
5038
5039 /* Set inst.__dict__ from the state dict (if any). */
5040 if (state != Py_None) {
5041 PyObject *dict;
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005042 PyObject *d_key, *d_value;
5043 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005044
5045 if (!PyDict_Check(state)) {
5046 PyErr_SetString(UnpicklingError, "state is not a dictionary");
5047 goto error;
5048 }
5049 dict = PyObject_GetAttrString(inst, "__dict__");
5050 if (dict == NULL)
5051 goto error;
5052
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005053 i = 0;
5054 while (PyDict_Next(state, &i, &d_key, &d_value)) {
5055 /* normally the keys for instance attributes are
5056 interned. we should try to do that here. */
5057 Py_INCREF(d_key);
5058 if (PyUnicode_CheckExact(d_key))
5059 PyUnicode_InternInPlace(&d_key);
5060 if (PyObject_SetItem(dict, d_key, d_value) < 0) {
5061 Py_DECREF(d_key);
5062 goto error;
5063 }
5064 Py_DECREF(d_key);
5065 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005066 Py_DECREF(dict);
5067 }
5068
5069 /* Also set instance attributes from the slotstate dict (if any). */
5070 if (slotstate != NULL) {
5071 PyObject *d_key, *d_value;
5072 Py_ssize_t i;
5073
5074 if (!PyDict_Check(slotstate)) {
5075 PyErr_SetString(UnpicklingError,
5076 "slot state is not a dictionary");
5077 goto error;
5078 }
5079 i = 0;
5080 while (PyDict_Next(slotstate, &i, &d_key, &d_value)) {
5081 if (PyObject_SetAttr(inst, d_key, d_value) < 0)
5082 goto error;
5083 }
5084 }
5085
5086 if (0) {
5087 error:
5088 status = -1;
5089 }
5090
5091 Py_DECREF(state);
5092 Py_XDECREF(slotstate);
5093 return status;
5094}
5095
5096static int
5097load_mark(UnpicklerObject *self)
5098{
5099
5100 /* Note that we split the (pickle.py) stack into two stacks, an
5101 * object stack and a mark stack. Here we push a mark onto the
5102 * mark stack.
5103 */
5104
5105 if ((self->num_marks + 1) >= self->marks_size) {
5106 size_t alloc;
5107 int *marks;
5108
5109 /* Use the size_t type to check for overflow. */
5110 alloc = ((size_t)self->num_marks << 1) + 20;
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00005111 if (alloc > PY_SSIZE_T_MAX ||
5112 alloc <= ((size_t)self->num_marks + 1)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005113 PyErr_NoMemory();
5114 return -1;
5115 }
5116
5117 if (self->marks == NULL)
5118 marks = (int *)PyMem_Malloc(alloc * sizeof(int));
5119 else
5120 marks = (int *)PyMem_Realloc(self->marks, alloc * sizeof(int));
5121 if (marks == NULL) {
5122 PyErr_NoMemory();
5123 return -1;
5124 }
5125 self->marks = marks;
5126 self->marks_size = (Py_ssize_t)alloc;
5127 }
5128
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005129 self->marks[self->num_marks++] = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005130
5131 return 0;
5132}
5133
5134static int
5135load_reduce(UnpicklerObject *self)
5136{
5137 PyObject *callable = NULL;
5138 PyObject *argtup = NULL;
5139 PyObject *obj = NULL;
5140
5141 PDATA_POP(self->stack, argtup);
5142 if (argtup == NULL)
5143 return -1;
5144 PDATA_POP(self->stack, callable);
5145 if (callable) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005146 obj = PyObject_CallObject(callable, argtup);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005147 Py_DECREF(callable);
5148 }
5149 Py_DECREF(argtup);
5150
5151 if (obj == NULL)
5152 return -1;
5153
5154 PDATA_PUSH(self->stack, obj, -1);
5155 return 0;
5156}
5157
5158/* Just raises an error if we don't know the protocol specified. PROTO
5159 * is the first opcode for protocols >= 2.
5160 */
5161static int
5162load_proto(UnpicklerObject *self)
5163{
5164 char *s;
5165 int i;
5166
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005167 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005168 return -1;
5169
5170 i = (unsigned char)s[0];
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005171 if (i <= HIGHEST_PROTOCOL) {
5172 self->proto = i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005173 return 0;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005174 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005175
5176 PyErr_Format(PyExc_ValueError, "unsupported pickle protocol: %d", i);
5177 return -1;
5178}
5179
5180static PyObject *
5181load(UnpicklerObject *self)
5182{
5183 PyObject *err;
5184 PyObject *value = NULL;
5185 char *s;
5186
5187 self->num_marks = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005188 if (Py_SIZE(self->stack))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005189 Pdata_clear(self->stack, 0);
5190
5191 /* Convenient macros for the dispatch while-switch loop just below. */
5192#define OP(opcode, load_func) \
5193 case opcode: if (load_func(self) < 0) break; continue;
5194
5195#define OP_ARG(opcode, load_func, arg) \
5196 case opcode: if (load_func(self, (arg)) < 0) break; continue;
5197
5198 while (1) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005199 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005200 break;
5201
5202 switch ((enum opcode)s[0]) {
5203 OP(NONE, load_none)
5204 OP(BININT, load_binint)
5205 OP(BININT1, load_binint1)
5206 OP(BININT2, load_binint2)
5207 OP(INT, load_int)
5208 OP(LONG, load_long)
5209 OP_ARG(LONG1, load_counted_long, 1)
5210 OP_ARG(LONG4, load_counted_long, 4)
5211 OP(FLOAT, load_float)
5212 OP(BINFLOAT, load_binfloat)
5213 OP(BINBYTES, load_binbytes)
5214 OP(SHORT_BINBYTES, load_short_binbytes)
5215 OP(BINSTRING, load_binstring)
5216 OP(SHORT_BINSTRING, load_short_binstring)
5217 OP(STRING, load_string)
5218 OP(UNICODE, load_unicode)
5219 OP(BINUNICODE, load_binunicode)
5220 OP_ARG(EMPTY_TUPLE, load_counted_tuple, 0)
5221 OP_ARG(TUPLE1, load_counted_tuple, 1)
5222 OP_ARG(TUPLE2, load_counted_tuple, 2)
5223 OP_ARG(TUPLE3, load_counted_tuple, 3)
5224 OP(TUPLE, load_tuple)
5225 OP(EMPTY_LIST, load_empty_list)
5226 OP(LIST, load_list)
5227 OP(EMPTY_DICT, load_empty_dict)
5228 OP(DICT, load_dict)
5229 OP(OBJ, load_obj)
5230 OP(INST, load_inst)
5231 OP(NEWOBJ, load_newobj)
5232 OP(GLOBAL, load_global)
5233 OP(APPEND, load_append)
5234 OP(APPENDS, load_appends)
5235 OP(BUILD, load_build)
5236 OP(DUP, load_dup)
5237 OP(BINGET, load_binget)
5238 OP(LONG_BINGET, load_long_binget)
5239 OP(GET, load_get)
5240 OP(MARK, load_mark)
5241 OP(BINPUT, load_binput)
5242 OP(LONG_BINPUT, load_long_binput)
5243 OP(PUT, load_put)
5244 OP(POP, load_pop)
5245 OP(POP_MARK, load_pop_mark)
5246 OP(SETITEM, load_setitem)
5247 OP(SETITEMS, load_setitems)
5248 OP(PERSID, load_persid)
5249 OP(BINPERSID, load_binpersid)
5250 OP(REDUCE, load_reduce)
5251 OP(PROTO, load_proto)
5252 OP_ARG(EXT1, load_extension, 1)
5253 OP_ARG(EXT2, load_extension, 2)
5254 OP_ARG(EXT4, load_extension, 4)
5255 OP_ARG(NEWTRUE, load_bool, Py_True)
5256 OP_ARG(NEWFALSE, load_bool, Py_False)
5257
5258 case STOP:
5259 break;
5260
5261 case '\0':
5262 PyErr_SetNone(PyExc_EOFError);
5263 return NULL;
5264
5265 default:
5266 PyErr_Format(UnpicklingError,
5267 "invalid load key, '%c'.", s[0]);
5268 return NULL;
5269 }
5270
5271 break; /* and we are done! */
5272 }
5273
Antoine Pitrou04248a82010-10-12 20:51:21 +00005274 if (_Unpickler_SkipConsumed(self) < 0)
5275 return NULL;
5276
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005277 /* XXX: It is not clear what this is actually for. */
5278 if ((err = PyErr_Occurred())) {
5279 if (err == PyExc_EOFError) {
5280 PyErr_SetNone(PyExc_EOFError);
5281 }
5282 return NULL;
5283 }
5284
5285 PDATA_POP(self->stack, value);
5286 return value;
5287}
5288
5289PyDoc_STRVAR(Unpickler_load_doc,
5290"load() -> object. Load a pickle."
5291"\n"
5292"Read a pickled object representation from the open file object given in\n"
5293"the constructor, and return the reconstituted object hierarchy specified\n"
5294"therein.\n");
5295
5296static PyObject *
5297Unpickler_load(UnpicklerObject *self)
5298{
5299 /* Check whether the Unpickler was initialized correctly. This prevents
5300 segfaulting if a subclass overridden __init__ with a function that does
5301 not call Unpickler.__init__(). Here, we simply ensure that self->read
5302 is not NULL. */
5303 if (self->read == NULL) {
5304 PyErr_Format(UnpicklingError,
5305 "Unpickler.__init__() was not called by %s.__init__()",
5306 Py_TYPE(self)->tp_name);
5307 return NULL;
5308 }
5309
5310 return load(self);
5311}
5312
5313/* The name of find_class() is misleading. In newer pickle protocols, this
5314 function is used for loading any global (i.e., functions), not just
5315 classes. The name is kept only for backward compatibility. */
5316
5317PyDoc_STRVAR(Unpickler_find_class_doc,
5318"find_class(module_name, global_name) -> object.\n"
5319"\n"
5320"Return an object from a specified module, importing the module if\n"
5321"necessary. Subclasses may override this method (e.g. to restrict\n"
5322"unpickling of arbitrary classes and functions).\n"
5323"\n"
5324"This method is called whenever a class or a function object is\n"
5325"needed. Both arguments passed are str objects.\n");
5326
5327static PyObject *
5328Unpickler_find_class(UnpicklerObject *self, PyObject *args)
5329{
5330 PyObject *global;
5331 PyObject *modules_dict;
5332 PyObject *module;
5333 PyObject *module_name, *global_name;
5334
5335 if (!PyArg_UnpackTuple(args, "find_class", 2, 2,
5336 &module_name, &global_name))
5337 return NULL;
5338
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005339 /* Try to map the old names used in Python 2.x to the new ones used in
5340 Python 3.x. We do this only with old pickle protocols and when the
5341 user has not disabled the feature. */
5342 if (self->proto < 3 && self->fix_imports) {
5343 PyObject *key;
5344 PyObject *item;
5345
5346 /* Check if the global (i.e., a function or a class) was renamed
5347 or moved to another module. */
5348 key = PyTuple_Pack(2, module_name, global_name);
5349 if (key == NULL)
5350 return NULL;
5351 item = PyDict_GetItemWithError(name_mapping_2to3, key);
5352 Py_DECREF(key);
5353 if (item) {
5354 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
5355 PyErr_Format(PyExc_RuntimeError,
5356 "_compat_pickle.NAME_MAPPING values should be "
5357 "2-tuples, not %.200s", Py_TYPE(item)->tp_name);
5358 return NULL;
5359 }
5360 module_name = PyTuple_GET_ITEM(item, 0);
5361 global_name = PyTuple_GET_ITEM(item, 1);
5362 if (!PyUnicode_Check(module_name) ||
5363 !PyUnicode_Check(global_name)) {
5364 PyErr_Format(PyExc_RuntimeError,
5365 "_compat_pickle.NAME_MAPPING values should be "
5366 "pairs of str, not (%.200s, %.200s)",
5367 Py_TYPE(module_name)->tp_name,
5368 Py_TYPE(global_name)->tp_name);
5369 return NULL;
5370 }
5371 }
5372 else if (PyErr_Occurred()) {
5373 return NULL;
5374 }
5375
5376 /* Check if the module was renamed. */
5377 item = PyDict_GetItemWithError(import_mapping_2to3, module_name);
5378 if (item) {
5379 if (!PyUnicode_Check(item)) {
5380 PyErr_Format(PyExc_RuntimeError,
5381 "_compat_pickle.IMPORT_MAPPING values should be "
5382 "strings, not %.200s", Py_TYPE(item)->tp_name);
5383 return NULL;
5384 }
5385 module_name = item;
5386 }
5387 else if (PyErr_Occurred()) {
5388 return NULL;
5389 }
5390 }
5391
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005392 modules_dict = PySys_GetObject("modules");
5393 if (modules_dict == NULL)
5394 return NULL;
5395
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005396 module = PyDict_GetItemWithError(modules_dict, module_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005397 if (module == NULL) {
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005398 if (PyErr_Occurred())
5399 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005400 module = PyImport_Import(module_name);
5401 if (module == NULL)
5402 return NULL;
5403 global = PyObject_GetAttr(module, global_name);
5404 Py_DECREF(module);
5405 }
5406 else {
5407 global = PyObject_GetAttr(module, global_name);
5408 }
5409 return global;
5410}
5411
5412static struct PyMethodDef Unpickler_methods[] = {
5413 {"load", (PyCFunction)Unpickler_load, METH_NOARGS,
5414 Unpickler_load_doc},
5415 {"find_class", (PyCFunction)Unpickler_find_class, METH_VARARGS,
5416 Unpickler_find_class_doc},
5417 {NULL, NULL} /* sentinel */
5418};
5419
5420static void
5421Unpickler_dealloc(UnpicklerObject *self)
5422{
5423 PyObject_GC_UnTrack((PyObject *)self);
5424 Py_XDECREF(self->readline);
5425 Py_XDECREF(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005426 Py_XDECREF(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005427 Py_XDECREF(self->stack);
5428 Py_XDECREF(self->pers_func);
5429 Py_XDECREF(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005430 if (self->buffer.buf != NULL) {
5431 PyBuffer_Release(&self->buffer);
5432 self->buffer.buf = NULL;
5433 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005434
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005435 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005436 PyMem_Free(self->marks);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005437 PyMem_Free(self->input_line);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005438 free(self->encoding);
5439 free(self->errors);
5440
5441 Py_TYPE(self)->tp_free((PyObject *)self);
5442}
5443
5444static int
5445Unpickler_traverse(UnpicklerObject *self, visitproc visit, void *arg)
5446{
5447 Py_VISIT(self->readline);
5448 Py_VISIT(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005449 Py_VISIT(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005450 Py_VISIT(self->stack);
5451 Py_VISIT(self->pers_func);
5452 Py_VISIT(self->arg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005453 return 0;
5454}
5455
5456static int
5457Unpickler_clear(UnpicklerObject *self)
5458{
5459 Py_CLEAR(self->readline);
5460 Py_CLEAR(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005461 Py_CLEAR(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005462 Py_CLEAR(self->stack);
5463 Py_CLEAR(self->pers_func);
5464 Py_CLEAR(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005465 if (self->buffer.buf != NULL) {
5466 PyBuffer_Release(&self->buffer);
5467 self->buffer.buf = NULL;
5468 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005469
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005470 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005471 PyMem_Free(self->marks);
5472 self->marks = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005473 PyMem_Free(self->input_line);
5474 self->input_line = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005475 free(self->encoding);
5476 self->encoding = NULL;
5477 free(self->errors);
5478 self->errors = NULL;
5479
5480 return 0;
5481}
5482
5483PyDoc_STRVAR(Unpickler_doc,
5484"Unpickler(file, *, encoding='ASCII', errors='strict')"
5485"\n"
5486"This takes a binary file for reading a pickle data stream.\n"
5487"\n"
5488"The protocol version of the pickle is detected automatically, so no\n"
5489"proto argument is needed.\n"
5490"\n"
5491"The file-like object must have two methods, a read() method\n"
5492"that takes an integer argument, and a readline() method that\n"
5493"requires no arguments. Both methods should return bytes.\n"
5494"Thus file-like object can be a binary file object opened for\n"
5495"reading, a BytesIO object, or any other custom object that\n"
5496"meets this interface.\n"
5497"\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005498"Optional keyword arguments are *fix_imports*, *encoding* and *errors*,\n"
5499"which are used to control compatiblity support for pickle stream\n"
5500"generated by Python 2.x. If *fix_imports* is True, pickle will try to\n"
5501"map the old Python 2.x names to the new names used in Python 3.x. The\n"
5502"*encoding* and *errors* tell pickle how to decode 8-bit string\n"
5503"instances pickled by Python 2.x; these default to 'ASCII' and\n"
5504"'strict', respectively.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005505
5506static int
5507Unpickler_init(UnpicklerObject *self, PyObject *args, PyObject *kwds)
5508{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005509 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005510 PyObject *file;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005511 PyObject *fix_imports = Py_True;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005512 char *encoding = NULL;
5513 char *errors = NULL;
5514
5515 /* XXX: That is an horrible error message. But, I don't know how to do
5516 better... */
5517 if (Py_SIZE(args) != 1) {
5518 PyErr_Format(PyExc_TypeError,
5519 "%s takes exactly one positional argument (%zd given)",
5520 Py_TYPE(self)->tp_name, Py_SIZE(args));
5521 return -1;
5522 }
5523
5524 /* Arguments parsing needs to be done in the __init__() method to allow
5525 subclasses to define their own __init__() method, which may (or may
5526 not) support Unpickler arguments. However, this means we need to be
5527 extra careful in the other Unpickler methods, since a subclass could
5528 forget to call Unpickler.__init__() thus breaking our internal
5529 invariants. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005530 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:Unpickler", kwlist,
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005531 &file, &fix_imports, &encoding, &errors))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005532 return -1;
5533
5534 /* In case of multiple __init__() calls, clear previous content. */
5535 if (self->read != NULL)
5536 (void)Unpickler_clear(self);
5537
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005538 if (_Unpickler_SetInputStream(self, file) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005539 return -1;
5540
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005541 if (_Unpickler_SetInputEncoding(self, encoding, errors) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005542 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005543
5544 self->fix_imports = PyObject_IsTrue(fix_imports);
5545 if (self->fix_imports == -1)
5546 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005547
5548 if (PyObject_HasAttrString((PyObject *)self, "persistent_load")) {
5549 self->pers_func = PyObject_GetAttrString((PyObject *)self,
5550 "persistent_load");
5551 if (self->pers_func == NULL)
5552 return -1;
5553 }
5554 else {
5555 self->pers_func = NULL;
5556 }
5557
5558 self->stack = (Pdata *)Pdata_New();
5559 if (self->stack == NULL)
5560 return -1;
5561
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005562 self->memo_size = 32;
5563 self->memo = _Unpickler_NewMemo(self->memo_size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005564 if (self->memo == NULL)
5565 return -1;
5566
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005567 self->arg = NULL;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005568 self->proto = 0;
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005569
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005570 return 0;
5571}
5572
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005573/* Define a proxy object for the Unpickler's internal memo object. This is to
5574 * avoid breaking code like:
5575 * unpickler.memo.clear()
5576 * and
5577 * unpickler.memo = saved_memo
5578 * Is this a good idea? Not really, but we don't want to break code that uses
5579 * it. Note that we don't implement the entire mapping API here. This is
5580 * intentional, as these should be treated as black-box implementation details.
5581 *
5582 * We do, however, have to implement pickling/unpickling support because of
5583 * real-world code like cvs2svn.
5584 */
5585
5586typedef struct {
5587 PyObject_HEAD
5588 UnpicklerObject *unpickler;
5589} UnpicklerMemoProxyObject;
5590
5591PyDoc_STRVAR(ump_clear_doc,
5592"memo.clear() -> None. Remove all items from memo.");
5593
5594static PyObject *
5595ump_clear(UnpicklerMemoProxyObject *self)
5596{
5597 _Unpickler_MemoCleanup(self->unpickler);
5598 self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size);
5599 if (self->unpickler->memo == NULL)
5600 return NULL;
5601 Py_RETURN_NONE;
5602}
5603
5604PyDoc_STRVAR(ump_copy_doc,
5605"memo.copy() -> new_memo. Copy the memo to a new object.");
5606
5607static PyObject *
5608ump_copy(UnpicklerMemoProxyObject *self)
5609{
5610 Py_ssize_t i;
5611 PyObject *new_memo = PyDict_New();
5612 if (new_memo == NULL)
5613 return NULL;
5614
5615 for (i = 0; i < self->unpickler->memo_size; i++) {
5616 int status;
5617 PyObject *key, *value;
5618
5619 value = self->unpickler->memo[i];
5620 if (value == NULL)
5621 continue;
5622
5623 key = PyLong_FromSsize_t(i);
5624 if (key == NULL)
5625 goto error;
5626 status = PyDict_SetItem(new_memo, key, value);
5627 Py_DECREF(key);
5628 if (status < 0)
5629 goto error;
5630 }
5631 return new_memo;
5632
5633error:
5634 Py_DECREF(new_memo);
5635 return NULL;
5636}
5637
5638PyDoc_STRVAR(ump_reduce_doc,
5639"memo.__reduce__(). Pickling support.");
5640
5641static PyObject *
5642ump_reduce(UnpicklerMemoProxyObject *self, PyObject *args)
5643{
5644 PyObject *reduce_value;
5645 PyObject *constructor_args;
5646 PyObject *contents = ump_copy(self);
5647 if (contents == NULL)
5648 return NULL;
5649
5650 reduce_value = PyTuple_New(2);
5651 if (reduce_value == NULL) {
5652 Py_DECREF(contents);
5653 return NULL;
5654 }
5655 constructor_args = PyTuple_New(1);
5656 if (constructor_args == NULL) {
5657 Py_DECREF(contents);
5658 Py_DECREF(reduce_value);
5659 return NULL;
5660 }
5661 PyTuple_SET_ITEM(constructor_args, 0, contents);
5662 Py_INCREF((PyObject *)&PyDict_Type);
5663 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
5664 PyTuple_SET_ITEM(reduce_value, 1, constructor_args);
5665 return reduce_value;
5666}
5667
5668static PyMethodDef unpicklerproxy_methods[] = {
5669 {"clear", (PyCFunction)ump_clear, METH_NOARGS, ump_clear_doc},
5670 {"copy", (PyCFunction)ump_copy, METH_NOARGS, ump_copy_doc},
5671 {"__reduce__", (PyCFunction)ump_reduce, METH_VARARGS, ump_reduce_doc},
5672 {NULL, NULL} /* sentinel */
5673};
5674
5675static void
5676UnpicklerMemoProxy_dealloc(UnpicklerMemoProxyObject *self)
5677{
5678 PyObject_GC_UnTrack(self);
5679 Py_XDECREF(self->unpickler);
5680 PyObject_GC_Del((PyObject *)self);
5681}
5682
5683static int
5684UnpicklerMemoProxy_traverse(UnpicklerMemoProxyObject *self,
5685 visitproc visit, void *arg)
5686{
5687 Py_VISIT(self->unpickler);
5688 return 0;
5689}
5690
5691static int
5692UnpicklerMemoProxy_clear(UnpicklerMemoProxyObject *self)
5693{
5694 Py_CLEAR(self->unpickler);
5695 return 0;
5696}
5697
5698static PyTypeObject UnpicklerMemoProxyType = {
5699 PyVarObject_HEAD_INIT(NULL, 0)
5700 "_pickle.UnpicklerMemoProxy", /*tp_name*/
5701 sizeof(UnpicklerMemoProxyObject), /*tp_basicsize*/
5702 0,
5703 (destructor)UnpicklerMemoProxy_dealloc, /* tp_dealloc */
5704 0, /* tp_print */
5705 0, /* tp_getattr */
5706 0, /* tp_setattr */
5707 0, /* tp_compare */
5708 0, /* tp_repr */
5709 0, /* tp_as_number */
5710 0, /* tp_as_sequence */
5711 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00005712 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005713 0, /* tp_call */
5714 0, /* tp_str */
5715 PyObject_GenericGetAttr, /* tp_getattro */
5716 PyObject_GenericSetAttr, /* tp_setattro */
5717 0, /* tp_as_buffer */
5718 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5719 0, /* tp_doc */
5720 (traverseproc)UnpicklerMemoProxy_traverse, /* tp_traverse */
5721 (inquiry)UnpicklerMemoProxy_clear, /* tp_clear */
5722 0, /* tp_richcompare */
5723 0, /* tp_weaklistoffset */
5724 0, /* tp_iter */
5725 0, /* tp_iternext */
5726 unpicklerproxy_methods, /* tp_methods */
5727};
5728
5729static PyObject *
5730UnpicklerMemoProxy_New(UnpicklerObject *unpickler)
5731{
5732 UnpicklerMemoProxyObject *self;
5733
5734 self = PyObject_GC_New(UnpicklerMemoProxyObject,
5735 &UnpicklerMemoProxyType);
5736 if (self == NULL)
5737 return NULL;
5738 Py_INCREF(unpickler);
5739 self->unpickler = unpickler;
5740 PyObject_GC_Track(self);
5741 return (PyObject *)self;
5742}
5743
5744/*****************************************************************************/
5745
5746
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005747static PyObject *
5748Unpickler_get_memo(UnpicklerObject *self)
5749{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005750 return UnpicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005751}
5752
5753static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005754Unpickler_set_memo(UnpicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005755{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005756 PyObject **new_memo;
5757 Py_ssize_t new_memo_size = 0;
5758 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005759
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005760 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005761 PyErr_SetString(PyExc_TypeError,
5762 "attribute deletion is not supported");
5763 return -1;
5764 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005765
5766 if (Py_TYPE(obj) == &UnpicklerMemoProxyType) {
5767 UnpicklerObject *unpickler =
5768 ((UnpicklerMemoProxyObject *)obj)->unpickler;
5769
5770 new_memo_size = unpickler->memo_size;
5771 new_memo = _Unpickler_NewMemo(new_memo_size);
5772 if (new_memo == NULL)
5773 return -1;
5774
5775 for (i = 0; i < new_memo_size; i++) {
5776 Py_XINCREF(unpickler->memo[i]);
5777 new_memo[i] = unpickler->memo[i];
5778 }
5779 }
5780 else if (PyDict_Check(obj)) {
5781 Py_ssize_t i = 0;
5782 PyObject *key, *value;
5783
5784 new_memo_size = PyDict_Size(obj);
5785 new_memo = _Unpickler_NewMemo(new_memo_size);
5786 if (new_memo == NULL)
5787 return -1;
5788
5789 while (PyDict_Next(obj, &i, &key, &value)) {
5790 Py_ssize_t idx;
5791 if (!PyLong_Check(key)) {
5792 PyErr_SetString(PyExc_TypeError,
5793 "memo key must be integers");
5794 goto error;
5795 }
5796 idx = PyLong_AsSsize_t(key);
5797 if (idx == -1 && PyErr_Occurred())
5798 goto error;
5799 if (_Unpickler_MemoPut(self, idx, value) < 0)
5800 goto error;
5801 }
5802 }
5803 else {
5804 PyErr_Format(PyExc_TypeError,
5805 "'memo' attribute must be an UnpicklerMemoProxy object"
5806 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005807 return -1;
5808 }
5809
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005810 _Unpickler_MemoCleanup(self);
5811 self->memo_size = new_memo_size;
5812 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005813
5814 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005815
5816 error:
5817 if (new_memo_size) {
5818 i = new_memo_size;
5819 while (--i >= 0) {
5820 Py_XDECREF(new_memo[i]);
5821 }
5822 PyMem_FREE(new_memo);
5823 }
5824 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005825}
5826
5827static PyObject *
5828Unpickler_get_persload(UnpicklerObject *self)
5829{
5830 if (self->pers_func == NULL)
5831 PyErr_SetString(PyExc_AttributeError, "persistent_load");
5832 else
5833 Py_INCREF(self->pers_func);
5834 return self->pers_func;
5835}
5836
5837static int
5838Unpickler_set_persload(UnpicklerObject *self, PyObject *value)
5839{
5840 PyObject *tmp;
5841
5842 if (value == NULL) {
5843 PyErr_SetString(PyExc_TypeError,
5844 "attribute deletion is not supported");
5845 return -1;
5846 }
5847 if (!PyCallable_Check(value)) {
5848 PyErr_SetString(PyExc_TypeError,
5849 "persistent_load must be a callable taking "
5850 "one argument");
5851 return -1;
5852 }
5853
5854 tmp = self->pers_func;
5855 Py_INCREF(value);
5856 self->pers_func = value;
5857 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
5858
5859 return 0;
5860}
5861
5862static PyGetSetDef Unpickler_getsets[] = {
5863 {"memo", (getter)Unpickler_get_memo, (setter)Unpickler_set_memo},
5864 {"persistent_load", (getter)Unpickler_get_persload,
5865 (setter)Unpickler_set_persload},
5866 {NULL}
5867};
5868
5869static PyTypeObject Unpickler_Type = {
5870 PyVarObject_HEAD_INIT(NULL, 0)
5871 "_pickle.Unpickler", /*tp_name*/
5872 sizeof(UnpicklerObject), /*tp_basicsize*/
5873 0, /*tp_itemsize*/
5874 (destructor)Unpickler_dealloc, /*tp_dealloc*/
5875 0, /*tp_print*/
5876 0, /*tp_getattr*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005877 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00005878 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005879 0, /*tp_repr*/
5880 0, /*tp_as_number*/
5881 0, /*tp_as_sequence*/
5882 0, /*tp_as_mapping*/
5883 0, /*tp_hash*/
5884 0, /*tp_call*/
5885 0, /*tp_str*/
5886 0, /*tp_getattro*/
5887 0, /*tp_setattro*/
5888 0, /*tp_as_buffer*/
5889 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5890 Unpickler_doc, /*tp_doc*/
5891 (traverseproc)Unpickler_traverse, /*tp_traverse*/
5892 (inquiry)Unpickler_clear, /*tp_clear*/
5893 0, /*tp_richcompare*/
5894 0, /*tp_weaklistoffset*/
5895 0, /*tp_iter*/
5896 0, /*tp_iternext*/
5897 Unpickler_methods, /*tp_methods*/
5898 0, /*tp_members*/
5899 Unpickler_getsets, /*tp_getset*/
5900 0, /*tp_base*/
5901 0, /*tp_dict*/
5902 0, /*tp_descr_get*/
5903 0, /*tp_descr_set*/
5904 0, /*tp_dictoffset*/
5905 (initproc)Unpickler_init, /*tp_init*/
5906 PyType_GenericAlloc, /*tp_alloc*/
5907 PyType_GenericNew, /*tp_new*/
5908 PyObject_GC_Del, /*tp_free*/
5909 0, /*tp_is_gc*/
5910};
5911
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005912PyDoc_STRVAR(pickle_dump_doc,
5913"dump(obj, file, protocol=None, *, fix_imports=True) -> None\n"
5914"\n"
5915"Write a pickled representation of obj to the open file object file. This\n"
5916"is equivalent to ``Pickler(file, protocol).dump(obj)``, but may be more\n"
5917"efficient.\n"
5918"\n"
5919"The optional protocol argument tells the pickler to use the given protocol;\n"
5920"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
5921"backward-incompatible protocol designed for Python 3.0.\n"
5922"\n"
5923"Specifying a negative protocol version selects the highest protocol version\n"
5924"supported. The higher the protocol used, the more recent the version of\n"
5925"Python needed to read the pickle produced.\n"
5926"\n"
5927"The file argument must have a write() method that accepts a single bytes\n"
5928"argument. It can thus be a file object opened for binary writing, a\n"
5929"io.BytesIO instance, or any other custom object that meets this interface.\n"
5930"\n"
5931"If fix_imports is True and protocol is less than 3, pickle will try to\n"
5932"map the new Python 3.x names to the old module names used in Python 2.x,\n"
5933"so that the pickle data stream is readable with Python 2.x.\n");
5934
5935static PyObject *
5936pickle_dump(PyObject *self, PyObject *args, PyObject *kwds)
5937{
5938 static char *kwlist[] = {"obj", "file", "protocol", "fix_imports", 0};
5939 PyObject *obj;
5940 PyObject *file;
5941 PyObject *proto = NULL;
5942 PyObject *fix_imports = Py_True;
5943 PicklerObject *pickler;
5944
5945 /* fix_imports is a keyword-only argument. */
5946 if (Py_SIZE(args) > 3) {
5947 PyErr_Format(PyExc_TypeError,
5948 "pickle.dump() takes at most 3 positional "
5949 "argument (%zd given)", Py_SIZE(args));
5950 return NULL;
5951 }
5952
5953 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|OO:dump", kwlist,
5954 &obj, &file, &proto, &fix_imports))
5955 return NULL;
5956
5957 pickler = _Pickler_New();
5958 if (pickler == NULL)
5959 return NULL;
5960
5961 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
5962 goto error;
5963
5964 if (_Pickler_SetOutputStream(pickler, file) < 0)
5965 goto error;
5966
5967 if (dump(pickler, obj) < 0)
5968 goto error;
5969
5970 if (_Pickler_FlushToFile(pickler) < 0)
5971 goto error;
5972
5973 Py_DECREF(pickler);
5974 Py_RETURN_NONE;
5975
5976 error:
5977 Py_XDECREF(pickler);
5978 return NULL;
5979}
5980
5981PyDoc_STRVAR(pickle_dumps_doc,
5982"dumps(obj, protocol=None, *, fix_imports=True) -> bytes\n"
5983"\n"
5984"Return the pickled representation of the object as a bytes\n"
5985"object, instead of writing it to a file.\n"
5986"\n"
5987"The optional protocol argument tells the pickler to use the given protocol;\n"
5988"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
5989"backward-incompatible protocol designed for Python 3.0.\n"
5990"\n"
5991"Specifying a negative protocol version selects the highest protocol version\n"
5992"supported. The higher the protocol used, the more recent the version of\n"
5993"Python needed to read the pickle produced.\n"
5994"\n"
5995"If fix_imports is True and *protocol* is less than 3, pickle will try to\n"
5996"map the new Python 3.x names to the old module names used in Python 2.x,\n"
5997"so that the pickle data stream is readable with Python 2.x.\n");
5998
5999static PyObject *
6000pickle_dumps(PyObject *self, PyObject *args, PyObject *kwds)
6001{
6002 static char *kwlist[] = {"obj", "protocol", "fix_imports", 0};
6003 PyObject *obj;
6004 PyObject *proto = NULL;
6005 PyObject *result;
6006 PyObject *fix_imports = Py_True;
6007 PicklerObject *pickler;
6008
6009 /* fix_imports is a keyword-only argument. */
6010 if (Py_SIZE(args) > 2) {
6011 PyErr_Format(PyExc_TypeError,
6012 "pickle.dumps() takes at most 2 positional "
6013 "argument (%zd given)", Py_SIZE(args));
6014 return NULL;
6015 }
6016
6017 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:dumps", kwlist,
6018 &obj, &proto, &fix_imports))
6019 return NULL;
6020
6021 pickler = _Pickler_New();
6022 if (pickler == NULL)
6023 return NULL;
6024
6025 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6026 goto error;
6027
6028 if (dump(pickler, obj) < 0)
6029 goto error;
6030
6031 result = _Pickler_GetString(pickler);
6032 Py_DECREF(pickler);
6033 return result;
6034
6035 error:
6036 Py_XDECREF(pickler);
6037 return NULL;
6038}
6039
6040PyDoc_STRVAR(pickle_load_doc,
6041"load(file, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6042"\n"
6043"Read a pickled object representation from the open file object file and\n"
6044"return the reconstituted object hierarchy specified therein. This is\n"
6045"equivalent to ``Unpickler(file).load()``, but may be more efficient.\n"
6046"\n"
6047"The protocol version of the pickle is detected automatically, so no protocol\n"
6048"argument is needed. Bytes past the pickled object's representation are\n"
6049"ignored.\n"
6050"\n"
6051"The argument file must have two methods, a read() method that takes an\n"
6052"integer argument, and a readline() method that requires no arguments. Both\n"
6053"methods should return bytes. Thus *file* can be a binary file object opened\n"
6054"for reading, a BytesIO object, or any other custom object that meets this\n"
6055"interface.\n"
6056"\n"
6057"Optional keyword arguments are fix_imports, encoding and errors,\n"
6058"which are used to control compatiblity support for pickle stream generated\n"
6059"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6060"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6061"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6062"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6063
6064static PyObject *
6065pickle_load(PyObject *self, PyObject *args, PyObject *kwds)
6066{
6067 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
6068 PyObject *file;
6069 PyObject *fix_imports = Py_True;
6070 PyObject *result;
6071 char *encoding = NULL;
6072 char *errors = NULL;
6073 UnpicklerObject *unpickler;
6074
6075 /* fix_imports, encoding and errors are a keyword-only argument. */
6076 if (Py_SIZE(args) != 1) {
6077 PyErr_Format(PyExc_TypeError,
6078 "pickle.load() takes exactly one positional "
6079 "argument (%zd given)", Py_SIZE(args));
6080 return NULL;
6081 }
6082
6083 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:load", kwlist,
6084 &file, &fix_imports, &encoding, &errors))
6085 return NULL;
6086
6087 unpickler = _Unpickler_New();
6088 if (unpickler == NULL)
6089 return NULL;
6090
6091 if (_Unpickler_SetInputStream(unpickler, file) < 0)
6092 goto error;
6093
6094 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6095 goto error;
6096
6097 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6098 if (unpickler->fix_imports == -1)
6099 goto error;
6100
6101 result = load(unpickler);
6102 Py_DECREF(unpickler);
6103 return result;
6104
6105 error:
6106 Py_XDECREF(unpickler);
6107 return NULL;
6108}
6109
6110PyDoc_STRVAR(pickle_loads_doc,
6111"loads(input, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6112"\n"
6113"Read a pickled object hierarchy from a bytes object and return the\n"
6114"reconstituted object hierarchy specified therein\n"
6115"\n"
6116"The protocol version of the pickle is detected automatically, so no protocol\n"
6117"argument is needed. Bytes past the pickled object's representation are\n"
6118"ignored.\n"
6119"\n"
6120"Optional keyword arguments are fix_imports, encoding and errors, which\n"
6121"are used to control compatiblity support for pickle stream generated\n"
6122"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6123"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6124"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6125"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6126
6127static PyObject *
6128pickle_loads(PyObject *self, PyObject *args, PyObject *kwds)
6129{
6130 static char *kwlist[] = {"input", "fix_imports", "encoding", "errors", 0};
6131 PyObject *input;
6132 PyObject *fix_imports = Py_True;
6133 PyObject *result;
6134 char *encoding = NULL;
6135 char *errors = NULL;
6136 UnpicklerObject *unpickler;
6137
6138 /* fix_imports, encoding and errors are a keyword-only argument. */
6139 if (Py_SIZE(args) != 1) {
6140 PyErr_Format(PyExc_TypeError,
6141 "pickle.loads() takes exactly one positional "
6142 "argument (%zd given)", Py_SIZE(args));
6143 return NULL;
6144 }
6145
6146 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:loads", kwlist,
6147 &input, &fix_imports, &encoding, &errors))
6148 return NULL;
6149
6150 unpickler = _Unpickler_New();
6151 if (unpickler == NULL)
6152 return NULL;
6153
6154 if (_Unpickler_SetStringInput(unpickler, input) < 0)
6155 goto error;
6156
6157 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6158 goto error;
6159
6160 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6161 if (unpickler->fix_imports == -1)
6162 goto error;
6163
6164 result = load(unpickler);
6165 Py_DECREF(unpickler);
6166 return result;
6167
6168 error:
6169 Py_XDECREF(unpickler);
6170 return NULL;
6171}
6172
6173
6174static struct PyMethodDef pickle_methods[] = {
6175 {"dump", (PyCFunction)pickle_dump, METH_VARARGS|METH_KEYWORDS,
6176 pickle_dump_doc},
6177 {"dumps", (PyCFunction)pickle_dumps, METH_VARARGS|METH_KEYWORDS,
6178 pickle_dumps_doc},
6179 {"load", (PyCFunction)pickle_load, METH_VARARGS|METH_KEYWORDS,
6180 pickle_load_doc},
6181 {"loads", (PyCFunction)pickle_loads, METH_VARARGS|METH_KEYWORDS,
6182 pickle_loads_doc},
6183 {NULL, NULL} /* sentinel */
6184};
6185
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006186static int
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006187initmodule(void)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006188{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006189 PyObject *copyreg = NULL;
6190 PyObject *compat_pickle = NULL;
6191
6192 /* XXX: We should ensure that the types of the dictionaries imported are
6193 exactly PyDict objects. Otherwise, it is possible to crash the pickle
6194 since we use the PyDict API directly to access these dictionaries. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006195
6196 copyreg = PyImport_ImportModule("copyreg");
6197 if (!copyreg)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006198 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006199 dispatch_table = PyObject_GetAttrString(copyreg, "dispatch_table");
6200 if (!dispatch_table)
6201 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006202 extension_registry = \
6203 PyObject_GetAttrString(copyreg, "_extension_registry");
6204 if (!extension_registry)
6205 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006206 inverted_registry = PyObject_GetAttrString(copyreg, "_inverted_registry");
6207 if (!inverted_registry)
6208 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006209 extension_cache = PyObject_GetAttrString(copyreg, "_extension_cache");
6210 if (!extension_cache)
6211 goto error;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006212 Py_CLEAR(copyreg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006213
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006214 /* Load the 2.x -> 3.x stdlib module mapping tables */
6215 compat_pickle = PyImport_ImportModule("_compat_pickle");
6216 if (!compat_pickle)
6217 goto error;
6218 name_mapping_2to3 = PyObject_GetAttrString(compat_pickle, "NAME_MAPPING");
6219 if (!name_mapping_2to3)
6220 goto error;
6221 if (!PyDict_CheckExact(name_mapping_2to3)) {
6222 PyErr_Format(PyExc_RuntimeError,
6223 "_compat_pickle.NAME_MAPPING should be a dict, not %.200s",
6224 Py_TYPE(name_mapping_2to3)->tp_name);
6225 goto error;
6226 }
6227 import_mapping_2to3 = PyObject_GetAttrString(compat_pickle,
6228 "IMPORT_MAPPING");
6229 if (!import_mapping_2to3)
6230 goto error;
6231 if (!PyDict_CheckExact(import_mapping_2to3)) {
6232 PyErr_Format(PyExc_RuntimeError,
6233 "_compat_pickle.IMPORT_MAPPING should be a dict, "
6234 "not %.200s", Py_TYPE(import_mapping_2to3)->tp_name);
6235 goto error;
6236 }
6237 /* ... and the 3.x -> 2.x mapping tables */
6238 name_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6239 "REVERSE_NAME_MAPPING");
6240 if (!name_mapping_3to2)
6241 goto error;
6242 if (!PyDict_CheckExact(name_mapping_3to2)) {
6243 PyErr_Format(PyExc_RuntimeError,
Ezio Melotti13925002011-03-16 11:05:33 +02006244 "_compat_pickle.REVERSE_NAME_MAPPING should be a dict, "
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006245 "not %.200s", Py_TYPE(name_mapping_3to2)->tp_name);
6246 goto error;
6247 }
6248 import_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6249 "REVERSE_IMPORT_MAPPING");
6250 if (!import_mapping_3to2)
6251 goto error;
6252 if (!PyDict_CheckExact(import_mapping_3to2)) {
6253 PyErr_Format(PyExc_RuntimeError,
6254 "_compat_pickle.REVERSE_IMPORT_MAPPING should be a dict, "
6255 "not %.200s", Py_TYPE(import_mapping_3to2)->tp_name);
6256 goto error;
6257 }
6258 Py_CLEAR(compat_pickle);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006259
6260 empty_tuple = PyTuple_New(0);
6261 if (empty_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006262 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006263 two_tuple = PyTuple_New(2);
6264 if (two_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006265 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006266 /* We use this temp container with no regard to refcounts, or to
6267 * keeping containees alive. Exempt from GC, because we don't
6268 * want anything looking at two_tuple() by magic.
6269 */
6270 PyObject_GC_UnTrack(two_tuple);
6271
6272 return 0;
6273
6274 error:
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006275 Py_CLEAR(copyreg);
6276 Py_CLEAR(dispatch_table);
6277 Py_CLEAR(extension_registry);
6278 Py_CLEAR(inverted_registry);
6279 Py_CLEAR(extension_cache);
6280 Py_CLEAR(compat_pickle);
6281 Py_CLEAR(name_mapping_2to3);
6282 Py_CLEAR(import_mapping_2to3);
6283 Py_CLEAR(name_mapping_3to2);
6284 Py_CLEAR(import_mapping_3to2);
6285 Py_CLEAR(empty_tuple);
6286 Py_CLEAR(two_tuple);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006287 return -1;
6288}
6289
6290static struct PyModuleDef _picklemodule = {
6291 PyModuleDef_HEAD_INIT,
6292 "_pickle",
6293 pickle_module_doc,
6294 -1,
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006295 pickle_methods,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006296 NULL,
6297 NULL,
6298 NULL,
6299 NULL
6300};
6301
6302PyMODINIT_FUNC
6303PyInit__pickle(void)
6304{
6305 PyObject *m;
6306
6307 if (PyType_Ready(&Unpickler_Type) < 0)
6308 return NULL;
6309 if (PyType_Ready(&Pickler_Type) < 0)
6310 return NULL;
6311 if (PyType_Ready(&Pdata_Type) < 0)
6312 return NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006313 if (PyType_Ready(&PicklerMemoProxyType) < 0)
6314 return NULL;
6315 if (PyType_Ready(&UnpicklerMemoProxyType) < 0)
6316 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006317
6318 /* Create the module and add the functions. */
6319 m = PyModule_Create(&_picklemodule);
6320 if (m == NULL)
6321 return NULL;
6322
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006323 Py_INCREF(&Pickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006324 if (PyModule_AddObject(m, "Pickler", (PyObject *)&Pickler_Type) < 0)
6325 return NULL;
Antoine Pitrou8391cf42011-07-15 21:01:21 +02006326 Py_INCREF(&Unpickler_Type);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006327 if (PyModule_AddObject(m, "Unpickler", (PyObject *)&Unpickler_Type) < 0)
6328 return NULL;
6329
6330 /* Initialize the exceptions. */
6331 PickleError = PyErr_NewException("_pickle.PickleError", NULL, NULL);
6332 if (PickleError == NULL)
6333 return NULL;
6334 PicklingError = \
6335 PyErr_NewException("_pickle.PicklingError", PickleError, NULL);
6336 if (PicklingError == NULL)
6337 return NULL;
6338 UnpicklingError = \
6339 PyErr_NewException("_pickle.UnpicklingError", PickleError, NULL);
6340 if (UnpicklingError == NULL)
6341 return NULL;
6342
6343 if (PyModule_AddObject(m, "PickleError", PickleError) < 0)
6344 return NULL;
6345 if (PyModule_AddObject(m, "PicklingError", PicklingError) < 0)
6346 return NULL;
6347 if (PyModule_AddObject(m, "UnpicklingError", UnpicklingError) < 0)
6348 return NULL;
6349
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006350 if (initmodule() < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006351 return NULL;
6352
6353 return m;
6354}