blob: 0ef1f22afa287a7e4ef30beaff79aca0aa4cdc6b [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 Pitrouea99c5c2010-09-09 18:33:21 +0000980 if (n == 0) {
981 *s = NULL;
982 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +0000983 }
984
Antoine Pitrou04248a82010-10-12 20:51:21 +0000985 if (self->next_read_idx + n <= self->input_len) {
986 *s = self->input_buffer + self->next_read_idx;
987 self->next_read_idx += n;
988 return n;
989 }
990 if (!self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +0000991 PyErr_Format(PyExc_EOFError, "Ran out of input");
Amaury Forgeot d'Arc3e4e72f2008-11-11 20:05:06 +0000992 return -1;
993 }
Antoine Pitrou04248a82010-10-12 20:51:21 +0000994 num_read = _Unpickler_ReadFromFile(self, n);
995 if (num_read < 0)
996 return -1;
997 if (num_read < n) {
998 PyErr_Format(PyExc_EOFError, "Ran out of input");
999 return -1;
1000 }
1001 *s = self->input_buffer;
1002 self->next_read_idx = n;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001003 return n;
1004}
1005
1006static Py_ssize_t
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001007_Unpickler_CopyLine(UnpicklerObject *self, char *line, Py_ssize_t len,
1008 char **result)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001009{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001010 char *input_line = PyMem_Realloc(self->input_line, len + 1);
1011 if (input_line == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001012 return -1;
1013
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001014 memcpy(input_line, line, len);
1015 input_line[len] = '\0';
1016 self->input_line = input_line;
1017 *result = self->input_line;
1018 return len;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001019}
1020
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001021/* Read a line from the input stream/buffer. If we run off the end of the input
1022 before hitting \n, return the data we found.
1023
1024 Returns the number of chars read, or -1 on failure. */
1025static Py_ssize_t
1026_Unpickler_Readline(UnpicklerObject *self, char **result)
1027{
1028 Py_ssize_t i, num_read;
1029
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001030 for (i = self->next_read_idx; i < self->input_len; i++) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001031 if (self->input_buffer[i] == '\n') {
1032 char *line_start = self->input_buffer + self->next_read_idx;
1033 num_read = i - self->next_read_idx + 1;
1034 self->next_read_idx = i + 1;
1035 return _Unpickler_CopyLine(self, line_start, num_read, result);
1036 }
1037 }
1038 if (self->read) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001039 num_read = _Unpickler_ReadFromFile(self, READ_WHOLE_LINE);
1040 if (num_read < 0)
1041 return -1;
1042 *result = self->input_buffer;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001043 self->next_read_idx = num_read;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001044 return num_read;
1045 }
1046
1047 /* If we get here, we've run off the end of the input string. Return the
1048 remaining string and let the caller figure it out. */
1049 *result = self->input_buffer + self->next_read_idx;
1050 num_read = i - self->next_read_idx;
1051 self->next_read_idx = i;
1052 return num_read;
1053}
1054
1055/* Returns -1 (with an exception set) on failure, 0 on success. The memo array
1056 will be modified in place. */
1057static int
1058_Unpickler_ResizeMemoList(UnpicklerObject *self, Py_ssize_t new_size)
1059{
1060 Py_ssize_t i;
1061 PyObject **memo;
1062
1063 assert(new_size > self->memo_size);
1064
1065 memo = PyMem_REALLOC(self->memo, new_size * sizeof(PyObject *));
1066 if (memo == NULL) {
1067 PyErr_NoMemory();
1068 return -1;
1069 }
1070 self->memo = memo;
1071 for (i = self->memo_size; i < new_size; i++)
1072 self->memo[i] = NULL;
1073 self->memo_size = new_size;
1074 return 0;
1075}
1076
1077/* Returns NULL if idx is out of bounds. */
1078static PyObject *
1079_Unpickler_MemoGet(UnpicklerObject *self, Py_ssize_t idx)
1080{
1081 if (idx < 0 || idx >= self->memo_size)
1082 return NULL;
1083
1084 return self->memo[idx];
1085}
1086
1087/* Returns -1 (with an exception set) on failure, 0 on success.
1088 This takes its own reference to `value`. */
1089static int
1090_Unpickler_MemoPut(UnpicklerObject *self, Py_ssize_t idx, PyObject *value)
1091{
1092 PyObject *old_item;
1093
1094 if (idx >= self->memo_size) {
1095 if (_Unpickler_ResizeMemoList(self, idx * 2) < 0)
1096 return -1;
1097 assert(idx < self->memo_size);
1098 }
1099 Py_INCREF(value);
1100 old_item = self->memo[idx];
1101 self->memo[idx] = value;
1102 Py_XDECREF(old_item);
1103 return 0;
1104}
1105
1106static PyObject **
1107_Unpickler_NewMemo(Py_ssize_t new_size)
1108{
1109 PyObject **memo = PyMem_MALLOC(new_size * sizeof(PyObject *));
1110 if (memo == NULL)
1111 return NULL;
1112 memset(memo, 0, new_size * sizeof(PyObject *));
1113 return memo;
1114}
1115
1116/* Free the unpickler's memo, taking care to decref any items left in it. */
1117static void
1118_Unpickler_MemoCleanup(UnpicklerObject *self)
1119{
1120 Py_ssize_t i;
1121 PyObject **memo = self->memo;
1122
1123 if (self->memo == NULL)
1124 return;
1125 self->memo = NULL;
1126 i = self->memo_size;
1127 while (--i >= 0) {
1128 Py_XDECREF(memo[i]);
1129 }
1130 PyMem_FREE(memo);
1131}
1132
1133static UnpicklerObject *
1134_Unpickler_New(void)
1135{
1136 UnpicklerObject *self;
1137
1138 self = PyObject_GC_New(UnpicklerObject, &Unpickler_Type);
1139 if (self == NULL)
1140 return NULL;
1141
1142 self->stack = (Pdata *)Pdata_New();
1143 if (self->stack == NULL) {
1144 Py_DECREF(self);
1145 return NULL;
1146 }
1147 memset(&self->buffer, 0, sizeof(Py_buffer));
1148
1149 self->memo_size = 32;
1150 self->memo = _Unpickler_NewMemo(self->memo_size);
1151 if (self->memo == NULL) {
1152 Py_DECREF(self);
1153 return NULL;
1154 }
1155
1156 self->arg = NULL;
1157 self->pers_func = NULL;
1158 self->input_buffer = NULL;
1159 self->input_line = NULL;
1160 self->input_len = 0;
1161 self->next_read_idx = 0;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001162 self->prefetched_idx = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001163 self->read = NULL;
1164 self->readline = NULL;
Antoine Pitrou04248a82010-10-12 20:51:21 +00001165 self->peek = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001166 self->encoding = NULL;
1167 self->errors = NULL;
1168 self->marks = NULL;
1169 self->num_marks = 0;
1170 self->marks_size = 0;
1171 self->proto = 0;
1172 self->fix_imports = 0;
1173
1174 return self;
1175}
1176
1177/* Returns -1 (with an exception set) on failure, 0 on success. This may
1178 be called once on a freshly created Pickler. */
1179static int
1180_Unpickler_SetInputStream(UnpicklerObject *self, PyObject *file)
1181{
Antoine Pitrou04248a82010-10-12 20:51:21 +00001182 self->peek = PyObject_GetAttrString(file, "peek");
1183 if (self->peek == NULL) {
1184 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1185 PyErr_Clear();
1186 else
1187 return -1;
1188 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001189 self->read = PyObject_GetAttrString(file, "read");
1190 self->readline = PyObject_GetAttrString(file, "readline");
1191 if (self->readline == NULL || self->read == NULL) {
1192 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1193 PyErr_SetString(PyExc_TypeError,
1194 "file must have 'read' and 'readline' attributes");
1195 Py_CLEAR(self->read);
1196 Py_CLEAR(self->readline);
Antoine Pitrou04248a82010-10-12 20:51:21 +00001197 Py_CLEAR(self->peek);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001198 return -1;
1199 }
1200 return 0;
1201}
1202
1203/* Returns -1 (with an exception set) on failure, 0 on success. This may
1204 be called once on a freshly created Pickler. */
1205static int
1206_Unpickler_SetInputEncoding(UnpicklerObject *self,
1207 const char *encoding,
1208 const char *errors)
1209{
1210 if (encoding == NULL)
1211 encoding = "ASCII";
1212 if (errors == NULL)
1213 errors = "strict";
1214
1215 self->encoding = strdup(encoding);
1216 self->errors = strdup(errors);
1217 if (self->encoding == NULL || self->errors == NULL) {
1218 PyErr_NoMemory();
1219 return -1;
1220 }
1221 return 0;
1222}
1223
1224/* Generate a GET opcode for an object stored in the memo. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001225static int
1226memo_get(PicklerObject *self, PyObject *key)
1227{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001228 long *value;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001229 char pdata[30];
1230 int len;
1231
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001232 value = PyMemoTable_Get(self->memo, key);
1233 if (value == NULL) {
1234 PyErr_SetObject(PyExc_KeyError, key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001235 return -1;
1236 }
1237
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001238 if (!self->bin) {
1239 pdata[0] = GET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001240 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1, "%ld\n", *value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001241 len = (int)strlen(pdata);
1242 }
1243 else {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001244 if (*value < 256) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001245 pdata[0] = BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001246 pdata[1] = (unsigned char)(*value & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001247 len = 2;
1248 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001249 else if (*value <= 0xffffffffL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001250 pdata[0] = LONG_BINGET;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001251 pdata[1] = (unsigned char)(*value & 0xff);
1252 pdata[2] = (unsigned char)((*value >> 8) & 0xff);
1253 pdata[3] = (unsigned char)((*value >> 16) & 0xff);
1254 pdata[4] = (unsigned char)((*value >> 24) & 0xff);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001255 len = 5;
1256 }
1257 else { /* unlikely */
1258 PyErr_SetString(PicklingError,
1259 "memo id too large for LONG_BINGET");
1260 return -1;
1261 }
1262 }
1263
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001264 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001265 return -1;
1266
1267 return 0;
1268}
1269
1270/* Store an object in the memo, assign it a new unique ID based on the number
1271 of objects currently stored in the memo and generate a PUT opcode. */
1272static int
1273memo_put(PicklerObject *self, PyObject *obj)
1274{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001275 long x;
1276 char pdata[30];
1277 int len;
1278 int status = 0;
1279
1280 if (self->fast)
1281 return 0;
1282
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001283 x = PyMemoTable_Size(self->memo);
1284 if (PyMemoTable_Set(self->memo, obj, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001285 goto error;
1286
1287 if (!self->bin) {
1288 pdata[0] = PUT;
1289 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1, "%ld\n", x);
1290 len = strlen(pdata);
1291 }
1292 else {
1293 if (x < 256) {
1294 pdata[0] = BINPUT;
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00001295 pdata[1] = (unsigned char)x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001296 len = 2;
1297 }
1298 else if (x <= 0xffffffffL) {
1299 pdata[0] = LONG_BINPUT;
1300 pdata[1] = (unsigned char)(x & 0xff);
1301 pdata[2] = (unsigned char)((x >> 8) & 0xff);
1302 pdata[3] = (unsigned char)((x >> 16) & 0xff);
1303 pdata[4] = (unsigned char)((x >> 24) & 0xff);
1304 len = 5;
1305 }
1306 else { /* unlikely */
1307 PyErr_SetString(PicklingError,
1308 "memo id too large for LONG_BINPUT");
1309 return -1;
1310 }
1311 }
1312
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001313 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001314 goto error;
1315
1316 if (0) {
1317 error:
1318 status = -1;
1319 }
1320
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001321 return status;
1322}
1323
1324static PyObject *
1325whichmodule(PyObject *global, PyObject *global_name)
1326{
1327 Py_ssize_t i, j;
1328 static PyObject *module_str = NULL;
1329 static PyObject *main_str = NULL;
1330 PyObject *module_name;
1331 PyObject *modules_dict;
1332 PyObject *module;
1333 PyObject *obj;
1334
1335 if (module_str == NULL) {
1336 module_str = PyUnicode_InternFromString("__module__");
1337 if (module_str == NULL)
1338 return NULL;
1339 main_str = PyUnicode_InternFromString("__main__");
1340 if (main_str == NULL)
1341 return NULL;
1342 }
1343
1344 module_name = PyObject_GetAttr(global, module_str);
1345
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00001346 /* In some rare cases (e.g., bound methods of extension types),
1347 __module__ can be None. If it is so, then search sys.modules
1348 for the module of global. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001349 if (module_name == Py_None) {
1350 Py_DECREF(module_name);
1351 goto search;
1352 }
1353
1354 if (module_name) {
1355 return module_name;
1356 }
1357 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1358 PyErr_Clear();
1359 else
1360 return NULL;
1361
1362 search:
1363 modules_dict = PySys_GetObject("modules");
1364 if (modules_dict == NULL)
1365 return NULL;
1366
1367 i = 0;
1368 module_name = NULL;
1369 while ((j = PyDict_Next(modules_dict, &i, &module_name, &module))) {
Mark Dickinson211c6252009-02-01 10:28:51 +00001370 if (PyObject_RichCompareBool(module_name, main_str, Py_EQ) == 1)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001371 continue;
1372
1373 obj = PyObject_GetAttr(module, global_name);
1374 if (obj == NULL) {
1375 if (PyErr_ExceptionMatches(PyExc_AttributeError))
1376 PyErr_Clear();
1377 else
1378 return NULL;
1379 continue;
1380 }
1381
1382 if (obj != global) {
1383 Py_DECREF(obj);
1384 continue;
1385 }
1386
1387 Py_DECREF(obj);
1388 break;
1389 }
1390
1391 /* If no module is found, use __main__. */
1392 if (!j) {
1393 module_name = main_str;
1394 }
1395
1396 Py_INCREF(module_name);
1397 return module_name;
1398}
1399
1400/* fast_save_enter() and fast_save_leave() are guards against recursive
1401 objects when Pickler is used with the "fast mode" (i.e., with object
1402 memoization disabled). If the nesting of a list or dict object exceed
1403 FAST_NESTING_LIMIT, these guards will start keeping an internal
1404 reference to the seen list or dict objects and check whether these objects
1405 are recursive. These are not strictly necessary, since save() has a
1406 hard-coded recursion limit, but they give a nicer error message than the
1407 typical RuntimeError. */
1408static int
1409fast_save_enter(PicklerObject *self, PyObject *obj)
1410{
1411 /* if fast_nesting < 0, we're doing an error exit. */
1412 if (++self->fast_nesting >= FAST_NESTING_LIMIT) {
1413 PyObject *key = NULL;
1414 if (self->fast_memo == NULL) {
1415 self->fast_memo = PyDict_New();
1416 if (self->fast_memo == NULL) {
1417 self->fast_nesting = -1;
1418 return 0;
1419 }
1420 }
1421 key = PyLong_FromVoidPtr(obj);
1422 if (key == NULL)
1423 return 0;
1424 if (PyDict_GetItem(self->fast_memo, key)) {
1425 Py_DECREF(key);
1426 PyErr_Format(PyExc_ValueError,
1427 "fast mode: can't pickle cyclic objects "
1428 "including object type %.200s at %p",
1429 obj->ob_type->tp_name, obj);
1430 self->fast_nesting = -1;
1431 return 0;
1432 }
1433 if (PyDict_SetItem(self->fast_memo, key, Py_None) < 0) {
1434 Py_DECREF(key);
1435 self->fast_nesting = -1;
1436 return 0;
1437 }
1438 Py_DECREF(key);
1439 }
1440 return 1;
1441}
1442
1443static int
1444fast_save_leave(PicklerObject *self, PyObject *obj)
1445{
1446 if (self->fast_nesting-- >= FAST_NESTING_LIMIT) {
1447 PyObject *key = PyLong_FromVoidPtr(obj);
1448 if (key == NULL)
1449 return 0;
1450 if (PyDict_DelItem(self->fast_memo, key) < 0) {
1451 Py_DECREF(key);
1452 return 0;
1453 }
1454 Py_DECREF(key);
1455 }
1456 return 1;
1457}
1458
1459static int
1460save_none(PicklerObject *self, PyObject *obj)
1461{
1462 const char none_op = NONE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001463 if (_Pickler_Write(self, &none_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001464 return -1;
1465
1466 return 0;
1467}
1468
1469static int
1470save_bool(PicklerObject *self, PyObject *obj)
1471{
1472 static const char *buf[2] = { FALSE, TRUE };
1473 const char len[2] = {sizeof(FALSE) - 1, sizeof(TRUE) - 1};
1474 int p = (obj == Py_True);
1475
1476 if (self->proto >= 2) {
1477 const char bool_op = p ? NEWTRUE : NEWFALSE;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001478 if (_Pickler_Write(self, &bool_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001479 return -1;
1480 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001481 else if (_Pickler_Write(self, buf[p], len[p]) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001482 return -1;
1483
1484 return 0;
1485}
1486
1487static int
1488save_int(PicklerObject *self, long x)
1489{
1490 char pdata[32];
1491 int len = 0;
1492
1493 if (!self->bin
1494#if SIZEOF_LONG > 4
1495 || x > 0x7fffffffL || x < -0x80000000L
1496#endif
1497 ) {
1498 /* Text-mode pickle, or long too big to fit in the 4-byte
1499 * signed BININT format: store as a string.
1500 */
Mark Dickinson8dd05142009-01-20 20:43:58 +00001501 pdata[0] = LONG; /* use LONG for consistency with pickle.py */
1502 PyOS_snprintf(pdata + 1, sizeof(pdata) - 1, "%ldL\n", x);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001503 if (_Pickler_Write(self, pdata, strlen(pdata)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001504 return -1;
1505 }
1506 else {
1507 /* Binary pickle and x fits in a signed 4-byte int. */
1508 pdata[1] = (unsigned char)(x & 0xff);
1509 pdata[2] = (unsigned char)((x >> 8) & 0xff);
1510 pdata[3] = (unsigned char)((x >> 16) & 0xff);
1511 pdata[4] = (unsigned char)((x >> 24) & 0xff);
1512
1513 if ((pdata[4] == 0) && (pdata[3] == 0)) {
1514 if (pdata[2] == 0) {
1515 pdata[0] = BININT1;
1516 len = 2;
1517 }
1518 else {
1519 pdata[0] = BININT2;
1520 len = 3;
1521 }
1522 }
1523 else {
1524 pdata[0] = BININT;
1525 len = 5;
1526 }
1527
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001528 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001529 return -1;
1530 }
1531
1532 return 0;
1533}
1534
1535static int
1536save_long(PicklerObject *self, PyObject *obj)
1537{
1538 PyObject *repr = NULL;
1539 Py_ssize_t size;
1540 long val = PyLong_AsLong(obj);
1541 int status = 0;
1542
1543 const char long_op = LONG;
1544
1545 if (val == -1 && PyErr_Occurred()) {
1546 /* out of range for int pickling */
1547 PyErr_Clear();
1548 }
1549 else
1550 return save_int(self, val);
1551
1552 if (self->proto >= 2) {
1553 /* Linear-time pickling. */
1554 size_t nbits;
1555 size_t nbytes;
1556 unsigned char *pdata;
1557 char header[5];
1558 int i;
1559 int sign = _PyLong_Sign(obj);
1560
1561 if (sign == 0) {
1562 header[0] = LONG1;
1563 header[1] = 0; /* It's 0 -- an empty bytestring. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001564 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001565 goto error;
1566 return 0;
1567 }
1568 nbits = _PyLong_NumBits(obj);
1569 if (nbits == (size_t)-1 && PyErr_Occurred())
1570 goto error;
1571 /* How many bytes do we need? There are nbits >> 3 full
1572 * bytes of data, and nbits & 7 leftover bits. If there
1573 * are any leftover bits, then we clearly need another
1574 * byte. Wnat's not so obvious is that we *probably*
1575 * need another byte even if there aren't any leftovers:
1576 * the most-significant bit of the most-significant byte
1577 * acts like a sign bit, and it's usually got a sense
1578 * opposite of the one we need. The exception is longs
1579 * of the form -(2**(8*j-1)) for j > 0. Such a long is
1580 * its own 256's-complement, so has the right sign bit
1581 * even without the extra byte. That's a pain to check
1582 * for in advance, though, so we always grab an extra
1583 * byte at the start, and cut it back later if possible.
1584 */
1585 nbytes = (nbits >> 3) + 1;
1586 if (nbytes > INT_MAX) {
1587 PyErr_SetString(PyExc_OverflowError,
1588 "long too large to pickle");
1589 goto error;
1590 }
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001591 repr = PyBytes_FromStringAndSize(NULL, (Py_ssize_t)nbytes);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001592 if (repr == NULL)
1593 goto error;
Neal Norwitz6ae2eb22008-08-24 23:50:08 +00001594 pdata = (unsigned char *)PyBytes_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001595 i = _PyLong_AsByteArray((PyLongObject *)obj,
1596 pdata, nbytes,
1597 1 /* little endian */ , 1 /* signed */ );
1598 if (i < 0)
1599 goto error;
1600 /* If the long is negative, this may be a byte more than
1601 * needed. This is so iff the MSB is all redundant sign
1602 * bits.
1603 */
1604 if (sign < 0 &&
1605 nbytes > 1 &&
1606 pdata[nbytes - 1] == 0xff &&
1607 (pdata[nbytes - 2] & 0x80) != 0) {
1608 nbytes--;
1609 }
1610
1611 if (nbytes < 256) {
1612 header[0] = LONG1;
1613 header[1] = (unsigned char)nbytes;
1614 size = 2;
1615 }
1616 else {
1617 header[0] = LONG4;
1618 size = (int)nbytes;
1619 for (i = 1; i < 5; i++) {
1620 header[i] = (unsigned char)(size & 0xff);
1621 size >>= 8;
1622 }
1623 size = 5;
1624 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001625 if (_Pickler_Write(self, header, size) < 0 ||
1626 _Pickler_Write(self, (char *)pdata, (int)nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001627 goto error;
1628 }
1629 else {
1630 char *string;
1631
Mark Dickinson8dd05142009-01-20 20:43:58 +00001632 /* proto < 2: write the repr and newline. This is quadratic-time (in
1633 the number of digits), in both directions. We add a trailing 'L'
1634 to the repr, for compatibility with Python 2.x. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001635
1636 repr = PyObject_Repr(obj);
1637 if (repr == NULL)
1638 goto error;
1639
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00001640 string = _PyUnicode_AsStringAndSize(repr, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001641 if (string == NULL)
1642 goto error;
1643
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001644 if (_Pickler_Write(self, &long_op, 1) < 0 ||
1645 _Pickler_Write(self, string, size) < 0 ||
1646 _Pickler_Write(self, "L\n", 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001647 goto error;
1648 }
1649
1650 if (0) {
1651 error:
1652 status = -1;
1653 }
1654 Py_XDECREF(repr);
1655
1656 return status;
1657}
1658
1659static int
1660save_float(PicklerObject *self, PyObject *obj)
1661{
1662 double x = PyFloat_AS_DOUBLE((PyFloatObject *)obj);
1663
1664 if (self->bin) {
1665 char pdata[9];
1666 pdata[0] = BINFLOAT;
1667 if (_PyFloat_Pack8(x, (unsigned char *)&pdata[1], 0) < 0)
1668 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001669 if (_Pickler_Write(self, pdata, 9) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001670 return -1;
Eric Smith0923d1d2009-04-16 20:16:10 +00001671 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001672 else {
Eric Smith0923d1d2009-04-16 20:16:10 +00001673 int result = -1;
1674 char *buf = NULL;
1675 char op = FLOAT;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001676
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001677 if (_Pickler_Write(self, &op, 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001678 goto done;
1679
Mark Dickinson3e09f432009-04-17 08:41:23 +00001680 buf = PyOS_double_to_string(x, 'g', 17, 0, NULL);
Eric Smith0923d1d2009-04-16 20:16:10 +00001681 if (!buf) {
1682 PyErr_NoMemory();
1683 goto done;
1684 }
1685
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001686 if (_Pickler_Write(self, buf, strlen(buf)) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001687 goto done;
1688
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001689 if (_Pickler_Write(self, "\n", 1) < 0)
Eric Smith0923d1d2009-04-16 20:16:10 +00001690 goto done;
1691
1692 result = 0;
1693done:
1694 PyMem_Free(buf);
1695 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001696 }
1697
1698 return 0;
1699}
1700
1701static int
1702save_bytes(PicklerObject *self, PyObject *obj)
1703{
1704 if (self->proto < 3) {
1705 /* Older pickle protocols do not have an opcode for pickling bytes
1706 objects. Therefore, we need to fake the copy protocol (i.e.,
1707 the __reduce__ method) to permit bytes object unpickling. */
1708 PyObject *reduce_value = NULL;
1709 PyObject *bytelist = NULL;
1710 int status;
1711
1712 bytelist = PySequence_List(obj);
1713 if (bytelist == NULL)
1714 return -1;
1715
1716 reduce_value = Py_BuildValue("(O(O))", (PyObject *)&PyBytes_Type,
1717 bytelist);
1718 if (reduce_value == NULL) {
1719 Py_DECREF(bytelist);
1720 return -1;
1721 }
1722
1723 /* save_reduce() will memoize the object automatically. */
1724 status = save_reduce(self, reduce_value, obj);
1725 Py_DECREF(reduce_value);
1726 Py_DECREF(bytelist);
1727 return status;
1728 }
1729 else {
1730 Py_ssize_t size;
1731 char header[5];
1732 int len;
1733
1734 size = PyBytes_Size(obj);
1735 if (size < 0)
1736 return -1;
1737
1738 if (size < 256) {
1739 header[0] = SHORT_BINBYTES;
1740 header[1] = (unsigned char)size;
1741 len = 2;
1742 }
1743 else if (size <= 0xffffffffL) {
1744 header[0] = BINBYTES;
1745 header[1] = (unsigned char)(size & 0xff);
1746 header[2] = (unsigned char)((size >> 8) & 0xff);
1747 header[3] = (unsigned char)((size >> 16) & 0xff);
1748 header[4] = (unsigned char)((size >> 24) & 0xff);
1749 len = 5;
1750 }
1751 else {
1752 return -1; /* string too large */
1753 }
1754
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001755 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001756 return -1;
1757
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001758 if (_Pickler_Write(self, PyBytes_AS_STRING(obj), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001759 return -1;
1760
1761 if (memo_put(self, obj) < 0)
1762 return -1;
1763
1764 return 0;
1765 }
1766}
1767
1768/* A copy of PyUnicode_EncodeRawUnicodeEscape() that also translates
1769 backslash and newline characters to \uXXXX escapes. */
1770static PyObject *
1771raw_unicode_escape(const Py_UNICODE *s, Py_ssize_t size)
1772{
1773 PyObject *repr, *result;
1774 char *p;
1775 char *q;
1776
1777 static const char *hexdigits = "0123456789abcdef";
1778
1779#ifdef Py_UNICODE_WIDE
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001780 const Py_ssize_t expandsize = 10;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001781#else
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001782 const Py_ssize_t expandsize = 6;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001783#endif
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001784
1785 if (size > PY_SSIZE_T_MAX / expandsize)
1786 return PyErr_NoMemory();
1787
1788 repr = PyByteArray_FromStringAndSize(NULL, expandsize * size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001789 if (repr == NULL)
1790 return NULL;
1791 if (size == 0)
1792 goto done;
1793
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001794 p = q = PyByteArray_AS_STRING(repr);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001795 while (size-- > 0) {
1796 Py_UNICODE ch = *s++;
1797#ifdef Py_UNICODE_WIDE
1798 /* Map 32-bit characters to '\Uxxxxxxxx' */
1799 if (ch >= 0x10000) {
1800 *p++ = '\\';
1801 *p++ = 'U';
1802 *p++ = hexdigits[(ch >> 28) & 0xf];
1803 *p++ = hexdigits[(ch >> 24) & 0xf];
1804 *p++ = hexdigits[(ch >> 20) & 0xf];
1805 *p++ = hexdigits[(ch >> 16) & 0xf];
1806 *p++ = hexdigits[(ch >> 12) & 0xf];
1807 *p++ = hexdigits[(ch >> 8) & 0xf];
1808 *p++ = hexdigits[(ch >> 4) & 0xf];
1809 *p++ = hexdigits[ch & 15];
1810 }
1811 else
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001812#else
1813 /* Map UTF-16 surrogate pairs to '\U00xxxxxx' */
1814 if (ch >= 0xD800 && ch < 0xDC00) {
1815 Py_UNICODE ch2;
1816 Py_UCS4 ucs;
1817
1818 ch2 = *s++;
1819 size--;
1820 if (ch2 >= 0xDC00 && ch2 <= 0xDFFF) {
1821 ucs = (((ch & 0x03FF) << 10) | (ch2 & 0x03FF)) + 0x00010000;
1822 *p++ = '\\';
1823 *p++ = 'U';
1824 *p++ = hexdigits[(ucs >> 28) & 0xf];
1825 *p++ = hexdigits[(ucs >> 24) & 0xf];
1826 *p++ = hexdigits[(ucs >> 20) & 0xf];
1827 *p++ = hexdigits[(ucs >> 16) & 0xf];
1828 *p++ = hexdigits[(ucs >> 12) & 0xf];
1829 *p++ = hexdigits[(ucs >> 8) & 0xf];
1830 *p++ = hexdigits[(ucs >> 4) & 0xf];
1831 *p++ = hexdigits[ucs & 0xf];
1832 continue;
1833 }
1834 /* Fall through: isolated surrogates are copied as-is */
1835 s--;
1836 size++;
1837 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001838#endif
1839 /* Map 16-bit characters to '\uxxxx' */
1840 if (ch >= 256 || ch == '\\' || ch == '\n') {
1841 *p++ = '\\';
1842 *p++ = 'u';
1843 *p++ = hexdigits[(ch >> 12) & 0xf];
1844 *p++ = hexdigits[(ch >> 8) & 0xf];
1845 *p++ = hexdigits[(ch >> 4) & 0xf];
1846 *p++ = hexdigits[ch & 15];
1847 }
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001848 /* Copy everything else as-is */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001849 else
1850 *p++ = (char) ch;
1851 }
1852 size = p - q;
1853
1854 done:
Alexandre Vassalotti554d8782008-12-27 07:32:41 +00001855 result = PyBytes_FromStringAndSize(PyByteArray_AS_STRING(repr), size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001856 Py_DECREF(repr);
1857 return result;
1858}
1859
1860static int
1861save_unicode(PicklerObject *self, PyObject *obj)
1862{
1863 Py_ssize_t size;
1864 PyObject *encoded = NULL;
1865
1866 if (self->bin) {
1867 char pdata[5];
1868
Victor Stinner485fb562010-04-13 11:07:24 +00001869 encoded = PyUnicode_EncodeUTF8(PyUnicode_AS_UNICODE(obj),
1870 PyUnicode_GET_SIZE(obj),
1871 "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001872 if (encoded == NULL)
1873 goto error;
1874
1875 size = PyBytes_GET_SIZE(encoded);
1876 if (size < 0 || size > 0xffffffffL)
1877 goto error; /* string too large */
1878
1879 pdata[0] = BINUNICODE;
1880 pdata[1] = (unsigned char)(size & 0xff);
1881 pdata[2] = (unsigned char)((size >> 8) & 0xff);
1882 pdata[3] = (unsigned char)((size >> 16) & 0xff);
1883 pdata[4] = (unsigned char)((size >> 24) & 0xff);
1884
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001885 if (_Pickler_Write(self, pdata, 5) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001886 goto error;
1887
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001888 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001889 goto error;
1890 }
1891 else {
1892 const char unicode_op = UNICODE;
1893
1894 encoded = raw_unicode_escape(PyUnicode_AS_UNICODE(obj),
1895 PyUnicode_GET_SIZE(obj));
1896 if (encoded == NULL)
1897 goto error;
1898
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001899 if (_Pickler_Write(self, &unicode_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001900 goto error;
1901
1902 size = PyBytes_GET_SIZE(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001903 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded), size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001904 goto error;
1905
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001906 if (_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001907 goto error;
1908 }
1909 if (memo_put(self, obj) < 0)
1910 goto error;
1911
1912 Py_DECREF(encoded);
1913 return 0;
1914
1915 error:
1916 Py_XDECREF(encoded);
1917 return -1;
1918}
1919
1920/* A helper for save_tuple. Push the len elements in tuple t on the stack. */
1921static int
1922store_tuple_elements(PicklerObject *self, PyObject *t, int len)
1923{
1924 int i;
1925
1926 assert(PyTuple_Size(t) == len);
1927
1928 for (i = 0; i < len; i++) {
1929 PyObject *element = PyTuple_GET_ITEM(t, i);
1930
1931 if (element == NULL)
1932 return -1;
1933 if (save(self, element, 0) < 0)
1934 return -1;
1935 }
1936
1937 return 0;
1938}
1939
1940/* Tuples are ubiquitous in the pickle protocols, so many techniques are
1941 * used across protocols to minimize the space needed to pickle them.
1942 * Tuples are also the only builtin immutable type that can be recursive
1943 * (a tuple can be reached from itself), and that requires some subtle
1944 * magic so that it works in all cases. IOW, this is a long routine.
1945 */
1946static int
1947save_tuple(PicklerObject *self, PyObject *obj)
1948{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001949 int len, i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001950
1951 const char mark_op = MARK;
1952 const char tuple_op = TUPLE;
1953 const char pop_op = POP;
1954 const char pop_mark_op = POP_MARK;
1955 const char len2opcode[] = {EMPTY_TUPLE, TUPLE1, TUPLE2, TUPLE3};
1956
1957 if ((len = PyTuple_Size(obj)) < 0)
1958 return -1;
1959
1960 if (len == 0) {
1961 char pdata[2];
1962
1963 if (self->proto) {
1964 pdata[0] = EMPTY_TUPLE;
1965 len = 1;
1966 }
1967 else {
1968 pdata[0] = MARK;
1969 pdata[1] = TUPLE;
1970 len = 2;
1971 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001972 if (_Pickler_Write(self, pdata, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001973 return -1;
1974 return 0;
1975 }
1976
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001977 /* The tuple isn't in the memo now. If it shows up there after
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001978 * saving the tuple elements, the tuple must be recursive, in
1979 * which case we'll pop everything we put on the stack, and fetch
1980 * its value from the memo.
1981 */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001982 if (len <= 3 && self->proto >= 2) {
1983 /* Use TUPLE{1,2,3} opcodes. */
1984 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001985 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001986
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001987 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001988 /* pop the len elements */
1989 for (i = 0; i < len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001990 if (_Pickler_Write(self, &pop_op, 1) < 0)
1991 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001992 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001993 if (memo_get(self, obj) < 0)
1994 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001995
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00001996 return 0;
1997 }
1998 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00001999 if (_Pickler_Write(self, len2opcode + len, 1) < 0)
2000 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002001 }
2002 goto memoize;
2003 }
2004
2005 /* proto < 2 and len > 0, or proto >= 2 and len > 3.
2006 * Generate MARK e1 e2 ... TUPLE
2007 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002008 if (_Pickler_Write(self, &mark_op, 1) < 0)
2009 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002010
2011 if (store_tuple_elements(self, obj, len) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002012 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002013
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002014 if (PyMemoTable_Get(self->memo, obj)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002015 /* pop the stack stuff we pushed */
2016 if (self->bin) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002017 if (_Pickler_Write(self, &pop_mark_op, 1) < 0)
2018 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002019 }
2020 else {
2021 /* Note that we pop one more than len, to remove
2022 * the MARK too.
2023 */
2024 for (i = 0; i <= len; i++)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002025 if (_Pickler_Write(self, &pop_op, 1) < 0)
2026 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002027 }
2028 /* fetch from memo */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002029 if (memo_get(self, obj) < 0)
2030 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002031
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002032 return 0;
2033 }
2034 else { /* Not recursive. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002035 if (_Pickler_Write(self, &tuple_op, 1) < 0)
2036 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002037 }
2038
2039 memoize:
2040 if (memo_put(self, obj) < 0)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002041 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002042
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002043 return 0;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002044}
2045
2046/* iter is an iterator giving items, and we batch up chunks of
2047 * MARK item item ... item APPENDS
2048 * opcode sequences. Calling code should have arranged to first create an
2049 * empty list, or list-like object, for the APPENDS to operate on.
2050 * Returns 0 on success, <0 on error.
2051 */
2052static int
2053batch_list(PicklerObject *self, PyObject *iter)
2054{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002055 PyObject *obj = NULL;
2056 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002057 int i, n;
2058
2059 const char mark_op = MARK;
2060 const char append_op = APPEND;
2061 const char appends_op = APPENDS;
2062
2063 assert(iter != NULL);
2064
2065 /* XXX: I think this function could be made faster by avoiding the
2066 iterator interface and fetching objects directly from list using
2067 PyList_GET_ITEM.
2068 */
2069
2070 if (self->proto == 0) {
2071 /* APPENDS isn't available; do one at a time. */
2072 for (;;) {
2073 obj = PyIter_Next(iter);
2074 if (obj == NULL) {
2075 if (PyErr_Occurred())
2076 return -1;
2077 break;
2078 }
2079 i = save(self, obj, 0);
2080 Py_DECREF(obj);
2081 if (i < 0)
2082 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002083 if (_Pickler_Write(self, &append_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002084 return -1;
2085 }
2086 return 0;
2087 }
2088
2089 /* proto > 0: write in batches of BATCHSIZE. */
2090 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002091 /* Get first item */
2092 firstitem = PyIter_Next(iter);
2093 if (firstitem == NULL) {
2094 if (PyErr_Occurred())
2095 goto error;
2096
2097 /* nothing more to add */
2098 break;
2099 }
2100
2101 /* Try to get a second item */
2102 obj = PyIter_Next(iter);
2103 if (obj == NULL) {
2104 if (PyErr_Occurred())
2105 goto error;
2106
2107 /* Only one item to write */
2108 if (save(self, firstitem, 0) < 0)
2109 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002110 if (_Pickler_Write(self, &append_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002111 goto error;
2112 Py_CLEAR(firstitem);
2113 break;
2114 }
2115
2116 /* More than one item to write */
2117
2118 /* Pump out MARK, items, APPENDS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002119 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002120 goto error;
2121
2122 if (save(self, firstitem, 0) < 0)
2123 goto error;
2124 Py_CLEAR(firstitem);
2125 n = 1;
2126
2127 /* Fetch and save up to BATCHSIZE items */
2128 while (obj) {
2129 if (save(self, obj, 0) < 0)
2130 goto error;
2131 Py_CLEAR(obj);
2132 n += 1;
2133
2134 if (n == BATCHSIZE)
2135 break;
2136
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002137 obj = PyIter_Next(iter);
2138 if (obj == NULL) {
2139 if (PyErr_Occurred())
2140 goto error;
2141 break;
2142 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002143 }
2144
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002145 if (_Pickler_Write(self, &appends_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002146 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002147
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002148 } while (n == BATCHSIZE);
2149 return 0;
2150
2151 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002152 Py_XDECREF(firstitem);
2153 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002154 return -1;
2155}
2156
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002157/* This is a variant of batch_list() above, specialized for lists (with no
2158 * support for list subclasses). Like batch_list(), we batch up chunks of
2159 * MARK item item ... item APPENDS
2160 * opcode sequences. Calling code should have arranged to first create an
2161 * empty list, or list-like object, for the APPENDS to operate on.
2162 * Returns 0 on success, -1 on error.
2163 *
2164 * This version is considerably faster than batch_list(), if less general.
2165 *
2166 * Note that this only works for protocols > 0.
2167 */
2168static int
2169batch_list_exact(PicklerObject *self, PyObject *obj)
2170{
2171 PyObject *item = NULL;
2172 int this_batch, total;
2173
2174 const char append_op = APPEND;
2175 const char appends_op = APPENDS;
2176 const char mark_op = MARK;
2177
2178 assert(obj != NULL);
2179 assert(self->proto > 0);
2180 assert(PyList_CheckExact(obj));
2181
2182 if (PyList_GET_SIZE(obj) == 1) {
2183 item = PyList_GET_ITEM(obj, 0);
2184 if (save(self, item, 0) < 0)
2185 return -1;
2186 if (_Pickler_Write(self, &append_op, 1) < 0)
2187 return -1;
2188 return 0;
2189 }
2190
2191 /* Write in batches of BATCHSIZE. */
2192 total = 0;
2193 do {
2194 this_batch = 0;
2195 if (_Pickler_Write(self, &mark_op, 1) < 0)
2196 return -1;
2197 while (total < PyList_GET_SIZE(obj)) {
2198 item = PyList_GET_ITEM(obj, total);
2199 if (save(self, item, 0) < 0)
2200 return -1;
2201 total++;
2202 if (++this_batch == BATCHSIZE)
2203 break;
2204 }
2205 if (_Pickler_Write(self, &appends_op, 1) < 0)
2206 return -1;
2207
2208 } while (total < PyList_GET_SIZE(obj));
2209
2210 return 0;
2211}
2212
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002213static int
2214save_list(PicklerObject *self, PyObject *obj)
2215{
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002216 char header[3];
2217 int len;
2218 int status = 0;
2219
2220 if (self->fast && !fast_save_enter(self, obj))
2221 goto error;
2222
2223 /* Create an empty list. */
2224 if (self->bin) {
2225 header[0] = EMPTY_LIST;
2226 len = 1;
2227 }
2228 else {
2229 header[0] = MARK;
2230 header[1] = LIST;
2231 len = 2;
2232 }
2233
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002234 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002235 goto error;
2236
2237 /* Get list length, and bow out early if empty. */
2238 if ((len = PyList_Size(obj)) < 0)
2239 goto error;
2240
2241 if (memo_put(self, obj) < 0)
2242 goto error;
2243
2244 if (len != 0) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002245 /* Materialize the list elements. */
2246 if (PyList_CheckExact(obj) && self->proto > 0) {
2247 if (Py_EnterRecursiveCall(" while pickling an object") == 0) {
2248 status = batch_list_exact(self, obj);
2249 Py_LeaveRecursiveCall();
2250 }
2251 } else {
2252 PyObject *iter = PyObject_GetIter(obj);
2253 if (iter == NULL)
2254 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002255
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002256 if (Py_EnterRecursiveCall(" while pickling an object") == 0) {
2257 status = batch_list(self, iter);
2258 Py_LeaveRecursiveCall();
2259 }
2260 Py_DECREF(iter);
2261 }
2262 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002263 if (0) {
2264 error:
2265 status = -1;
2266 }
2267
2268 if (self->fast && !fast_save_leave(self, obj))
2269 status = -1;
2270
2271 return status;
2272}
2273
2274/* iter is an iterator giving (key, value) pairs, and we batch up chunks of
2275 * MARK key value ... key value SETITEMS
2276 * opcode sequences. Calling code should have arranged to first create an
2277 * empty dict, or dict-like object, for the SETITEMS to operate on.
2278 * Returns 0 on success, <0 on error.
2279 *
2280 * This is very much like batch_list(). The difference between saving
2281 * elements directly, and picking apart two-tuples, is so long-winded at
2282 * the C level, though, that attempts to combine these routines were too
2283 * ugly to bear.
2284 */
2285static int
2286batch_dict(PicklerObject *self, PyObject *iter)
2287{
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002288 PyObject *obj = NULL;
2289 PyObject *firstitem = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002290 int i, n;
2291
2292 const char mark_op = MARK;
2293 const char setitem_op = SETITEM;
2294 const char setitems_op = SETITEMS;
2295
2296 assert(iter != NULL);
2297
2298 if (self->proto == 0) {
2299 /* SETITEMS isn't available; do one at a time. */
2300 for (;;) {
2301 obj = PyIter_Next(iter);
2302 if (obj == NULL) {
2303 if (PyErr_Occurred())
2304 return -1;
2305 break;
2306 }
2307 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2308 PyErr_SetString(PyExc_TypeError, "dict items "
2309 "iterator must return 2-tuples");
2310 return -1;
2311 }
2312 i = save(self, PyTuple_GET_ITEM(obj, 0), 0);
2313 if (i >= 0)
2314 i = save(self, PyTuple_GET_ITEM(obj, 1), 0);
2315 Py_DECREF(obj);
2316 if (i < 0)
2317 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002318 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002319 return -1;
2320 }
2321 return 0;
2322 }
2323
2324 /* proto > 0: write in batches of BATCHSIZE. */
2325 do {
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002326 /* Get first item */
2327 firstitem = PyIter_Next(iter);
2328 if (firstitem == NULL) {
2329 if (PyErr_Occurred())
2330 goto error;
2331
2332 /* nothing more to add */
2333 break;
2334 }
2335 if (!PyTuple_Check(firstitem) || PyTuple_Size(firstitem) != 2) {
2336 PyErr_SetString(PyExc_TypeError, "dict items "
2337 "iterator must return 2-tuples");
2338 goto error;
2339 }
2340
2341 /* Try to get a second item */
2342 obj = PyIter_Next(iter);
2343 if (obj == NULL) {
2344 if (PyErr_Occurred())
2345 goto error;
2346
2347 /* Only one item to write */
2348 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2349 goto error;
2350 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2351 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002352 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002353 goto error;
2354 Py_CLEAR(firstitem);
2355 break;
2356 }
2357
2358 /* More than one item to write */
2359
2360 /* Pump out MARK, items, SETITEMS. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002361 if (_Pickler_Write(self, &mark_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002362 goto error;
2363
2364 if (save(self, PyTuple_GET_ITEM(firstitem, 0), 0) < 0)
2365 goto error;
2366 if (save(self, PyTuple_GET_ITEM(firstitem, 1), 0) < 0)
2367 goto error;
2368 Py_CLEAR(firstitem);
2369 n = 1;
2370
2371 /* Fetch and save up to BATCHSIZE items */
2372 while (obj) {
2373 if (!PyTuple_Check(obj) || PyTuple_Size(obj) != 2) {
2374 PyErr_SetString(PyExc_TypeError, "dict items "
2375 "iterator must return 2-tuples");
2376 goto error;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002377 }
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002378 if (save(self, PyTuple_GET_ITEM(obj, 0), 0) < 0 ||
2379 save(self, PyTuple_GET_ITEM(obj, 1), 0) < 0)
2380 goto error;
2381 Py_CLEAR(obj);
2382 n += 1;
2383
2384 if (n == BATCHSIZE)
2385 break;
2386
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002387 obj = PyIter_Next(iter);
2388 if (obj == NULL) {
2389 if (PyErr_Occurred())
2390 goto error;
2391 break;
2392 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002393 }
2394
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002395 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002396 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002397
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002398 } while (n == BATCHSIZE);
2399 return 0;
2400
2401 error:
Amaury Forgeot d'Arcfb1a5eb2008-09-11 21:03:37 +00002402 Py_XDECREF(firstitem);
2403 Py_XDECREF(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002404 return -1;
2405}
2406
Collin Winter5c9b02d2009-05-25 05:43:30 +00002407/* This is a variant of batch_dict() above that specializes for dicts, with no
2408 * support for dict subclasses. Like batch_dict(), we batch up chunks of
2409 * MARK key value ... key value SETITEMS
2410 * opcode sequences. Calling code should have arranged to first create an
2411 * empty dict, or dict-like object, for the SETITEMS to operate on.
2412 * Returns 0 on success, -1 on error.
2413 *
2414 * Note that this currently doesn't work for protocol 0.
2415 */
2416static int
2417batch_dict_exact(PicklerObject *self, PyObject *obj)
2418{
2419 PyObject *key = NULL, *value = NULL;
2420 int i;
2421 Py_ssize_t dict_size, ppos = 0;
2422
Alexandre Vassalottif70b1292009-05-25 18:00:52 +00002423 const char mark_op = MARK;
2424 const char setitem_op = SETITEM;
2425 const char setitems_op = SETITEMS;
Collin Winter5c9b02d2009-05-25 05:43:30 +00002426
2427 assert(obj != NULL);
2428 assert(self->proto > 0);
2429
2430 dict_size = PyDict_Size(obj);
2431
2432 /* Special-case len(d) == 1 to save space. */
2433 if (dict_size == 1) {
2434 PyDict_Next(obj, &ppos, &key, &value);
2435 if (save(self, key, 0) < 0)
2436 return -1;
2437 if (save(self, value, 0) < 0)
2438 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002439 if (_Pickler_Write(self, &setitem_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002440 return -1;
2441 return 0;
2442 }
2443
2444 /* Write in batches of BATCHSIZE. */
2445 do {
2446 i = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002447 if (_Pickler_Write(self, &mark_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002448 return -1;
2449 while (PyDict_Next(obj, &ppos, &key, &value)) {
2450 if (save(self, key, 0) < 0)
2451 return -1;
2452 if (save(self, value, 0) < 0)
2453 return -1;
2454 if (++i == BATCHSIZE)
2455 break;
2456 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002457 if (_Pickler_Write(self, &setitems_op, 1) < 0)
Collin Winter5c9b02d2009-05-25 05:43:30 +00002458 return -1;
2459 if (PyDict_Size(obj) != dict_size) {
2460 PyErr_Format(
2461 PyExc_RuntimeError,
2462 "dictionary changed size during iteration");
2463 return -1;
2464 }
2465
2466 } while (i == BATCHSIZE);
2467 return 0;
2468}
2469
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002470static int
2471save_dict(PicklerObject *self, PyObject *obj)
2472{
2473 PyObject *items, *iter;
2474 char header[3];
2475 int len;
2476 int status = 0;
2477
2478 if (self->fast && !fast_save_enter(self, obj))
2479 goto error;
2480
2481 /* Create an empty dict. */
2482 if (self->bin) {
2483 header[0] = EMPTY_DICT;
2484 len = 1;
2485 }
2486 else {
2487 header[0] = MARK;
2488 header[1] = DICT;
2489 len = 2;
2490 }
2491
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002492 if (_Pickler_Write(self, header, len) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002493 goto error;
2494
2495 /* Get dict size, and bow out early if empty. */
2496 if ((len = PyDict_Size(obj)) < 0)
2497 goto error;
2498
2499 if (memo_put(self, obj) < 0)
2500 goto error;
2501
2502 if (len != 0) {
2503 /* Save the dict items. */
Collin Winter5c9b02d2009-05-25 05:43:30 +00002504 if (PyDict_CheckExact(obj) && self->proto > 0) {
2505 /* We can take certain shortcuts if we know this is a dict and
2506 not a dict subclass. */
2507 if (Py_EnterRecursiveCall(" while pickling an object") == 0) {
2508 status = batch_dict_exact(self, obj);
2509 Py_LeaveRecursiveCall();
2510 }
2511 } else {
2512 items = PyObject_CallMethod(obj, "items", "()");
2513 if (items == NULL)
2514 goto error;
2515 iter = PyObject_GetIter(items);
2516 Py_DECREF(items);
2517 if (iter == NULL)
2518 goto error;
2519 status = batch_dict(self, iter);
2520 Py_DECREF(iter);
2521 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002522 }
2523
2524 if (0) {
2525 error:
2526 status = -1;
2527 }
2528
2529 if (self->fast && !fast_save_leave(self, obj))
2530 status = -1;
2531
2532 return status;
2533}
2534
2535static int
2536save_global(PicklerObject *self, PyObject *obj, PyObject *name)
2537{
2538 static PyObject *name_str = NULL;
2539 PyObject *global_name = NULL;
2540 PyObject *module_name = NULL;
2541 PyObject *module = NULL;
2542 PyObject *cls;
2543 int status = 0;
2544
2545 const char global_op = GLOBAL;
2546
2547 if (name_str == NULL) {
2548 name_str = PyUnicode_InternFromString("__name__");
2549 if (name_str == NULL)
2550 goto error;
2551 }
2552
2553 if (name) {
2554 global_name = name;
2555 Py_INCREF(global_name);
2556 }
2557 else {
2558 global_name = PyObject_GetAttr(obj, name_str);
2559 if (global_name == NULL)
2560 goto error;
2561 }
2562
2563 module_name = whichmodule(obj, global_name);
2564 if (module_name == NULL)
2565 goto error;
2566
2567 /* XXX: Change to use the import C API directly with level=0 to disallow
2568 relative imports.
2569
2570 XXX: PyImport_ImportModuleLevel could be used. However, this bypasses
2571 builtins.__import__. Therefore, _pickle, unlike pickle.py, will ignore
2572 custom import functions (IMHO, this would be a nice security
2573 feature). The import C API would need to be extended to support the
2574 extra parameters of __import__ to fix that. */
2575 module = PyImport_Import(module_name);
2576 if (module == NULL) {
2577 PyErr_Format(PicklingError,
2578 "Can't pickle %R: import of module %R failed",
2579 obj, module_name);
2580 goto error;
2581 }
2582 cls = PyObject_GetAttr(module, global_name);
2583 if (cls == NULL) {
2584 PyErr_Format(PicklingError,
2585 "Can't pickle %R: attribute lookup %S.%S failed",
2586 obj, module_name, global_name);
2587 goto error;
2588 }
2589 if (cls != obj) {
2590 Py_DECREF(cls);
2591 PyErr_Format(PicklingError,
2592 "Can't pickle %R: it's not the same object as %S.%S",
2593 obj, module_name, global_name);
2594 goto error;
2595 }
2596 Py_DECREF(cls);
2597
2598 if (self->proto >= 2) {
2599 /* See whether this is in the extension registry, and if
2600 * so generate an EXT opcode.
2601 */
2602 PyObject *code_obj; /* extension code as Python object */
2603 long code; /* extension code as C value */
2604 char pdata[5];
2605 int n;
2606
2607 PyTuple_SET_ITEM(two_tuple, 0, module_name);
2608 PyTuple_SET_ITEM(two_tuple, 1, global_name);
2609 code_obj = PyDict_GetItem(extension_registry, two_tuple);
2610 /* The object is not registered in the extension registry.
2611 This is the most likely code path. */
2612 if (code_obj == NULL)
2613 goto gen_global;
2614
2615 /* XXX: pickle.py doesn't check neither the type, nor the range
2616 of the value returned by the extension_registry. It should for
2617 consistency. */
2618
2619 /* Verify code_obj has the right type and value. */
2620 if (!PyLong_Check(code_obj)) {
2621 PyErr_Format(PicklingError,
2622 "Can't pickle %R: extension code %R isn't an integer",
2623 obj, code_obj);
2624 goto error;
2625 }
2626 code = PyLong_AS_LONG(code_obj);
2627 if (code <= 0 || code > 0x7fffffffL) {
2628 PyErr_Format(PicklingError,
2629 "Can't pickle %R: extension code %ld is out of range",
2630 obj, code);
2631 goto error;
2632 }
2633
2634 /* Generate an EXT opcode. */
2635 if (code <= 0xff) {
2636 pdata[0] = EXT1;
2637 pdata[1] = (unsigned char)code;
2638 n = 2;
2639 }
2640 else if (code <= 0xffff) {
2641 pdata[0] = EXT2;
2642 pdata[1] = (unsigned char)(code & 0xff);
2643 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2644 n = 3;
2645 }
2646 else {
2647 pdata[0] = EXT4;
2648 pdata[1] = (unsigned char)(code & 0xff);
2649 pdata[2] = (unsigned char)((code >> 8) & 0xff);
2650 pdata[3] = (unsigned char)((code >> 16) & 0xff);
2651 pdata[4] = (unsigned char)((code >> 24) & 0xff);
2652 n = 5;
2653 }
2654
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002655 if (_Pickler_Write(self, pdata, n) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002656 goto error;
2657 }
2658 else {
2659 /* Generate a normal global opcode if we are using a pickle
2660 protocol <= 2, or if the object is not registered in the
2661 extension registry. */
2662 PyObject *encoded;
2663 PyObject *(*unicode_encoder)(PyObject *);
2664
2665 gen_global:
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002666 if (_Pickler_Write(self, &global_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002667 goto error;
2668
2669 /* Since Python 3.0 now supports non-ASCII identifiers, we encode both
2670 the module name and the global name using UTF-8. We do so only when
2671 we are using the pickle protocol newer than version 3. This is to
2672 ensure compatibility with older Unpickler running on Python 2.x. */
2673 if (self->proto >= 3) {
2674 unicode_encoder = PyUnicode_AsUTF8String;
2675 }
2676 else {
2677 unicode_encoder = PyUnicode_AsASCIIString;
2678 }
2679
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00002680 /* For protocol < 3 and if the user didn't request against doing so,
2681 we convert module names to the old 2.x module names. */
2682 if (self->fix_imports) {
2683 PyObject *key;
2684 PyObject *item;
2685
2686 key = PyTuple_Pack(2, module_name, global_name);
2687 if (key == NULL)
2688 goto error;
2689 item = PyDict_GetItemWithError(name_mapping_3to2, key);
2690 Py_DECREF(key);
2691 if (item) {
2692 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
2693 PyErr_Format(PyExc_RuntimeError,
2694 "_compat_pickle.REVERSE_NAME_MAPPING values "
2695 "should be 2-tuples, not %.200s",
2696 Py_TYPE(item)->tp_name);
2697 goto error;
2698 }
2699 Py_CLEAR(module_name);
2700 Py_CLEAR(global_name);
2701 module_name = PyTuple_GET_ITEM(item, 0);
2702 global_name = PyTuple_GET_ITEM(item, 1);
2703 if (!PyUnicode_Check(module_name) ||
2704 !PyUnicode_Check(global_name)) {
2705 PyErr_Format(PyExc_RuntimeError,
2706 "_compat_pickle.REVERSE_NAME_MAPPING values "
2707 "should be pairs of str, not (%.200s, %.200s)",
2708 Py_TYPE(module_name)->tp_name,
2709 Py_TYPE(global_name)->tp_name);
2710 goto error;
2711 }
2712 Py_INCREF(module_name);
2713 Py_INCREF(global_name);
2714 }
2715 else if (PyErr_Occurred()) {
2716 goto error;
2717 }
2718
2719 item = PyDict_GetItemWithError(import_mapping_3to2, module_name);
2720 if (item) {
2721 if (!PyUnicode_Check(item)) {
2722 PyErr_Format(PyExc_RuntimeError,
2723 "_compat_pickle.REVERSE_IMPORT_MAPPING values "
2724 "should be strings, not %.200s",
2725 Py_TYPE(item)->tp_name);
2726 goto error;
2727 }
2728 Py_CLEAR(module_name);
2729 module_name = item;
2730 Py_INCREF(module_name);
2731 }
2732 else if (PyErr_Occurred()) {
2733 goto error;
2734 }
2735 }
2736
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002737 /* Save the name of the module. */
2738 encoded = unicode_encoder(module_name);
2739 if (encoded == NULL) {
2740 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2741 PyErr_Format(PicklingError,
2742 "can't pickle module identifier '%S' using "
2743 "pickle protocol %i", module_name, self->proto);
2744 goto error;
2745 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002746 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002747 PyBytes_GET_SIZE(encoded)) < 0) {
2748 Py_DECREF(encoded);
2749 goto error;
2750 }
2751 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002752 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002753 goto error;
2754
2755 /* Save the name of the module. */
2756 encoded = unicode_encoder(global_name);
2757 if (encoded == NULL) {
2758 if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError))
2759 PyErr_Format(PicklingError,
2760 "can't pickle global identifier '%S' using "
2761 "pickle protocol %i", global_name, self->proto);
2762 goto error;
2763 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002764 if (_Pickler_Write(self, PyBytes_AS_STRING(encoded),
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002765 PyBytes_GET_SIZE(encoded)) < 0) {
2766 Py_DECREF(encoded);
2767 goto error;
2768 }
2769 Py_DECREF(encoded);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002770 if(_Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002771 goto error;
2772
2773 /* Memoize the object. */
2774 if (memo_put(self, obj) < 0)
2775 goto error;
2776 }
2777
2778 if (0) {
2779 error:
2780 status = -1;
2781 }
2782 Py_XDECREF(module_name);
2783 Py_XDECREF(global_name);
2784 Py_XDECREF(module);
2785
2786 return status;
2787}
2788
2789static int
2790save_pers(PicklerObject *self, PyObject *obj, PyObject *func)
2791{
2792 PyObject *pid = NULL;
2793 int status = 0;
2794
2795 const char persid_op = PERSID;
2796 const char binpersid_op = BINPERSID;
2797
2798 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002799 pid = _Pickler_FastCall(self, func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002800 if (pid == NULL)
2801 return -1;
2802
2803 if (pid != Py_None) {
2804 if (self->bin) {
2805 if (save(self, pid, 1) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002806 _Pickler_Write(self, &binpersid_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002807 goto error;
2808 }
2809 else {
2810 PyObject *pid_str = NULL;
2811 char *pid_ascii_bytes;
2812 Py_ssize_t size;
2813
2814 pid_str = PyObject_Str(pid);
2815 if (pid_str == NULL)
2816 goto error;
2817
2818 /* XXX: Should it check whether the persistent id only contains
2819 ASCII characters? And what if the pid contains embedded
2820 newlines? */
Marc-André Lemburg4cc0f242008-08-07 18:54:33 +00002821 pid_ascii_bytes = _PyUnicode_AsStringAndSize(pid_str, &size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002822 Py_DECREF(pid_str);
2823 if (pid_ascii_bytes == NULL)
2824 goto error;
2825
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002826 if (_Pickler_Write(self, &persid_op, 1) < 0 ||
2827 _Pickler_Write(self, pid_ascii_bytes, size) < 0 ||
2828 _Pickler_Write(self, "\n", 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002829 goto error;
2830 }
2831 status = 1;
2832 }
2833
2834 if (0) {
2835 error:
2836 status = -1;
2837 }
2838 Py_XDECREF(pid);
2839
2840 return status;
2841}
2842
2843/* We're saving obj, and args is the 2-thru-5 tuple returned by the
2844 * appropriate __reduce__ method for obj.
2845 */
2846static int
2847save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
2848{
2849 PyObject *callable;
2850 PyObject *argtup;
2851 PyObject *state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002852 PyObject *listitems = Py_None;
2853 PyObject *dictitems = Py_None;
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002854 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002855
2856 int use_newobj = self->proto >= 2;
2857
2858 const char reduce_op = REDUCE;
2859 const char build_op = BUILD;
2860 const char newobj_op = NEWOBJ;
2861
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002862 size = PyTuple_Size(args);
2863 if (size < 2 || size > 5) {
2864 PyErr_SetString(PicklingError, "tuple returned by "
2865 "__reduce__ must contain 2 through 5 elements");
2866 return -1;
2867 }
2868
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002869 if (!PyArg_UnpackTuple(args, "save_reduce", 2, 5,
2870 &callable, &argtup, &state, &listitems, &dictitems))
2871 return -1;
2872
2873 if (!PyCallable_Check(callable)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002874 PyErr_SetString(PicklingError, "first item of the tuple "
2875 "returned by __reduce__ must be callable");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002876 return -1;
2877 }
2878 if (!PyTuple_Check(argtup)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002879 PyErr_SetString(PicklingError, "second item of the tuple "
2880 "returned by __reduce__ must be a tuple");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002881 return -1;
2882 }
2883
2884 if (state == Py_None)
2885 state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002886
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002887 if (listitems == Py_None)
2888 listitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002889 else if (!PyIter_Check(listitems)) {
2890 PyErr_Format(PicklingError, "Fourth element of tuple"
2891 "returned by __reduce__ must be an iterator, not %s",
2892 Py_TYPE(listitems)->tp_name);
2893 return -1;
2894 }
2895
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002896 if (dictitems == Py_None)
2897 dictitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002898 else if (!PyIter_Check(dictitems)) {
2899 PyErr_Format(PicklingError, "Fifth element of tuple"
2900 "returned by __reduce__ must be an iterator, not %s",
2901 Py_TYPE(dictitems)->tp_name);
2902 return -1;
2903 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002904
2905 /* Protocol 2 special case: if callable's name is __newobj__, use
2906 NEWOBJ. */
2907 if (use_newobj) {
Antoine Pitrouff150f22010-10-22 21:41:05 +00002908 static PyObject *newobj_str = NULL;
2909 PyObject *name_str;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002910
2911 if (newobj_str == NULL) {
2912 newobj_str = PyUnicode_InternFromString("__newobj__");
Antoine Pitrouff150f22010-10-22 21:41:05 +00002913 if (newobj_str == NULL)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002914 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002915 }
2916
Antoine Pitrouff150f22010-10-22 21:41:05 +00002917 name_str = PyObject_GetAttrString(callable, "__name__");
2918 if (name_str == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002919 if (PyErr_ExceptionMatches(PyExc_AttributeError))
2920 PyErr_Clear();
2921 else
2922 return -1;
2923 use_newobj = 0;
2924 }
2925 else {
Antoine Pitrouff150f22010-10-22 21:41:05 +00002926 use_newobj = PyUnicode_Check(name_str) &&
2927 PyUnicode_Compare(name_str, newobj_str) == 0;
2928 Py_DECREF(name_str);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002929 }
2930 }
2931 if (use_newobj) {
2932 PyObject *cls;
2933 PyObject *newargtup;
2934 PyObject *obj_class;
2935 int p;
2936
2937 /* Sanity checks. */
2938 if (Py_SIZE(argtup) < 1) {
2939 PyErr_SetString(PicklingError, "__newobj__ arglist is empty");
2940 return -1;
2941 }
2942
2943 cls = PyTuple_GET_ITEM(argtup, 0);
Antoine Pitrouff150f22010-10-22 21:41:05 +00002944 if (!PyObject_HasAttrString(cls, "__new__")) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002945 PyErr_SetString(PicklingError, "args[0] from "
Antoine Pitrouff150f22010-10-22 21:41:05 +00002946 "__newobj__ args has no __new__");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002947 return -1;
2948 }
2949
2950 if (obj != NULL) {
Antoine Pitrouff150f22010-10-22 21:41:05 +00002951 obj_class = PyObject_GetAttrString(obj, "__class__");
2952 if (obj_class == NULL) {
2953 if (PyErr_ExceptionMatches(PyExc_AttributeError))
2954 PyErr_Clear();
2955 else
2956 return -1;
2957 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002958 p = obj_class != cls; /* true iff a problem */
2959 Py_DECREF(obj_class);
2960 if (p) {
2961 PyErr_SetString(PicklingError, "args[0] from "
2962 "__newobj__ args has the wrong class");
2963 return -1;
2964 }
2965 }
2966 /* XXX: These calls save() are prone to infinite recursion. Imagine
2967 what happen if the value returned by the __reduce__() method of
2968 some extension type contains another object of the same type. Ouch!
2969
2970 Here is a quick example, that I ran into, to illustrate what I
2971 mean:
2972
2973 >>> import pickle, copyreg
2974 >>> copyreg.dispatch_table.pop(complex)
2975 >>> pickle.dumps(1+2j)
2976 Traceback (most recent call last):
2977 ...
2978 RuntimeError: maximum recursion depth exceeded
2979
2980 Removing the complex class from copyreg.dispatch_table made the
2981 __reduce_ex__() method emit another complex object:
2982
2983 >>> (1+1j).__reduce_ex__(2)
2984 (<function __newobj__ at 0xb7b71c3c>,
2985 (<class 'complex'>, (1+1j)), None, None, None)
2986
2987 Thus when save() was called on newargstup (the 2nd item) recursion
2988 ensued. Of course, the bug was in the complex class which had a
2989 broken __getnewargs__() that emitted another complex object. But,
2990 the point, here, is it is quite easy to end up with a broken reduce
2991 function. */
2992
2993 /* Save the class and its __new__ arguments. */
2994 if (save(self, cls, 0) < 0)
2995 return -1;
2996
2997 newargtup = PyTuple_GetSlice(argtup, 1, Py_SIZE(argtup));
2998 if (newargtup == NULL)
2999 return -1;
3000
3001 p = save(self, newargtup, 0);
3002 Py_DECREF(newargtup);
3003 if (p < 0)
3004 return -1;
3005
3006 /* Add NEWOBJ opcode. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003007 if (_Pickler_Write(self, &newobj_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003008 return -1;
3009 }
3010 else { /* Not using NEWOBJ. */
3011 if (save(self, callable, 0) < 0 ||
3012 save(self, argtup, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003013 _Pickler_Write(self, &reduce_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003014 return -1;
3015 }
3016
3017 /* obj can be NULL when save_reduce() is used directly. A NULL obj means
3018 the caller do not want to memoize the object. Not particularly useful,
3019 but that is to mimic the behavior save_reduce() in pickle.py when
3020 obj is None. */
3021 if (obj && memo_put(self, obj) < 0)
3022 return -1;
3023
3024 if (listitems && batch_list(self, listitems) < 0)
3025 return -1;
3026
3027 if (dictitems && batch_dict(self, dictitems) < 0)
3028 return -1;
3029
3030 if (state) {
3031 if (save(self, state, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003032 _Pickler_Write(self, &build_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003033 return -1;
3034 }
3035
3036 return 0;
3037}
3038
3039static int
3040save(PicklerObject *self, PyObject *obj, int pers_save)
3041{
3042 PyTypeObject *type;
3043 PyObject *reduce_func = NULL;
3044 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003045 int status = 0;
3046
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003047 if (Py_EnterRecursiveCall(" while pickling an object") < 0)
3048 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003049
3050 /* The extra pers_save argument is necessary to avoid calling save_pers()
3051 on its returned object. */
3052 if (!pers_save && self->pers_func) {
3053 /* save_pers() returns:
3054 -1 to signal an error;
3055 0 if it did nothing successfully;
3056 1 if a persistent id was saved.
3057 */
3058 if ((status = save_pers(self, obj, self->pers_func)) != 0)
3059 goto done;
3060 }
3061
3062 type = Py_TYPE(obj);
3063
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003064 /* The old cPickle had an optimization that used switch-case statement
3065 dispatching on the first letter of the type name. This has was removed
3066 since benchmarks shown that this optimization was actually slowing
3067 things down. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003068
3069 /* Atom types; these aren't memoized, so don't check the memo. */
3070
3071 if (obj == Py_None) {
3072 status = save_none(self, obj);
3073 goto done;
3074 }
3075 else if (obj == Py_False || obj == Py_True) {
3076 status = save_bool(self, obj);
3077 goto done;
3078 }
3079 else if (type == &PyLong_Type) {
3080 status = save_long(self, obj);
3081 goto done;
3082 }
3083 else if (type == &PyFloat_Type) {
3084 status = save_float(self, obj);
3085 goto done;
3086 }
3087
3088 /* Check the memo to see if it has the object. If so, generate
3089 a GET (or BINGET) opcode, instead of pickling the object
3090 once again. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003091 if (PyMemoTable_Get(self->memo, obj)) {
3092 if (memo_get(self, obj) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003093 goto error;
3094 goto done;
3095 }
3096
3097 if (type == &PyBytes_Type) {
3098 status = save_bytes(self, obj);
3099 goto done;
3100 }
3101 else if (type == &PyUnicode_Type) {
3102 status = save_unicode(self, obj);
3103 goto done;
3104 }
3105 else if (type == &PyDict_Type) {
3106 status = save_dict(self, obj);
3107 goto done;
3108 }
3109 else if (type == &PyList_Type) {
3110 status = save_list(self, obj);
3111 goto done;
3112 }
3113 else if (type == &PyTuple_Type) {
3114 status = save_tuple(self, obj);
3115 goto done;
3116 }
3117 else if (type == &PyType_Type) {
3118 status = save_global(self, obj, NULL);
3119 goto done;
3120 }
3121 else if (type == &PyFunction_Type) {
3122 status = save_global(self, obj, NULL);
3123 if (status < 0 && PyErr_ExceptionMatches(PickleError)) {
3124 /* fall back to reduce */
3125 PyErr_Clear();
3126 }
3127 else {
3128 goto done;
3129 }
3130 }
3131 else if (type == &PyCFunction_Type) {
3132 status = save_global(self, obj, NULL);
3133 goto done;
3134 }
3135 else if (PyType_IsSubtype(type, &PyType_Type)) {
3136 status = save_global(self, obj, NULL);
3137 goto done;
3138 }
3139
3140 /* XXX: This part needs some unit tests. */
3141
3142 /* Get a reduction callable, and call it. This may come from
3143 * copyreg.dispatch_table, the object's __reduce_ex__ method,
3144 * or the object's __reduce__ method.
3145 */
3146 reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
3147 if (reduce_func != NULL) {
3148 /* Here, the reference count of the reduce_func object returned by
3149 PyDict_GetItem needs to be increased to be consistent with the one
3150 returned by PyObject_GetAttr. This is allow us to blindly DECREF
3151 reduce_func at the end of the save() routine.
3152 */
3153 Py_INCREF(reduce_func);
3154 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003155 reduce_value = _Pickler_FastCall(self, reduce_func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003156 }
3157 else {
3158 static PyObject *reduce_str = NULL;
3159 static PyObject *reduce_ex_str = NULL;
3160
3161 /* Cache the name of the reduce methods. */
3162 if (reduce_str == NULL) {
3163 reduce_str = PyUnicode_InternFromString("__reduce__");
3164 if (reduce_str == NULL)
3165 goto error;
3166 reduce_ex_str = PyUnicode_InternFromString("__reduce_ex__");
3167 if (reduce_ex_str == NULL)
3168 goto error;
3169 }
3170
3171 /* XXX: If the __reduce__ method is defined, __reduce_ex__ is
3172 automatically defined as __reduce__. While this is convenient, this
3173 make it impossible to know which method was actually called. Of
3174 course, this is not a big deal. But still, it would be nice to let
3175 the user know which method was called when something go
3176 wrong. Incidentally, this means if __reduce_ex__ is not defined, we
3177 don't actually have to check for a __reduce__ method. */
3178
3179 /* Check for a __reduce_ex__ method. */
3180 reduce_func = PyObject_GetAttr(obj, reduce_ex_str);
3181 if (reduce_func != NULL) {
3182 PyObject *proto;
3183 proto = PyLong_FromLong(self->proto);
3184 if (proto != NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003185 reduce_value = _Pickler_FastCall(self, reduce_func, proto);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003186 }
3187 }
3188 else {
3189 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3190 PyErr_Clear();
3191 else
3192 goto error;
3193 /* Check for a __reduce__ method. */
3194 reduce_func = PyObject_GetAttr(obj, reduce_str);
3195 if (reduce_func != NULL) {
3196 reduce_value = PyObject_Call(reduce_func, empty_tuple, NULL);
3197 }
3198 else {
3199 PyErr_Format(PicklingError, "can't pickle '%.200s' object: %R",
3200 type->tp_name, obj);
3201 goto error;
3202 }
3203 }
3204 }
3205
3206 if (reduce_value == NULL)
3207 goto error;
3208
3209 if (PyUnicode_Check(reduce_value)) {
3210 status = save_global(self, obj, reduce_value);
3211 goto done;
3212 }
3213
3214 if (!PyTuple_Check(reduce_value)) {
3215 PyErr_SetString(PicklingError,
3216 "__reduce__ must return a string or tuple");
3217 goto error;
3218 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003219
3220 status = save_reduce(self, reduce_value, obj);
3221
3222 if (0) {
3223 error:
3224 status = -1;
3225 }
3226 done:
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003227 Py_LeaveRecursiveCall();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003228 Py_XDECREF(reduce_func);
3229 Py_XDECREF(reduce_value);
3230
3231 return status;
3232}
3233
3234static int
3235dump(PicklerObject *self, PyObject *obj)
3236{
3237 const char stop_op = STOP;
3238
3239 if (self->proto >= 2) {
3240 char header[2];
3241
3242 header[0] = PROTO;
3243 assert(self->proto >= 0 && self->proto < 256);
3244 header[1] = (unsigned char)self->proto;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003245 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003246 return -1;
3247 }
3248
3249 if (save(self, obj, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003250 _Pickler_Write(self, &stop_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003251 return -1;
3252
3253 return 0;
3254}
3255
3256PyDoc_STRVAR(Pickler_clear_memo_doc,
3257"clear_memo() -> None. Clears the pickler's \"memo\"."
3258"\n"
3259"The memo is the data structure that remembers which objects the\n"
3260"pickler has already seen, so that shared or recursive objects are\n"
3261"pickled by reference and not by value. This method is useful when\n"
3262"re-using picklers.");
3263
3264static PyObject *
3265Pickler_clear_memo(PicklerObject *self)
3266{
3267 if (self->memo)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003268 PyMemoTable_Clear(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003269
3270 Py_RETURN_NONE;
3271}
3272
3273PyDoc_STRVAR(Pickler_dump_doc,
3274"dump(obj) -> None. Write a pickled representation of obj to the open file.");
3275
3276static PyObject *
3277Pickler_dump(PicklerObject *self, PyObject *args)
3278{
3279 PyObject *obj;
3280
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003281 /* Check whether the Pickler was initialized correctly (issue3664).
3282 Developers often forget to call __init__() in their subclasses, which
3283 would trigger a segfault without this check. */
3284 if (self->write == NULL) {
3285 PyErr_Format(PicklingError,
3286 "Pickler.__init__() was not called by %s.__init__()",
3287 Py_TYPE(self)->tp_name);
3288 return NULL;
3289 }
3290
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003291 if (!PyArg_ParseTuple(args, "O:dump", &obj))
3292 return NULL;
3293
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003294 if (_Pickler_ClearBuffer(self) < 0)
3295 return NULL;
3296
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003297 if (dump(self, obj) < 0)
3298 return NULL;
3299
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003300 if (_Pickler_FlushToFile(self) < 0)
3301 return NULL;
3302
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003303 Py_RETURN_NONE;
3304}
3305
3306static struct PyMethodDef Pickler_methods[] = {
3307 {"dump", (PyCFunction)Pickler_dump, METH_VARARGS,
3308 Pickler_dump_doc},
3309 {"clear_memo", (PyCFunction)Pickler_clear_memo, METH_NOARGS,
3310 Pickler_clear_memo_doc},
3311 {NULL, NULL} /* sentinel */
3312};
3313
3314static void
3315Pickler_dealloc(PicklerObject *self)
3316{
3317 PyObject_GC_UnTrack(self);
3318
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003319 Py_XDECREF(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003320 Py_XDECREF(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003321 Py_XDECREF(self->pers_func);
3322 Py_XDECREF(self->arg);
3323 Py_XDECREF(self->fast_memo);
3324
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003325 PyMemoTable_Del(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003326
3327 Py_TYPE(self)->tp_free((PyObject *)self);
3328}
3329
3330static int
3331Pickler_traverse(PicklerObject *self, visitproc visit, void *arg)
3332{
3333 Py_VISIT(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003334 Py_VISIT(self->pers_func);
3335 Py_VISIT(self->arg);
3336 Py_VISIT(self->fast_memo);
3337 return 0;
3338}
3339
3340static int
3341Pickler_clear(PicklerObject *self)
3342{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003343 Py_CLEAR(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003344 Py_CLEAR(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003345 Py_CLEAR(self->pers_func);
3346 Py_CLEAR(self->arg);
3347 Py_CLEAR(self->fast_memo);
3348
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003349 if (self->memo != NULL) {
3350 PyMemoTable *memo = self->memo;
3351 self->memo = NULL;
3352 PyMemoTable_Del(memo);
3353 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003354 return 0;
3355}
3356
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003357
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003358PyDoc_STRVAR(Pickler_doc,
3359"Pickler(file, protocol=None)"
3360"\n"
3361"This takes a binary file for writing a pickle data stream.\n"
3362"\n"
3363"The optional protocol argument tells the pickler to use the\n"
3364"given protocol; supported protocols are 0, 1, 2, 3. The default\n"
3365"protocol is 3; a backward-incompatible protocol designed for\n"
3366"Python 3.0.\n"
3367"\n"
3368"Specifying a negative protocol version selects the highest\n"
3369"protocol version supported. The higher the protocol used, the\n"
3370"more recent the version of Python needed to read the pickle\n"
3371"produced.\n"
3372"\n"
3373"The file argument must have a write() method that accepts a single\n"
3374"bytes argument. It can thus be a file object opened for binary\n"
3375"writing, a io.BytesIO instance, or any other custom object that\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003376"meets this interface.\n"
3377"\n"
3378"If fix_imports is True and protocol is less than 3, pickle will try to\n"
3379"map the new Python 3.x names to the old module names used in Python\n"
3380"2.x, so that the pickle data stream is readable with Python 2.x.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003381
3382static int
3383Pickler_init(PicklerObject *self, PyObject *args, PyObject *kwds)
3384{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003385 static char *kwlist[] = {"file", "protocol", "fix_imports", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003386 PyObject *file;
3387 PyObject *proto_obj = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003388 PyObject *fix_imports = Py_True;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003389
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003390 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:Pickler",
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003391 kwlist, &file, &proto_obj, &fix_imports))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003392 return -1;
3393
3394 /* In case of multiple __init__() calls, clear previous content. */
3395 if (self->write != NULL)
3396 (void)Pickler_clear(self);
3397
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003398 if (_Pickler_SetProtocol(self, proto_obj, fix_imports) < 0)
3399 return -1;
3400
3401 if (_Pickler_SetOutputStream(self, file) < 0)
3402 return -1;
3403
3404 /* memo and output_buffer may have already been created in _Pickler_New */
3405 if (self->memo == NULL) {
3406 self->memo = PyMemoTable_New();
3407 if (self->memo == NULL)
3408 return -1;
3409 }
3410 self->output_len = 0;
3411 if (self->output_buffer == NULL) {
3412 self->max_output_len = WRITE_BUF_SIZE;
3413 self->output_buffer = PyBytes_FromStringAndSize(NULL,
3414 self->max_output_len);
3415 if (self->output_buffer == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003416 return -1;
3417 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003418
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003419 self->arg = NULL;
3420 self->fast = 0;
3421 self->fast_nesting = 0;
3422 self->fast_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003423 self->pers_func = NULL;
3424 if (PyObject_HasAttrString((PyObject *)self, "persistent_id")) {
3425 self->pers_func = PyObject_GetAttrString((PyObject *)self,
3426 "persistent_id");
3427 if (self->pers_func == NULL)
3428 return -1;
3429 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003430 return 0;
3431}
3432
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003433/* Define a proxy object for the Pickler's internal memo object. This is to
3434 * avoid breaking code like:
3435 * pickler.memo.clear()
3436 * and
3437 * pickler.memo = saved_memo
3438 * Is this a good idea? Not really, but we don't want to break code that uses
3439 * it. Note that we don't implement the entire mapping API here. This is
3440 * intentional, as these should be treated as black-box implementation details.
3441 */
3442
3443typedef struct {
3444 PyObject_HEAD
3445 PicklerObject *pickler; /* Pickler whose memo table we're proxying. */
3446} PicklerMemoProxyObject;
3447
3448PyDoc_STRVAR(pmp_clear_doc,
3449"memo.clear() -> None. Remove all items from memo.");
3450
3451static PyObject *
3452pmp_clear(PicklerMemoProxyObject *self)
3453{
3454 if (self->pickler->memo)
3455 PyMemoTable_Clear(self->pickler->memo);
3456 Py_RETURN_NONE;
3457}
3458
3459PyDoc_STRVAR(pmp_copy_doc,
3460"memo.copy() -> new_memo. Copy the memo to a new object.");
3461
3462static PyObject *
3463pmp_copy(PicklerMemoProxyObject *self)
3464{
3465 Py_ssize_t i;
3466 PyMemoTable *memo;
3467 PyObject *new_memo = PyDict_New();
3468 if (new_memo == NULL)
3469 return NULL;
3470
3471 memo = self->pickler->memo;
3472 for (i = 0; i < memo->mt_allocated; ++i) {
3473 PyMemoEntry entry = memo->mt_table[i];
3474 if (entry.me_key != NULL) {
3475 int status;
3476 PyObject *key, *value;
3477
3478 key = PyLong_FromVoidPtr(entry.me_key);
3479 value = Py_BuildValue("lO", entry.me_value, entry.me_key);
3480
3481 if (key == NULL || value == NULL) {
3482 Py_XDECREF(key);
3483 Py_XDECREF(value);
3484 goto error;
3485 }
3486 status = PyDict_SetItem(new_memo, key, value);
3487 Py_DECREF(key);
3488 Py_DECREF(value);
3489 if (status < 0)
3490 goto error;
3491 }
3492 }
3493 return new_memo;
3494
3495 error:
3496 Py_XDECREF(new_memo);
3497 return NULL;
3498}
3499
3500PyDoc_STRVAR(pmp_reduce_doc,
3501"memo.__reduce__(). Pickling support.");
3502
3503static PyObject *
3504pmp_reduce(PicklerMemoProxyObject *self, PyObject *args)
3505{
3506 PyObject *reduce_value, *dict_args;
3507 PyObject *contents = pmp_copy(self);
3508 if (contents == NULL)
3509 return NULL;
3510
3511 reduce_value = PyTuple_New(2);
3512 if (reduce_value == NULL) {
3513 Py_DECREF(contents);
3514 return NULL;
3515 }
3516 dict_args = PyTuple_New(1);
3517 if (dict_args == NULL) {
3518 Py_DECREF(contents);
3519 Py_DECREF(reduce_value);
3520 return NULL;
3521 }
3522 PyTuple_SET_ITEM(dict_args, 0, contents);
3523 Py_INCREF((PyObject *)&PyDict_Type);
3524 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
3525 PyTuple_SET_ITEM(reduce_value, 1, dict_args);
3526 return reduce_value;
3527}
3528
3529static PyMethodDef picklerproxy_methods[] = {
3530 {"clear", (PyCFunction)pmp_clear, METH_NOARGS, pmp_clear_doc},
3531 {"copy", (PyCFunction)pmp_copy, METH_NOARGS, pmp_copy_doc},
3532 {"__reduce__", (PyCFunction)pmp_reduce, METH_VARARGS, pmp_reduce_doc},
3533 {NULL, NULL} /* sentinel */
3534};
3535
3536static void
3537PicklerMemoProxy_dealloc(PicklerMemoProxyObject *self)
3538{
3539 PyObject_GC_UnTrack(self);
3540 Py_XDECREF(self->pickler);
3541 PyObject_GC_Del((PyObject *)self);
3542}
3543
3544static int
3545PicklerMemoProxy_traverse(PicklerMemoProxyObject *self,
3546 visitproc visit, void *arg)
3547{
3548 Py_VISIT(self->pickler);
3549 return 0;
3550}
3551
3552static int
3553PicklerMemoProxy_clear(PicklerMemoProxyObject *self)
3554{
3555 Py_CLEAR(self->pickler);
3556 return 0;
3557}
3558
3559static PyTypeObject PicklerMemoProxyType = {
3560 PyVarObject_HEAD_INIT(NULL, 0)
3561 "_pickle.PicklerMemoProxy", /*tp_name*/
3562 sizeof(PicklerMemoProxyObject), /*tp_basicsize*/
3563 0,
3564 (destructor)PicklerMemoProxy_dealloc, /* tp_dealloc */
3565 0, /* tp_print */
3566 0, /* tp_getattr */
3567 0, /* tp_setattr */
3568 0, /* tp_compare */
3569 0, /* tp_repr */
3570 0, /* tp_as_number */
3571 0, /* tp_as_sequence */
3572 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00003573 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003574 0, /* tp_call */
3575 0, /* tp_str */
3576 PyObject_GenericGetAttr, /* tp_getattro */
3577 PyObject_GenericSetAttr, /* tp_setattro */
3578 0, /* tp_as_buffer */
3579 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3580 0, /* tp_doc */
3581 (traverseproc)PicklerMemoProxy_traverse, /* tp_traverse */
3582 (inquiry)PicklerMemoProxy_clear, /* tp_clear */
3583 0, /* tp_richcompare */
3584 0, /* tp_weaklistoffset */
3585 0, /* tp_iter */
3586 0, /* tp_iternext */
3587 picklerproxy_methods, /* tp_methods */
3588};
3589
3590static PyObject *
3591PicklerMemoProxy_New(PicklerObject *pickler)
3592{
3593 PicklerMemoProxyObject *self;
3594
3595 self = PyObject_GC_New(PicklerMemoProxyObject, &PicklerMemoProxyType);
3596 if (self == NULL)
3597 return NULL;
3598 Py_INCREF(pickler);
3599 self->pickler = pickler;
3600 PyObject_GC_Track(self);
3601 return (PyObject *)self;
3602}
3603
3604/*****************************************************************************/
3605
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003606static PyObject *
3607Pickler_get_memo(PicklerObject *self)
3608{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003609 return PicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003610}
3611
3612static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003613Pickler_set_memo(PicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003614{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003615 PyMemoTable *new_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003616
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003617 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003618 PyErr_SetString(PyExc_TypeError,
3619 "attribute deletion is not supported");
3620 return -1;
3621 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003622
3623 if (Py_TYPE(obj) == &PicklerMemoProxyType) {
3624 PicklerObject *pickler =
3625 ((PicklerMemoProxyObject *)obj)->pickler;
3626
3627 new_memo = PyMemoTable_Copy(pickler->memo);
3628 if (new_memo == NULL)
3629 return -1;
3630 }
3631 else if (PyDict_Check(obj)) {
3632 Py_ssize_t i = 0;
3633 PyObject *key, *value;
3634
3635 new_memo = PyMemoTable_New();
3636 if (new_memo == NULL)
3637 return -1;
3638
3639 while (PyDict_Next(obj, &i, &key, &value)) {
3640 long memo_id;
3641 PyObject *memo_obj;
3642
3643 if (!PyTuple_Check(value) || Py_SIZE(value) != 2) {
3644 PyErr_SetString(PyExc_TypeError,
3645 "'memo' values must be 2-item tuples");
3646 goto error;
3647 }
3648 memo_id = PyLong_AsLong(PyTuple_GET_ITEM(value, 0));
3649 if (memo_id == -1 && PyErr_Occurred())
3650 goto error;
3651 memo_obj = PyTuple_GET_ITEM(value, 1);
3652 if (PyMemoTable_Set(new_memo, memo_obj, memo_id) < 0)
3653 goto error;
3654 }
3655 }
3656 else {
3657 PyErr_Format(PyExc_TypeError,
3658 "'memo' attribute must be an PicklerMemoProxy object"
3659 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003660 return -1;
3661 }
3662
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003663 PyMemoTable_Del(self->memo);
3664 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003665
3666 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003667
3668 error:
3669 if (new_memo)
3670 PyMemoTable_Del(new_memo);
3671 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003672}
3673
3674static PyObject *
3675Pickler_get_persid(PicklerObject *self)
3676{
3677 if (self->pers_func == NULL)
3678 PyErr_SetString(PyExc_AttributeError, "persistent_id");
3679 else
3680 Py_INCREF(self->pers_func);
3681 return self->pers_func;
3682}
3683
3684static int
3685Pickler_set_persid(PicklerObject *self, PyObject *value)
3686{
3687 PyObject *tmp;
3688
3689 if (value == NULL) {
3690 PyErr_SetString(PyExc_TypeError,
3691 "attribute deletion is not supported");
3692 return -1;
3693 }
3694 if (!PyCallable_Check(value)) {
3695 PyErr_SetString(PyExc_TypeError,
3696 "persistent_id must be a callable taking one argument");
3697 return -1;
3698 }
3699
3700 tmp = self->pers_func;
3701 Py_INCREF(value);
3702 self->pers_func = value;
3703 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
3704
3705 return 0;
3706}
3707
3708static PyMemberDef Pickler_members[] = {
3709 {"bin", T_INT, offsetof(PicklerObject, bin)},
3710 {"fast", T_INT, offsetof(PicklerObject, fast)},
3711 {NULL}
3712};
3713
3714static PyGetSetDef Pickler_getsets[] = {
3715 {"memo", (getter)Pickler_get_memo,
3716 (setter)Pickler_set_memo},
3717 {"persistent_id", (getter)Pickler_get_persid,
3718 (setter)Pickler_set_persid},
3719 {NULL}
3720};
3721
3722static PyTypeObject Pickler_Type = {
3723 PyVarObject_HEAD_INIT(NULL, 0)
3724 "_pickle.Pickler" , /*tp_name*/
3725 sizeof(PicklerObject), /*tp_basicsize*/
3726 0, /*tp_itemsize*/
3727 (destructor)Pickler_dealloc, /*tp_dealloc*/
3728 0, /*tp_print*/
3729 0, /*tp_getattr*/
3730 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00003731 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003732 0, /*tp_repr*/
3733 0, /*tp_as_number*/
3734 0, /*tp_as_sequence*/
3735 0, /*tp_as_mapping*/
3736 0, /*tp_hash*/
3737 0, /*tp_call*/
3738 0, /*tp_str*/
3739 0, /*tp_getattro*/
3740 0, /*tp_setattro*/
3741 0, /*tp_as_buffer*/
3742 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3743 Pickler_doc, /*tp_doc*/
3744 (traverseproc)Pickler_traverse, /*tp_traverse*/
3745 (inquiry)Pickler_clear, /*tp_clear*/
3746 0, /*tp_richcompare*/
3747 0, /*tp_weaklistoffset*/
3748 0, /*tp_iter*/
3749 0, /*tp_iternext*/
3750 Pickler_methods, /*tp_methods*/
3751 Pickler_members, /*tp_members*/
3752 Pickler_getsets, /*tp_getset*/
3753 0, /*tp_base*/
3754 0, /*tp_dict*/
3755 0, /*tp_descr_get*/
3756 0, /*tp_descr_set*/
3757 0, /*tp_dictoffset*/
3758 (initproc)Pickler_init, /*tp_init*/
3759 PyType_GenericAlloc, /*tp_alloc*/
3760 PyType_GenericNew, /*tp_new*/
3761 PyObject_GC_Del, /*tp_free*/
3762 0, /*tp_is_gc*/
3763};
3764
3765/* Temporary helper for calling self.find_class().
3766
3767 XXX: It would be nice to able to avoid Python function call overhead, by
3768 using directly the C version of find_class(), when find_class() is not
3769 overridden by a subclass. Although, this could become rather hackish. A
3770 simpler optimization would be to call the C function when self is not a
3771 subclass instance. */
3772static PyObject *
3773find_class(UnpicklerObject *self, PyObject *module_name, PyObject *global_name)
3774{
3775 return PyObject_CallMethod((PyObject *)self, "find_class", "OO",
3776 module_name, global_name);
3777}
3778
3779static int
3780marker(UnpicklerObject *self)
3781{
3782 if (self->num_marks < 1) {
3783 PyErr_SetString(UnpicklingError, "could not find MARK");
3784 return -1;
3785 }
3786
3787 return self->marks[--self->num_marks];
3788}
3789
3790static int
3791load_none(UnpicklerObject *self)
3792{
3793 PDATA_APPEND(self->stack, Py_None, -1);
3794 return 0;
3795}
3796
3797static int
3798bad_readline(void)
3799{
3800 PyErr_SetString(UnpicklingError, "pickle data was truncated");
3801 return -1;
3802}
3803
3804static int
3805load_int(UnpicklerObject *self)
3806{
3807 PyObject *value;
3808 char *endptr, *s;
3809 Py_ssize_t len;
3810 long x;
3811
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003812 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003813 return -1;
3814 if (len < 2)
3815 return bad_readline();
3816
3817 errno = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003818 /* XXX: Should the base argument of strtol() be explicitly set to 10?
3819 XXX(avassalotti): Should this uses PyOS_strtol()? */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003820 x = strtol(s, &endptr, 0);
3821
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003822 if (errno || (*endptr != '\n' && *endptr != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003823 /* Hm, maybe we've got something long. Let's try reading
3824 * it as a Python long object. */
3825 errno = 0;
3826 /* XXX: Same thing about the base here. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003827 value = PyLong_FromString(s, NULL, 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003828 if (value == NULL) {
3829 PyErr_SetString(PyExc_ValueError,
3830 "could not convert string to int");
3831 return -1;
3832 }
3833 }
3834 else {
3835 if (len == 3 && (x == 0 || x == 1)) {
3836 if ((value = PyBool_FromLong(x)) == NULL)
3837 return -1;
3838 }
3839 else {
3840 if ((value = PyLong_FromLong(x)) == NULL)
3841 return -1;
3842 }
3843 }
3844
3845 PDATA_PUSH(self->stack, value, -1);
3846 return 0;
3847}
3848
3849static int
3850load_bool(UnpicklerObject *self, PyObject *boolean)
3851{
3852 assert(boolean == Py_True || boolean == Py_False);
3853 PDATA_APPEND(self->stack, boolean, -1);
3854 return 0;
3855}
3856
3857/* s contains x bytes of a little-endian integer. Return its value as a
3858 * C int. Obscure: when x is 1 or 2, this is an unsigned little-endian
3859 * int, but when x is 4 it's a signed one. This is an historical source
3860 * of x-platform bugs.
3861 */
3862static long
3863calc_binint(char *bytes, int size)
3864{
3865 unsigned char *s = (unsigned char *)bytes;
3866 int i = size;
3867 long x = 0;
3868
3869 for (i = 0; i < size; i++) {
3870 x |= (long)s[i] << (i * 8);
3871 }
3872
3873 /* Unlike BININT1 and BININT2, BININT (more accurately BININT4)
3874 * is signed, so on a box with longs bigger than 4 bytes we need
3875 * to extend a BININT's sign bit to the full width.
3876 */
3877 if (SIZEOF_LONG > 4 && size == 4) {
3878 x |= -(x & (1L << 31));
3879 }
3880
3881 return x;
3882}
3883
3884static int
3885load_binintx(UnpicklerObject *self, char *s, int size)
3886{
3887 PyObject *value;
3888 long x;
3889
3890 x = calc_binint(s, size);
3891
3892 if ((value = PyLong_FromLong(x)) == NULL)
3893 return -1;
3894
3895 PDATA_PUSH(self->stack, value, -1);
3896 return 0;
3897}
3898
3899static int
3900load_binint(UnpicklerObject *self)
3901{
3902 char *s;
3903
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003904 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003905 return -1;
3906
3907 return load_binintx(self, s, 4);
3908}
3909
3910static int
3911load_binint1(UnpicklerObject *self)
3912{
3913 char *s;
3914
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003915 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003916 return -1;
3917
3918 return load_binintx(self, s, 1);
3919}
3920
3921static int
3922load_binint2(UnpicklerObject *self)
3923{
3924 char *s;
3925
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003926 if (_Unpickler_Read(self, &s, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003927 return -1;
3928
3929 return load_binintx(self, s, 2);
3930}
3931
3932static int
3933load_long(UnpicklerObject *self)
3934{
3935 PyObject *value;
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00003936 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003937 Py_ssize_t len;
3938
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003939 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003940 return -1;
3941 if (len < 2)
3942 return bad_readline();
3943
Mark Dickinson8dd05142009-01-20 20:43:58 +00003944 /* s[len-2] will usually be 'L' (and s[len-1] is '\n'); we need to remove
3945 the 'L' before calling PyLong_FromString. In order to maintain
3946 compatibility with Python 3.0.0, we don't actually *require*
3947 the 'L' to be present. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003948 if (s[len-2] == 'L')
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00003949 s[len-2] = '\0';
Alexandre Vassalottie4bccb72009-01-24 01:47:57 +00003950 /* XXX: Should the base argument explicitly set to 10? */
3951 value = PyLong_FromString(s, NULL, 0);
Mark Dickinson8dd05142009-01-20 20:43:58 +00003952 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003953 return -1;
3954
3955 PDATA_PUSH(self->stack, value, -1);
3956 return 0;
3957}
3958
3959/* 'size' bytes contain the # of bytes of little-endian 256's-complement
3960 * data following.
3961 */
3962static int
3963load_counted_long(UnpicklerObject *self, int size)
3964{
3965 PyObject *value;
3966 char *nbytes;
3967 char *pdata;
3968
3969 assert(size == 1 || size == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003970 if (_Unpickler_Read(self, &nbytes, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003971 return -1;
3972
3973 size = calc_binint(nbytes, size);
3974 if (size < 0) {
3975 /* Corrupt or hostile pickle -- we never write one like this */
3976 PyErr_SetString(UnpicklingError,
3977 "LONG pickle has negative byte count");
3978 return -1;
3979 }
3980
3981 if (size == 0)
3982 value = PyLong_FromLong(0L);
3983 else {
3984 /* Read the raw little-endian bytes and convert. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003985 if (_Unpickler_Read(self, &pdata, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003986 return -1;
3987 value = _PyLong_FromByteArray((unsigned char *)pdata, (size_t)size,
3988 1 /* little endian */ , 1 /* signed */ );
3989 }
3990 if (value == NULL)
3991 return -1;
3992 PDATA_PUSH(self->stack, value, -1);
3993 return 0;
3994}
3995
3996static int
3997load_float(UnpicklerObject *self)
3998{
3999 PyObject *value;
4000 char *endptr, *s;
4001 Py_ssize_t len;
4002 double d;
4003
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004004 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004005 return -1;
4006 if (len < 2)
4007 return bad_readline();
4008
4009 errno = 0;
Mark Dickinson725bfd82009-05-03 20:33:40 +00004010 d = PyOS_string_to_double(s, &endptr, PyExc_OverflowError);
4011 if (d == -1.0 && PyErr_Occurred())
4012 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004013 if ((endptr[0] != '\n') && (endptr[0] != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004014 PyErr_SetString(PyExc_ValueError, "could not convert string to float");
4015 return -1;
4016 }
Mark Dickinson725bfd82009-05-03 20:33:40 +00004017 value = PyFloat_FromDouble(d);
4018 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004019 return -1;
4020
4021 PDATA_PUSH(self->stack, value, -1);
4022 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004023}
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004024
4025static int
4026load_binfloat(UnpicklerObject *self)
4027{
4028 PyObject *value;
4029 double x;
4030 char *s;
4031
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004032 if (_Unpickler_Read(self, &s, 8) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004033 return -1;
4034
4035 x = _PyFloat_Unpack8((unsigned char *)s, 0);
4036 if (x == -1.0 && PyErr_Occurred())
4037 return -1;
4038
4039 if ((value = PyFloat_FromDouble(x)) == NULL)
4040 return -1;
4041
4042 PDATA_PUSH(self->stack, value, -1);
4043 return 0;
4044}
4045
4046static int
4047load_string(UnpicklerObject *self)
4048{
4049 PyObject *bytes;
4050 PyObject *str = NULL;
4051 Py_ssize_t len;
4052 char *s, *p;
4053
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004054 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004055 return -1;
4056 if (len < 3)
4057 return bad_readline();
4058 if ((s = strdup(s)) == NULL) {
4059 PyErr_NoMemory();
4060 return -1;
4061 }
4062
4063 /* Strip outermost quotes */
4064 while (s[len - 1] <= ' ')
4065 len--;
4066 if (s[0] == '"' && s[len - 1] == '"') {
4067 s[len - 1] = '\0';
4068 p = s + 1;
4069 len -= 2;
4070 }
4071 else if (s[0] == '\'' && s[len - 1] == '\'') {
4072 s[len - 1] = '\0';
4073 p = s + 1;
4074 len -= 2;
4075 }
4076 else {
4077 free(s);
4078 PyErr_SetString(PyExc_ValueError, "insecure string pickle");
4079 return -1;
4080 }
4081
4082 /* Use the PyBytes API to decode the string, since that is what is used
4083 to encode, and then coerce the result to Unicode. */
4084 bytes = PyBytes_DecodeEscape(p, len, NULL, 0, NULL);
4085 free(s);
4086 if (bytes == NULL)
4087 return -1;
4088 str = PyUnicode_FromEncodedObject(bytes, self->encoding, self->errors);
4089 Py_DECREF(bytes);
4090 if (str == NULL)
4091 return -1;
4092
4093 PDATA_PUSH(self->stack, str, -1);
4094 return 0;
4095}
4096
4097static int
4098load_binbytes(UnpicklerObject *self)
4099{
4100 PyObject *bytes;
4101 long x;
4102 char *s;
4103
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004104 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004105 return -1;
4106
4107 x = calc_binint(s, 4);
4108 if (x < 0) {
4109 PyErr_SetString(UnpicklingError,
4110 "BINBYTES pickle has negative byte count");
4111 return -1;
4112 }
4113
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004114 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004115 return -1;
4116 bytes = PyBytes_FromStringAndSize(s, x);
4117 if (bytes == NULL)
4118 return -1;
4119
4120 PDATA_PUSH(self->stack, bytes, -1);
4121 return 0;
4122}
4123
4124static int
4125load_short_binbytes(UnpicklerObject *self)
4126{
4127 PyObject *bytes;
4128 unsigned char x;
4129 char *s;
4130
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004131 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004132 return -1;
4133
4134 x = (unsigned char)s[0];
4135
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004136 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004137 return -1;
4138
4139 bytes = PyBytes_FromStringAndSize(s, x);
4140 if (bytes == NULL)
4141 return -1;
4142
4143 PDATA_PUSH(self->stack, bytes, -1);
4144 return 0;
4145}
4146
4147static int
4148load_binstring(UnpicklerObject *self)
4149{
4150 PyObject *str;
4151 long x;
4152 char *s;
4153
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004154 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004155 return -1;
4156
4157 x = calc_binint(s, 4);
4158 if (x < 0) {
4159 PyErr_SetString(UnpicklingError,
4160 "BINSTRING pickle has negative byte count");
4161 return -1;
4162 }
4163
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004164 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004165 return -1;
4166
4167 /* Convert Python 2.x strings to unicode. */
4168 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4169 if (str == NULL)
4170 return -1;
4171
4172 PDATA_PUSH(self->stack, str, -1);
4173 return 0;
4174}
4175
4176static int
4177load_short_binstring(UnpicklerObject *self)
4178{
4179 PyObject *str;
4180 unsigned char x;
4181 char *s;
4182
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004183 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004184 return -1;
4185
4186 x = (unsigned char)s[0];
4187
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004188 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004189 return -1;
4190
4191 /* Convert Python 2.x strings to unicode. */
4192 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4193 if (str == NULL)
4194 return -1;
4195
4196 PDATA_PUSH(self->stack, str, -1);
4197 return 0;
4198}
4199
4200static int
4201load_unicode(UnpicklerObject *self)
4202{
4203 PyObject *str;
4204 Py_ssize_t len;
4205 char *s;
4206
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004207 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004208 return -1;
4209 if (len < 1)
4210 return bad_readline();
4211
4212 str = PyUnicode_DecodeRawUnicodeEscape(s, len - 1, NULL);
4213 if (str == NULL)
4214 return -1;
4215
4216 PDATA_PUSH(self->stack, str, -1);
4217 return 0;
4218}
4219
4220static int
4221load_binunicode(UnpicklerObject *self)
4222{
4223 PyObject *str;
4224 long size;
4225 char *s;
4226
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004227 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004228 return -1;
4229
4230 size = calc_binint(s, 4);
4231 if (size < 0) {
4232 PyErr_SetString(UnpicklingError,
4233 "BINUNICODE pickle has negative byte count");
4234 return -1;
4235 }
4236
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004237 if (_Unpickler_Read(self, &s, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004238 return -1;
4239
Victor Stinner485fb562010-04-13 11:07:24 +00004240 str = PyUnicode_DecodeUTF8(s, size, "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004241 if (str == NULL)
4242 return -1;
4243
4244 PDATA_PUSH(self->stack, str, -1);
4245 return 0;
4246}
4247
4248static int
4249load_tuple(UnpicklerObject *self)
4250{
4251 PyObject *tuple;
4252 int i;
4253
4254 if ((i = marker(self)) < 0)
4255 return -1;
4256
4257 tuple = Pdata_poptuple(self->stack, i);
4258 if (tuple == NULL)
4259 return -1;
4260 PDATA_PUSH(self->stack, tuple, -1);
4261 return 0;
4262}
4263
4264static int
4265load_counted_tuple(UnpicklerObject *self, int len)
4266{
4267 PyObject *tuple;
4268
4269 tuple = PyTuple_New(len);
4270 if (tuple == NULL)
4271 return -1;
4272
4273 while (--len >= 0) {
4274 PyObject *item;
4275
4276 PDATA_POP(self->stack, item);
4277 if (item == NULL)
4278 return -1;
4279 PyTuple_SET_ITEM(tuple, len, item);
4280 }
4281 PDATA_PUSH(self->stack, tuple, -1);
4282 return 0;
4283}
4284
4285static int
4286load_empty_list(UnpicklerObject *self)
4287{
4288 PyObject *list;
4289
4290 if ((list = PyList_New(0)) == NULL)
4291 return -1;
4292 PDATA_PUSH(self->stack, list, -1);
4293 return 0;
4294}
4295
4296static int
4297load_empty_dict(UnpicklerObject *self)
4298{
4299 PyObject *dict;
4300
4301 if ((dict = PyDict_New()) == NULL)
4302 return -1;
4303 PDATA_PUSH(self->stack, dict, -1);
4304 return 0;
4305}
4306
4307static int
4308load_list(UnpicklerObject *self)
4309{
4310 PyObject *list;
4311 int i;
4312
4313 if ((i = marker(self)) < 0)
4314 return -1;
4315
4316 list = Pdata_poplist(self->stack, i);
4317 if (list == NULL)
4318 return -1;
4319 PDATA_PUSH(self->stack, list, -1);
4320 return 0;
4321}
4322
4323static int
4324load_dict(UnpicklerObject *self)
4325{
4326 PyObject *dict, *key, *value;
4327 int i, j, k;
4328
4329 if ((i = marker(self)) < 0)
4330 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004331 j = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004332
4333 if ((dict = PyDict_New()) == NULL)
4334 return -1;
4335
4336 for (k = i + 1; k < j; k += 2) {
4337 key = self->stack->data[k - 1];
4338 value = self->stack->data[k];
4339 if (PyDict_SetItem(dict, key, value) < 0) {
4340 Py_DECREF(dict);
4341 return -1;
4342 }
4343 }
4344 Pdata_clear(self->stack, i);
4345 PDATA_PUSH(self->stack, dict, -1);
4346 return 0;
4347}
4348
4349static PyObject *
4350instantiate(PyObject *cls, PyObject *args)
4351{
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004352 PyObject *result = NULL;
4353 /* Caller must assure args are a tuple. Normally, args come from
4354 Pdata_poptuple which packs objects from the top of the stack
4355 into a newly created tuple. */
4356 assert(PyTuple_Check(args));
4357 if (Py_SIZE(args) > 0 || !PyType_Check(cls) ||
4358 PyObject_HasAttrString(cls, "__getinitargs__")) {
4359 result = PyObject_CallObject(cls, args);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004360 }
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004361 else {
4362 result = PyObject_CallMethod(cls, "__new__", "O", cls);
4363 }
4364 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004365}
4366
4367static int
4368load_obj(UnpicklerObject *self)
4369{
4370 PyObject *cls, *args, *obj = NULL;
4371 int i;
4372
4373 if ((i = marker(self)) < 0)
4374 return -1;
4375
4376 args = Pdata_poptuple(self->stack, i + 1);
4377 if (args == NULL)
4378 return -1;
4379
4380 PDATA_POP(self->stack, cls);
4381 if (cls) {
4382 obj = instantiate(cls, args);
4383 Py_DECREF(cls);
4384 }
4385 Py_DECREF(args);
4386 if (obj == NULL)
4387 return -1;
4388
4389 PDATA_PUSH(self->stack, obj, -1);
4390 return 0;
4391}
4392
4393static int
4394load_inst(UnpicklerObject *self)
4395{
4396 PyObject *cls = NULL;
4397 PyObject *args = NULL;
4398 PyObject *obj = NULL;
4399 PyObject *module_name;
4400 PyObject *class_name;
4401 Py_ssize_t len;
4402 int i;
4403 char *s;
4404
4405 if ((i = marker(self)) < 0)
4406 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004407 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004408 return -1;
4409 if (len < 2)
4410 return bad_readline();
4411
4412 /* Here it is safe to use PyUnicode_DecodeASCII(), even though non-ASCII
4413 identifiers are permitted in Python 3.0, since the INST opcode is only
4414 supported by older protocols on Python 2.x. */
4415 module_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
4416 if (module_name == NULL)
4417 return -1;
4418
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004419 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004420 if (len < 2)
4421 return bad_readline();
4422 class_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004423 if (class_name != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004424 cls = find_class(self, module_name, class_name);
4425 Py_DECREF(class_name);
4426 }
4427 }
4428 Py_DECREF(module_name);
4429
4430 if (cls == NULL)
4431 return -1;
4432
4433 if ((args = Pdata_poptuple(self->stack, i)) != NULL) {
4434 obj = instantiate(cls, args);
4435 Py_DECREF(args);
4436 }
4437 Py_DECREF(cls);
4438
4439 if (obj == NULL)
4440 return -1;
4441
4442 PDATA_PUSH(self->stack, obj, -1);
4443 return 0;
4444}
4445
4446static int
4447load_newobj(UnpicklerObject *self)
4448{
4449 PyObject *args = NULL;
4450 PyObject *clsraw = NULL;
4451 PyTypeObject *cls; /* clsraw cast to its true type */
4452 PyObject *obj;
4453
4454 /* Stack is ... cls argtuple, and we want to call
4455 * cls.__new__(cls, *argtuple).
4456 */
4457 PDATA_POP(self->stack, args);
4458 if (args == NULL)
4459 goto error;
4460 if (!PyTuple_Check(args)) {
4461 PyErr_SetString(UnpicklingError, "NEWOBJ expected an arg " "tuple.");
4462 goto error;
4463 }
4464
4465 PDATA_POP(self->stack, clsraw);
4466 cls = (PyTypeObject *)clsraw;
4467 if (cls == NULL)
4468 goto error;
4469 if (!PyType_Check(cls)) {
4470 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4471 "isn't a type object");
4472 goto error;
4473 }
4474 if (cls->tp_new == NULL) {
4475 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4476 "has NULL tp_new");
4477 goto error;
4478 }
4479
4480 /* Call __new__. */
4481 obj = cls->tp_new(cls, args, NULL);
4482 if (obj == NULL)
4483 goto error;
4484
4485 Py_DECREF(args);
4486 Py_DECREF(clsraw);
4487 PDATA_PUSH(self->stack, obj, -1);
4488 return 0;
4489
4490 error:
4491 Py_XDECREF(args);
4492 Py_XDECREF(clsraw);
4493 return -1;
4494}
4495
4496static int
4497load_global(UnpicklerObject *self)
4498{
4499 PyObject *global = NULL;
4500 PyObject *module_name;
4501 PyObject *global_name;
4502 Py_ssize_t len;
4503 char *s;
4504
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004505 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004506 return -1;
4507 if (len < 2)
4508 return bad_readline();
4509 module_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4510 if (!module_name)
4511 return -1;
4512
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004513 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004514 if (len < 2) {
4515 Py_DECREF(module_name);
4516 return bad_readline();
4517 }
4518 global_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4519 if (global_name) {
4520 global = find_class(self, module_name, global_name);
4521 Py_DECREF(global_name);
4522 }
4523 }
4524 Py_DECREF(module_name);
4525
4526 if (global == NULL)
4527 return -1;
4528 PDATA_PUSH(self->stack, global, -1);
4529 return 0;
4530}
4531
4532static int
4533load_persid(UnpicklerObject *self)
4534{
4535 PyObject *pid;
4536 Py_ssize_t len;
4537 char *s;
4538
4539 if (self->pers_func) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004540 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004541 return -1;
4542 if (len < 2)
4543 return bad_readline();
4544
4545 pid = PyBytes_FromStringAndSize(s, len - 1);
4546 if (pid == NULL)
4547 return -1;
4548
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004549 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004550 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004551 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004552 if (pid == NULL)
4553 return -1;
4554
4555 PDATA_PUSH(self->stack, pid, -1);
4556 return 0;
4557 }
4558 else {
4559 PyErr_SetString(UnpicklingError,
4560 "A load persistent id instruction was encountered,\n"
4561 "but no persistent_load function was specified.");
4562 return -1;
4563 }
4564}
4565
4566static int
4567load_binpersid(UnpicklerObject *self)
4568{
4569 PyObject *pid;
4570
4571 if (self->pers_func) {
4572 PDATA_POP(self->stack, pid);
4573 if (pid == NULL)
4574 return -1;
4575
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004576 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004577 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004578 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004579 if (pid == NULL)
4580 return -1;
4581
4582 PDATA_PUSH(self->stack, pid, -1);
4583 return 0;
4584 }
4585 else {
4586 PyErr_SetString(UnpicklingError,
4587 "A load persistent id instruction was encountered,\n"
4588 "but no persistent_load function was specified.");
4589 return -1;
4590 }
4591}
4592
4593static int
4594load_pop(UnpicklerObject *self)
4595{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004596 int len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004597
4598 /* Note that we split the (pickle.py) stack into two stacks,
4599 * an object stack and a mark stack. We have to be clever and
4600 * pop the right one. We do this by looking at the top of the
Collin Winter8ca69de2009-05-26 16:53:41 +00004601 * mark stack first, and only signalling a stack underflow if
4602 * the object stack is empty and the mark stack doesn't match
4603 * our expectations.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004604 */
Collin Winter8ca69de2009-05-26 16:53:41 +00004605 if (self->num_marks > 0 && self->marks[self->num_marks - 1] == len) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004606 self->num_marks--;
Antoine Pitrou01a15ea2010-01-07 17:57:31 +00004607 } else if (len > 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004608 len--;
4609 Py_DECREF(self->stack->data[len]);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004610 Py_SIZE(self->stack) = len;
Collin Winter8ca69de2009-05-26 16:53:41 +00004611 } else {
4612 return stack_underflow();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004613 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004614 return 0;
4615}
4616
4617static int
4618load_pop_mark(UnpicklerObject *self)
4619{
4620 int i;
4621
4622 if ((i = marker(self)) < 0)
4623 return -1;
4624
4625 Pdata_clear(self->stack, i);
4626
4627 return 0;
4628}
4629
4630static int
4631load_dup(UnpicklerObject *self)
4632{
4633 PyObject *last;
4634 int len;
4635
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004636 if ((len = Py_SIZE(self->stack)) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004637 return stack_underflow();
4638 last = self->stack->data[len - 1];
4639 PDATA_APPEND(self->stack, last, -1);
4640 return 0;
4641}
4642
4643static int
4644load_get(UnpicklerObject *self)
4645{
4646 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004647 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004648 Py_ssize_t len;
4649 char *s;
4650
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004651 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004652 return -1;
4653 if (len < 2)
4654 return bad_readline();
4655
4656 key = PyLong_FromString(s, NULL, 10);
4657 if (key == NULL)
4658 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004659 idx = PyLong_AsSsize_t(key);
4660 if (idx == -1 && PyErr_Occurred()) {
4661 Py_DECREF(key);
4662 return -1;
4663 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004664
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004665 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004666 if (value == NULL) {
4667 if (!PyErr_Occurred())
4668 PyErr_SetObject(PyExc_KeyError, key);
4669 Py_DECREF(key);
4670 return -1;
4671 }
4672 Py_DECREF(key);
4673
4674 PDATA_APPEND(self->stack, value, -1);
4675 return 0;
4676}
4677
4678static int
4679load_binget(UnpicklerObject *self)
4680{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004681 PyObject *value;
4682 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004683 char *s;
4684
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004685 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004686 return -1;
4687
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004688 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004689
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004690 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004691 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004692 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004693 if (!PyErr_Occurred())
4694 PyErr_SetObject(PyExc_KeyError, key);
4695 Py_DECREF(key);
4696 return -1;
4697 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004698
4699 PDATA_APPEND(self->stack, value, -1);
4700 return 0;
4701}
4702
4703static int
4704load_long_binget(UnpicklerObject *self)
4705{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004706 PyObject *value;
4707 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004708 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004709
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004710 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004711 return -1;
4712
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004713 idx = (long)Py_CHARMASK(s[0]);
4714 idx |= (long)Py_CHARMASK(s[1]) << 8;
4715 idx |= (long)Py_CHARMASK(s[2]) << 16;
4716 idx |= (long)Py_CHARMASK(s[3]) << 24;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004717
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004718 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004719 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004720 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004721 if (!PyErr_Occurred())
4722 PyErr_SetObject(PyExc_KeyError, key);
4723 Py_DECREF(key);
4724 return -1;
4725 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004726
4727 PDATA_APPEND(self->stack, value, -1);
4728 return 0;
4729}
4730
4731/* Push an object from the extension registry (EXT[124]). nbytes is
4732 * the number of bytes following the opcode, holding the index (code) value.
4733 */
4734static int
4735load_extension(UnpicklerObject *self, int nbytes)
4736{
4737 char *codebytes; /* the nbytes bytes after the opcode */
4738 long code; /* calc_binint returns long */
4739 PyObject *py_code; /* code as a Python int */
4740 PyObject *obj; /* the object to push */
4741 PyObject *pair; /* (module_name, class_name) */
4742 PyObject *module_name, *class_name;
4743
4744 assert(nbytes == 1 || nbytes == 2 || nbytes == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004745 if (_Unpickler_Read(self, &codebytes, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004746 return -1;
4747 code = calc_binint(codebytes, nbytes);
4748 if (code <= 0) { /* note that 0 is forbidden */
4749 /* Corrupt or hostile pickle. */
4750 PyErr_SetString(UnpicklingError, "EXT specifies code <= 0");
4751 return -1;
4752 }
4753
4754 /* Look for the code in the cache. */
4755 py_code = PyLong_FromLong(code);
4756 if (py_code == NULL)
4757 return -1;
4758 obj = PyDict_GetItem(extension_cache, py_code);
4759 if (obj != NULL) {
4760 /* Bingo. */
4761 Py_DECREF(py_code);
4762 PDATA_APPEND(self->stack, obj, -1);
4763 return 0;
4764 }
4765
4766 /* Look up the (module_name, class_name) pair. */
4767 pair = PyDict_GetItem(inverted_registry, py_code);
4768 if (pair == NULL) {
4769 Py_DECREF(py_code);
4770 PyErr_Format(PyExc_ValueError, "unregistered extension "
4771 "code %ld", code);
4772 return -1;
4773 }
4774 /* Since the extension registry is manipulable via Python code,
4775 * confirm that pair is really a 2-tuple of strings.
4776 */
4777 if (!PyTuple_Check(pair) || PyTuple_Size(pair) != 2 ||
4778 !PyUnicode_Check(module_name = PyTuple_GET_ITEM(pair, 0)) ||
4779 !PyUnicode_Check(class_name = PyTuple_GET_ITEM(pair, 1))) {
4780 Py_DECREF(py_code);
4781 PyErr_Format(PyExc_ValueError, "_inverted_registry[%ld] "
4782 "isn't a 2-tuple of strings", code);
4783 return -1;
4784 }
4785 /* Load the object. */
4786 obj = find_class(self, module_name, class_name);
4787 if (obj == NULL) {
4788 Py_DECREF(py_code);
4789 return -1;
4790 }
4791 /* Cache code -> obj. */
4792 code = PyDict_SetItem(extension_cache, py_code, obj);
4793 Py_DECREF(py_code);
4794 if (code < 0) {
4795 Py_DECREF(obj);
4796 return -1;
4797 }
4798 PDATA_PUSH(self->stack, obj, -1);
4799 return 0;
4800}
4801
4802static int
4803load_put(UnpicklerObject *self)
4804{
4805 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004806 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004807 Py_ssize_t len;
4808 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004809
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004810 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004811 return -1;
4812 if (len < 2)
4813 return bad_readline();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004814 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004815 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004816 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004817
4818 key = PyLong_FromString(s, NULL, 10);
4819 if (key == NULL)
4820 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004821 idx = PyLong_AsSsize_t(key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004822 Py_DECREF(key);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004823 if (idx == -1 && PyErr_Occurred())
4824 return -1;
4825
4826 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004827}
4828
4829static int
4830load_binput(UnpicklerObject *self)
4831{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004832 PyObject *value;
4833 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004834 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004835
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004836 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004837 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004838
4839 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004840 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004841 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004842
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004843 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004844
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004845 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004846}
4847
4848static int
4849load_long_binput(UnpicklerObject *self)
4850{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004851 PyObject *value;
4852 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004853 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004854
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004855 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004856 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004857
4858 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004859 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004860 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004861
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004862 idx = (long)Py_CHARMASK(s[0]);
4863 idx |= (long)Py_CHARMASK(s[1]) << 8;
4864 idx |= (long)Py_CHARMASK(s[2]) << 16;
4865 idx |= (long)Py_CHARMASK(s[3]) << 24;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004866
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004867 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004868}
4869
4870static int
4871do_append(UnpicklerObject *self, int x)
4872{
4873 PyObject *value;
4874 PyObject *list;
4875 int len, i;
4876
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004877 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004878 if (x > len || x <= 0)
4879 return stack_underflow();
4880 if (len == x) /* nothing to do */
4881 return 0;
4882
4883 list = self->stack->data[x - 1];
4884
4885 if (PyList_Check(list)) {
4886 PyObject *slice;
4887 Py_ssize_t list_len;
4888
4889 slice = Pdata_poplist(self->stack, x);
4890 if (!slice)
4891 return -1;
4892 list_len = PyList_GET_SIZE(list);
4893 i = PyList_SetSlice(list, list_len, list_len, slice);
4894 Py_DECREF(slice);
4895 return i;
4896 }
4897 else {
4898 PyObject *append_func;
4899
4900 append_func = PyObject_GetAttrString(list, "append");
4901 if (append_func == NULL)
4902 return -1;
4903 for (i = x; i < len; i++) {
4904 PyObject *result;
4905
4906 value = self->stack->data[i];
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004907 result = _Unpickler_FastCall(self, append_func, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004908 if (result == NULL) {
4909 Pdata_clear(self->stack, i + 1);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004910 Py_SIZE(self->stack) = x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004911 return -1;
4912 }
4913 Py_DECREF(result);
4914 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004915 Py_SIZE(self->stack) = x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004916 }
4917
4918 return 0;
4919}
4920
4921static int
4922load_append(UnpicklerObject *self)
4923{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004924 return do_append(self, Py_SIZE(self->stack) - 1);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004925}
4926
4927static int
4928load_appends(UnpicklerObject *self)
4929{
4930 return do_append(self, marker(self));
4931}
4932
4933static int
4934do_setitems(UnpicklerObject *self, int x)
4935{
4936 PyObject *value, *key;
4937 PyObject *dict;
4938 int len, i;
4939 int status = 0;
4940
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004941 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004942 if (x > len || x <= 0)
4943 return stack_underflow();
4944 if (len == x) /* nothing to do */
4945 return 0;
4946 if ((len - x) % 2 != 0) {
4947 /* Currupt or hostile pickle -- we never write one like this. */
4948 PyErr_SetString(UnpicklingError, "odd number of items for SETITEMS");
4949 return -1;
4950 }
4951
4952 /* Here, dict does not actually need to be a PyDict; it could be anything
4953 that supports the __setitem__ attribute. */
4954 dict = self->stack->data[x - 1];
4955
4956 for (i = x + 1; i < len; i += 2) {
4957 key = self->stack->data[i - 1];
4958 value = self->stack->data[i];
4959 if (PyObject_SetItem(dict, key, value) < 0) {
4960 status = -1;
4961 break;
4962 }
4963 }
4964
4965 Pdata_clear(self->stack, x);
4966 return status;
4967}
4968
4969static int
4970load_setitem(UnpicklerObject *self)
4971{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004972 return do_setitems(self, Py_SIZE(self->stack) - 2);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004973}
4974
4975static int
4976load_setitems(UnpicklerObject *self)
4977{
4978 return do_setitems(self, marker(self));
4979}
4980
4981static int
4982load_build(UnpicklerObject *self)
4983{
4984 PyObject *state, *inst, *slotstate;
4985 PyObject *setstate;
4986 int status = 0;
4987
4988 /* Stack is ... instance, state. We want to leave instance at
4989 * the stack top, possibly mutated via instance.__setstate__(state).
4990 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004991 if (Py_SIZE(self->stack) < 2)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004992 return stack_underflow();
4993
4994 PDATA_POP(self->stack, state);
4995 if (state == NULL)
4996 return -1;
4997
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004998 inst = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004999
5000 setstate = PyObject_GetAttrString(inst, "__setstate__");
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005001 if (setstate == NULL) {
5002 if (PyErr_ExceptionMatches(PyExc_AttributeError))
5003 PyErr_Clear();
Antoine Pitroud79dc622008-09-05 00:03:33 +00005004 else {
5005 Py_DECREF(state);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005006 return -1;
Antoine Pitroud79dc622008-09-05 00:03:33 +00005007 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005008 }
5009 else {
5010 PyObject *result;
5011
5012 /* The explicit __setstate__ is responsible for everything. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005013 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Antoine Pitroud79dc622008-09-05 00:03:33 +00005014 reference to state first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005015 result = _Unpickler_FastCall(self, setstate, state);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005016 Py_DECREF(setstate);
5017 if (result == NULL)
5018 return -1;
5019 Py_DECREF(result);
5020 return 0;
5021 }
5022
5023 /* A default __setstate__. First see whether state embeds a
5024 * slot state dict too (a proto 2 addition).
5025 */
5026 if (PyTuple_Check(state) && Py_SIZE(state) == 2) {
5027 PyObject *tmp = state;
5028
5029 state = PyTuple_GET_ITEM(tmp, 0);
5030 slotstate = PyTuple_GET_ITEM(tmp, 1);
5031 Py_INCREF(state);
5032 Py_INCREF(slotstate);
5033 Py_DECREF(tmp);
5034 }
5035 else
5036 slotstate = NULL;
5037
5038 /* Set inst.__dict__ from the state dict (if any). */
5039 if (state != Py_None) {
5040 PyObject *dict;
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005041 PyObject *d_key, *d_value;
5042 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005043
5044 if (!PyDict_Check(state)) {
5045 PyErr_SetString(UnpicklingError, "state is not a dictionary");
5046 goto error;
5047 }
5048 dict = PyObject_GetAttrString(inst, "__dict__");
5049 if (dict == NULL)
5050 goto error;
5051
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005052 i = 0;
5053 while (PyDict_Next(state, &i, &d_key, &d_value)) {
5054 /* normally the keys for instance attributes are
5055 interned. we should try to do that here. */
5056 Py_INCREF(d_key);
5057 if (PyUnicode_CheckExact(d_key))
5058 PyUnicode_InternInPlace(&d_key);
5059 if (PyObject_SetItem(dict, d_key, d_value) < 0) {
5060 Py_DECREF(d_key);
5061 goto error;
5062 }
5063 Py_DECREF(d_key);
5064 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005065 Py_DECREF(dict);
5066 }
5067
5068 /* Also set instance attributes from the slotstate dict (if any). */
5069 if (slotstate != NULL) {
5070 PyObject *d_key, *d_value;
5071 Py_ssize_t i;
5072
5073 if (!PyDict_Check(slotstate)) {
5074 PyErr_SetString(UnpicklingError,
5075 "slot state is not a dictionary");
5076 goto error;
5077 }
5078 i = 0;
5079 while (PyDict_Next(slotstate, &i, &d_key, &d_value)) {
5080 if (PyObject_SetAttr(inst, d_key, d_value) < 0)
5081 goto error;
5082 }
5083 }
5084
5085 if (0) {
5086 error:
5087 status = -1;
5088 }
5089
5090 Py_DECREF(state);
5091 Py_XDECREF(slotstate);
5092 return status;
5093}
5094
5095static int
5096load_mark(UnpicklerObject *self)
5097{
5098
5099 /* Note that we split the (pickle.py) stack into two stacks, an
5100 * object stack and a mark stack. Here we push a mark onto the
5101 * mark stack.
5102 */
5103
5104 if ((self->num_marks + 1) >= self->marks_size) {
5105 size_t alloc;
5106 int *marks;
5107
5108 /* Use the size_t type to check for overflow. */
5109 alloc = ((size_t)self->num_marks << 1) + 20;
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00005110 if (alloc > PY_SSIZE_T_MAX ||
5111 alloc <= ((size_t)self->num_marks + 1)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005112 PyErr_NoMemory();
5113 return -1;
5114 }
5115
5116 if (self->marks == NULL)
5117 marks = (int *)PyMem_Malloc(alloc * sizeof(int));
5118 else
5119 marks = (int *)PyMem_Realloc(self->marks, alloc * sizeof(int));
5120 if (marks == NULL) {
5121 PyErr_NoMemory();
5122 return -1;
5123 }
5124 self->marks = marks;
5125 self->marks_size = (Py_ssize_t)alloc;
5126 }
5127
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005128 self->marks[self->num_marks++] = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005129
5130 return 0;
5131}
5132
5133static int
5134load_reduce(UnpicklerObject *self)
5135{
5136 PyObject *callable = NULL;
5137 PyObject *argtup = NULL;
5138 PyObject *obj = NULL;
5139
5140 PDATA_POP(self->stack, argtup);
5141 if (argtup == NULL)
5142 return -1;
5143 PDATA_POP(self->stack, callable);
5144 if (callable) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005145 obj = PyObject_CallObject(callable, argtup);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005146 Py_DECREF(callable);
5147 }
5148 Py_DECREF(argtup);
5149
5150 if (obj == NULL)
5151 return -1;
5152
5153 PDATA_PUSH(self->stack, obj, -1);
5154 return 0;
5155}
5156
5157/* Just raises an error if we don't know the protocol specified. PROTO
5158 * is the first opcode for protocols >= 2.
5159 */
5160static int
5161load_proto(UnpicklerObject *self)
5162{
5163 char *s;
5164 int i;
5165
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005166 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005167 return -1;
5168
5169 i = (unsigned char)s[0];
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005170 if (i <= HIGHEST_PROTOCOL) {
5171 self->proto = i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005172 return 0;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005173 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005174
5175 PyErr_Format(PyExc_ValueError, "unsupported pickle protocol: %d", i);
5176 return -1;
5177}
5178
5179static PyObject *
5180load(UnpicklerObject *self)
5181{
5182 PyObject *err;
5183 PyObject *value = NULL;
5184 char *s;
5185
5186 self->num_marks = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005187 if (Py_SIZE(self->stack))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005188 Pdata_clear(self->stack, 0);
5189
5190 /* Convenient macros for the dispatch while-switch loop just below. */
5191#define OP(opcode, load_func) \
5192 case opcode: if (load_func(self) < 0) break; continue;
5193
5194#define OP_ARG(opcode, load_func, arg) \
5195 case opcode: if (load_func(self, (arg)) < 0) break; continue;
5196
5197 while (1) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005198 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005199 break;
5200
5201 switch ((enum opcode)s[0]) {
5202 OP(NONE, load_none)
5203 OP(BININT, load_binint)
5204 OP(BININT1, load_binint1)
5205 OP(BININT2, load_binint2)
5206 OP(INT, load_int)
5207 OP(LONG, load_long)
5208 OP_ARG(LONG1, load_counted_long, 1)
5209 OP_ARG(LONG4, load_counted_long, 4)
5210 OP(FLOAT, load_float)
5211 OP(BINFLOAT, load_binfloat)
5212 OP(BINBYTES, load_binbytes)
5213 OP(SHORT_BINBYTES, load_short_binbytes)
5214 OP(BINSTRING, load_binstring)
5215 OP(SHORT_BINSTRING, load_short_binstring)
5216 OP(STRING, load_string)
5217 OP(UNICODE, load_unicode)
5218 OP(BINUNICODE, load_binunicode)
5219 OP_ARG(EMPTY_TUPLE, load_counted_tuple, 0)
5220 OP_ARG(TUPLE1, load_counted_tuple, 1)
5221 OP_ARG(TUPLE2, load_counted_tuple, 2)
5222 OP_ARG(TUPLE3, load_counted_tuple, 3)
5223 OP(TUPLE, load_tuple)
5224 OP(EMPTY_LIST, load_empty_list)
5225 OP(LIST, load_list)
5226 OP(EMPTY_DICT, load_empty_dict)
5227 OP(DICT, load_dict)
5228 OP(OBJ, load_obj)
5229 OP(INST, load_inst)
5230 OP(NEWOBJ, load_newobj)
5231 OP(GLOBAL, load_global)
5232 OP(APPEND, load_append)
5233 OP(APPENDS, load_appends)
5234 OP(BUILD, load_build)
5235 OP(DUP, load_dup)
5236 OP(BINGET, load_binget)
5237 OP(LONG_BINGET, load_long_binget)
5238 OP(GET, load_get)
5239 OP(MARK, load_mark)
5240 OP(BINPUT, load_binput)
5241 OP(LONG_BINPUT, load_long_binput)
5242 OP(PUT, load_put)
5243 OP(POP, load_pop)
5244 OP(POP_MARK, load_pop_mark)
5245 OP(SETITEM, load_setitem)
5246 OP(SETITEMS, load_setitems)
5247 OP(PERSID, load_persid)
5248 OP(BINPERSID, load_binpersid)
5249 OP(REDUCE, load_reduce)
5250 OP(PROTO, load_proto)
5251 OP_ARG(EXT1, load_extension, 1)
5252 OP_ARG(EXT2, load_extension, 2)
5253 OP_ARG(EXT4, load_extension, 4)
5254 OP_ARG(NEWTRUE, load_bool, Py_True)
5255 OP_ARG(NEWFALSE, load_bool, Py_False)
5256
5257 case STOP:
5258 break;
5259
5260 case '\0':
5261 PyErr_SetNone(PyExc_EOFError);
5262 return NULL;
5263
5264 default:
5265 PyErr_Format(UnpicklingError,
5266 "invalid load key, '%c'.", s[0]);
5267 return NULL;
5268 }
5269
5270 break; /* and we are done! */
5271 }
5272
Antoine Pitrou04248a82010-10-12 20:51:21 +00005273 if (_Unpickler_SkipConsumed(self) < 0)
5274 return NULL;
5275
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005276 /* XXX: It is not clear what this is actually for. */
5277 if ((err = PyErr_Occurred())) {
5278 if (err == PyExc_EOFError) {
5279 PyErr_SetNone(PyExc_EOFError);
5280 }
5281 return NULL;
5282 }
5283
5284 PDATA_POP(self->stack, value);
5285 return value;
5286}
5287
5288PyDoc_STRVAR(Unpickler_load_doc,
5289"load() -> object. Load a pickle."
5290"\n"
5291"Read a pickled object representation from the open file object given in\n"
5292"the constructor, and return the reconstituted object hierarchy specified\n"
5293"therein.\n");
5294
5295static PyObject *
5296Unpickler_load(UnpicklerObject *self)
5297{
5298 /* Check whether the Unpickler was initialized correctly. This prevents
5299 segfaulting if a subclass overridden __init__ with a function that does
5300 not call Unpickler.__init__(). Here, we simply ensure that self->read
5301 is not NULL. */
5302 if (self->read == NULL) {
5303 PyErr_Format(UnpicklingError,
5304 "Unpickler.__init__() was not called by %s.__init__()",
5305 Py_TYPE(self)->tp_name);
5306 return NULL;
5307 }
5308
5309 return load(self);
5310}
5311
5312/* The name of find_class() is misleading. In newer pickle protocols, this
5313 function is used for loading any global (i.e., functions), not just
5314 classes. The name is kept only for backward compatibility. */
5315
5316PyDoc_STRVAR(Unpickler_find_class_doc,
5317"find_class(module_name, global_name) -> object.\n"
5318"\n"
5319"Return an object from a specified module, importing the module if\n"
5320"necessary. Subclasses may override this method (e.g. to restrict\n"
5321"unpickling of arbitrary classes and functions).\n"
5322"\n"
5323"This method is called whenever a class or a function object is\n"
5324"needed. Both arguments passed are str objects.\n");
5325
5326static PyObject *
5327Unpickler_find_class(UnpicklerObject *self, PyObject *args)
5328{
5329 PyObject *global;
5330 PyObject *modules_dict;
5331 PyObject *module;
5332 PyObject *module_name, *global_name;
5333
5334 if (!PyArg_UnpackTuple(args, "find_class", 2, 2,
5335 &module_name, &global_name))
5336 return NULL;
5337
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005338 /* Try to map the old names used in Python 2.x to the new ones used in
5339 Python 3.x. We do this only with old pickle protocols and when the
5340 user has not disabled the feature. */
5341 if (self->proto < 3 && self->fix_imports) {
5342 PyObject *key;
5343 PyObject *item;
5344
5345 /* Check if the global (i.e., a function or a class) was renamed
5346 or moved to another module. */
5347 key = PyTuple_Pack(2, module_name, global_name);
5348 if (key == NULL)
5349 return NULL;
5350 item = PyDict_GetItemWithError(name_mapping_2to3, key);
5351 Py_DECREF(key);
5352 if (item) {
5353 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
5354 PyErr_Format(PyExc_RuntimeError,
5355 "_compat_pickle.NAME_MAPPING values should be "
5356 "2-tuples, not %.200s", Py_TYPE(item)->tp_name);
5357 return NULL;
5358 }
5359 module_name = PyTuple_GET_ITEM(item, 0);
5360 global_name = PyTuple_GET_ITEM(item, 1);
5361 if (!PyUnicode_Check(module_name) ||
5362 !PyUnicode_Check(global_name)) {
5363 PyErr_Format(PyExc_RuntimeError,
5364 "_compat_pickle.NAME_MAPPING values should be "
5365 "pairs of str, not (%.200s, %.200s)",
5366 Py_TYPE(module_name)->tp_name,
5367 Py_TYPE(global_name)->tp_name);
5368 return NULL;
5369 }
5370 }
5371 else if (PyErr_Occurred()) {
5372 return NULL;
5373 }
5374
5375 /* Check if the module was renamed. */
5376 item = PyDict_GetItemWithError(import_mapping_2to3, module_name);
5377 if (item) {
5378 if (!PyUnicode_Check(item)) {
5379 PyErr_Format(PyExc_RuntimeError,
5380 "_compat_pickle.IMPORT_MAPPING values should be "
5381 "strings, not %.200s", Py_TYPE(item)->tp_name);
5382 return NULL;
5383 }
5384 module_name = item;
5385 }
5386 else if (PyErr_Occurred()) {
5387 return NULL;
5388 }
5389 }
5390
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005391 modules_dict = PySys_GetObject("modules");
5392 if (modules_dict == NULL)
5393 return NULL;
5394
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005395 module = PyDict_GetItemWithError(modules_dict, module_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005396 if (module == NULL) {
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005397 if (PyErr_Occurred())
5398 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005399 module = PyImport_Import(module_name);
5400 if (module == NULL)
5401 return NULL;
5402 global = PyObject_GetAttr(module, global_name);
5403 Py_DECREF(module);
5404 }
5405 else {
5406 global = PyObject_GetAttr(module, global_name);
5407 }
5408 return global;
5409}
5410
5411static struct PyMethodDef Unpickler_methods[] = {
5412 {"load", (PyCFunction)Unpickler_load, METH_NOARGS,
5413 Unpickler_load_doc},
5414 {"find_class", (PyCFunction)Unpickler_find_class, METH_VARARGS,
5415 Unpickler_find_class_doc},
5416 {NULL, NULL} /* sentinel */
5417};
5418
5419static void
5420Unpickler_dealloc(UnpicklerObject *self)
5421{
5422 PyObject_GC_UnTrack((PyObject *)self);
5423 Py_XDECREF(self->readline);
5424 Py_XDECREF(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005425 Py_XDECREF(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005426 Py_XDECREF(self->stack);
5427 Py_XDECREF(self->pers_func);
5428 Py_XDECREF(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005429 if (self->buffer.buf != NULL) {
5430 PyBuffer_Release(&self->buffer);
5431 self->buffer.buf = NULL;
5432 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005433
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005434 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005435 PyMem_Free(self->marks);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005436 PyMem_Free(self->input_line);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005437 free(self->encoding);
5438 free(self->errors);
5439
5440 Py_TYPE(self)->tp_free((PyObject *)self);
5441}
5442
5443static int
5444Unpickler_traverse(UnpicklerObject *self, visitproc visit, void *arg)
5445{
5446 Py_VISIT(self->readline);
5447 Py_VISIT(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005448 Py_VISIT(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005449 Py_VISIT(self->stack);
5450 Py_VISIT(self->pers_func);
5451 Py_VISIT(self->arg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005452 return 0;
5453}
5454
5455static int
5456Unpickler_clear(UnpicklerObject *self)
5457{
5458 Py_CLEAR(self->readline);
5459 Py_CLEAR(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005460 Py_CLEAR(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005461 Py_CLEAR(self->stack);
5462 Py_CLEAR(self->pers_func);
5463 Py_CLEAR(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005464 if (self->buffer.buf != NULL) {
5465 PyBuffer_Release(&self->buffer);
5466 self->buffer.buf = NULL;
5467 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005468
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005469 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005470 PyMem_Free(self->marks);
5471 self->marks = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005472 PyMem_Free(self->input_line);
5473 self->input_line = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005474 free(self->encoding);
5475 self->encoding = NULL;
5476 free(self->errors);
5477 self->errors = NULL;
5478
5479 return 0;
5480}
5481
5482PyDoc_STRVAR(Unpickler_doc,
5483"Unpickler(file, *, encoding='ASCII', errors='strict')"
5484"\n"
5485"This takes a binary file for reading a pickle data stream.\n"
5486"\n"
5487"The protocol version of the pickle is detected automatically, so no\n"
5488"proto argument is needed.\n"
5489"\n"
5490"The file-like object must have two methods, a read() method\n"
5491"that takes an integer argument, and a readline() method that\n"
5492"requires no arguments. Both methods should return bytes.\n"
5493"Thus file-like object can be a binary file object opened for\n"
5494"reading, a BytesIO object, or any other custom object that\n"
5495"meets this interface.\n"
5496"\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005497"Optional keyword arguments are *fix_imports*, *encoding* and *errors*,\n"
5498"which are used to control compatiblity support for pickle stream\n"
5499"generated by Python 2.x. If *fix_imports* is True, pickle will try to\n"
5500"map the old Python 2.x names to the new names used in Python 3.x. The\n"
5501"*encoding* and *errors* tell pickle how to decode 8-bit string\n"
5502"instances pickled by Python 2.x; these default to 'ASCII' and\n"
5503"'strict', respectively.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005504
5505static int
5506Unpickler_init(UnpicklerObject *self, PyObject *args, PyObject *kwds)
5507{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005508 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005509 PyObject *file;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005510 PyObject *fix_imports = Py_True;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005511 char *encoding = NULL;
5512 char *errors = NULL;
5513
5514 /* XXX: That is an horrible error message. But, I don't know how to do
5515 better... */
5516 if (Py_SIZE(args) != 1) {
5517 PyErr_Format(PyExc_TypeError,
5518 "%s takes exactly one positional argument (%zd given)",
5519 Py_TYPE(self)->tp_name, Py_SIZE(args));
5520 return -1;
5521 }
5522
5523 /* Arguments parsing needs to be done in the __init__() method to allow
5524 subclasses to define their own __init__() method, which may (or may
5525 not) support Unpickler arguments. However, this means we need to be
5526 extra careful in the other Unpickler methods, since a subclass could
5527 forget to call Unpickler.__init__() thus breaking our internal
5528 invariants. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005529 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:Unpickler", kwlist,
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005530 &file, &fix_imports, &encoding, &errors))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005531 return -1;
5532
5533 /* In case of multiple __init__() calls, clear previous content. */
5534 if (self->read != NULL)
5535 (void)Unpickler_clear(self);
5536
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005537 if (_Unpickler_SetInputStream(self, file) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005538 return -1;
5539
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005540 if (_Unpickler_SetInputEncoding(self, encoding, errors) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005541 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005542
5543 self->fix_imports = PyObject_IsTrue(fix_imports);
5544 if (self->fix_imports == -1)
5545 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005546
5547 if (PyObject_HasAttrString((PyObject *)self, "persistent_load")) {
5548 self->pers_func = PyObject_GetAttrString((PyObject *)self,
5549 "persistent_load");
5550 if (self->pers_func == NULL)
5551 return -1;
5552 }
5553 else {
5554 self->pers_func = NULL;
5555 }
5556
5557 self->stack = (Pdata *)Pdata_New();
5558 if (self->stack == NULL)
5559 return -1;
5560
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005561 self->memo_size = 32;
5562 self->memo = _Unpickler_NewMemo(self->memo_size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005563 if (self->memo == NULL)
5564 return -1;
5565
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005566 self->arg = NULL;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005567 self->proto = 0;
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005568
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005569 return 0;
5570}
5571
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005572/* Define a proxy object for the Unpickler's internal memo object. This is to
5573 * avoid breaking code like:
5574 * unpickler.memo.clear()
5575 * and
5576 * unpickler.memo = saved_memo
5577 * Is this a good idea? Not really, but we don't want to break code that uses
5578 * it. Note that we don't implement the entire mapping API here. This is
5579 * intentional, as these should be treated as black-box implementation details.
5580 *
5581 * We do, however, have to implement pickling/unpickling support because of
5582 * real-world code like cvs2svn.
5583 */
5584
5585typedef struct {
5586 PyObject_HEAD
5587 UnpicklerObject *unpickler;
5588} UnpicklerMemoProxyObject;
5589
5590PyDoc_STRVAR(ump_clear_doc,
5591"memo.clear() -> None. Remove all items from memo.");
5592
5593static PyObject *
5594ump_clear(UnpicklerMemoProxyObject *self)
5595{
5596 _Unpickler_MemoCleanup(self->unpickler);
5597 self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size);
5598 if (self->unpickler->memo == NULL)
5599 return NULL;
5600 Py_RETURN_NONE;
5601}
5602
5603PyDoc_STRVAR(ump_copy_doc,
5604"memo.copy() -> new_memo. Copy the memo to a new object.");
5605
5606static PyObject *
5607ump_copy(UnpicklerMemoProxyObject *self)
5608{
5609 Py_ssize_t i;
5610 PyObject *new_memo = PyDict_New();
5611 if (new_memo == NULL)
5612 return NULL;
5613
5614 for (i = 0; i < self->unpickler->memo_size; i++) {
5615 int status;
5616 PyObject *key, *value;
5617
5618 value = self->unpickler->memo[i];
5619 if (value == NULL)
5620 continue;
5621
5622 key = PyLong_FromSsize_t(i);
5623 if (key == NULL)
5624 goto error;
5625 status = PyDict_SetItem(new_memo, key, value);
5626 Py_DECREF(key);
5627 if (status < 0)
5628 goto error;
5629 }
5630 return new_memo;
5631
5632error:
5633 Py_DECREF(new_memo);
5634 return NULL;
5635}
5636
5637PyDoc_STRVAR(ump_reduce_doc,
5638"memo.__reduce__(). Pickling support.");
5639
5640static PyObject *
5641ump_reduce(UnpicklerMemoProxyObject *self, PyObject *args)
5642{
5643 PyObject *reduce_value;
5644 PyObject *constructor_args;
5645 PyObject *contents = ump_copy(self);
5646 if (contents == NULL)
5647 return NULL;
5648
5649 reduce_value = PyTuple_New(2);
5650 if (reduce_value == NULL) {
5651 Py_DECREF(contents);
5652 return NULL;
5653 }
5654 constructor_args = PyTuple_New(1);
5655 if (constructor_args == NULL) {
5656 Py_DECREF(contents);
5657 Py_DECREF(reduce_value);
5658 return NULL;
5659 }
5660 PyTuple_SET_ITEM(constructor_args, 0, contents);
5661 Py_INCREF((PyObject *)&PyDict_Type);
5662 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
5663 PyTuple_SET_ITEM(reduce_value, 1, constructor_args);
5664 return reduce_value;
5665}
5666
5667static PyMethodDef unpicklerproxy_methods[] = {
5668 {"clear", (PyCFunction)ump_clear, METH_NOARGS, ump_clear_doc},
5669 {"copy", (PyCFunction)ump_copy, METH_NOARGS, ump_copy_doc},
5670 {"__reduce__", (PyCFunction)ump_reduce, METH_VARARGS, ump_reduce_doc},
5671 {NULL, NULL} /* sentinel */
5672};
5673
5674static void
5675UnpicklerMemoProxy_dealloc(UnpicklerMemoProxyObject *self)
5676{
5677 PyObject_GC_UnTrack(self);
5678 Py_XDECREF(self->unpickler);
5679 PyObject_GC_Del((PyObject *)self);
5680}
5681
5682static int
5683UnpicklerMemoProxy_traverse(UnpicklerMemoProxyObject *self,
5684 visitproc visit, void *arg)
5685{
5686 Py_VISIT(self->unpickler);
5687 return 0;
5688}
5689
5690static int
5691UnpicklerMemoProxy_clear(UnpicklerMemoProxyObject *self)
5692{
5693 Py_CLEAR(self->unpickler);
5694 return 0;
5695}
5696
5697static PyTypeObject UnpicklerMemoProxyType = {
5698 PyVarObject_HEAD_INIT(NULL, 0)
5699 "_pickle.UnpicklerMemoProxy", /*tp_name*/
5700 sizeof(UnpicklerMemoProxyObject), /*tp_basicsize*/
5701 0,
5702 (destructor)UnpicklerMemoProxy_dealloc, /* tp_dealloc */
5703 0, /* tp_print */
5704 0, /* tp_getattr */
5705 0, /* tp_setattr */
5706 0, /* tp_compare */
5707 0, /* tp_repr */
5708 0, /* tp_as_number */
5709 0, /* tp_as_sequence */
5710 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00005711 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005712 0, /* tp_call */
5713 0, /* tp_str */
5714 PyObject_GenericGetAttr, /* tp_getattro */
5715 PyObject_GenericSetAttr, /* tp_setattro */
5716 0, /* tp_as_buffer */
5717 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5718 0, /* tp_doc */
5719 (traverseproc)UnpicklerMemoProxy_traverse, /* tp_traverse */
5720 (inquiry)UnpicklerMemoProxy_clear, /* tp_clear */
5721 0, /* tp_richcompare */
5722 0, /* tp_weaklistoffset */
5723 0, /* tp_iter */
5724 0, /* tp_iternext */
5725 unpicklerproxy_methods, /* tp_methods */
5726};
5727
5728static PyObject *
5729UnpicklerMemoProxy_New(UnpicklerObject *unpickler)
5730{
5731 UnpicklerMemoProxyObject *self;
5732
5733 self = PyObject_GC_New(UnpicklerMemoProxyObject,
5734 &UnpicklerMemoProxyType);
5735 if (self == NULL)
5736 return NULL;
5737 Py_INCREF(unpickler);
5738 self->unpickler = unpickler;
5739 PyObject_GC_Track(self);
5740 return (PyObject *)self;
5741}
5742
5743/*****************************************************************************/
5744
5745
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005746static PyObject *
5747Unpickler_get_memo(UnpicklerObject *self)
5748{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005749 return UnpicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005750}
5751
5752static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005753Unpickler_set_memo(UnpicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005754{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005755 PyObject **new_memo;
5756 Py_ssize_t new_memo_size = 0;
5757 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005758
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005759 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005760 PyErr_SetString(PyExc_TypeError,
5761 "attribute deletion is not supported");
5762 return -1;
5763 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005764
5765 if (Py_TYPE(obj) == &UnpicklerMemoProxyType) {
5766 UnpicklerObject *unpickler =
5767 ((UnpicklerMemoProxyObject *)obj)->unpickler;
5768
5769 new_memo_size = unpickler->memo_size;
5770 new_memo = _Unpickler_NewMemo(new_memo_size);
5771 if (new_memo == NULL)
5772 return -1;
5773
5774 for (i = 0; i < new_memo_size; i++) {
5775 Py_XINCREF(unpickler->memo[i]);
5776 new_memo[i] = unpickler->memo[i];
5777 }
5778 }
5779 else if (PyDict_Check(obj)) {
5780 Py_ssize_t i = 0;
5781 PyObject *key, *value;
5782
5783 new_memo_size = PyDict_Size(obj);
5784 new_memo = _Unpickler_NewMemo(new_memo_size);
5785 if (new_memo == NULL)
5786 return -1;
5787
5788 while (PyDict_Next(obj, &i, &key, &value)) {
5789 Py_ssize_t idx;
5790 if (!PyLong_Check(key)) {
5791 PyErr_SetString(PyExc_TypeError,
5792 "memo key must be integers");
5793 goto error;
5794 }
5795 idx = PyLong_AsSsize_t(key);
5796 if (idx == -1 && PyErr_Occurred())
5797 goto error;
5798 if (_Unpickler_MemoPut(self, idx, value) < 0)
5799 goto error;
5800 }
5801 }
5802 else {
5803 PyErr_Format(PyExc_TypeError,
5804 "'memo' attribute must be an UnpicklerMemoProxy object"
5805 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005806 return -1;
5807 }
5808
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005809 _Unpickler_MemoCleanup(self);
5810 self->memo_size = new_memo_size;
5811 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005812
5813 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005814
5815 error:
5816 if (new_memo_size) {
5817 i = new_memo_size;
5818 while (--i >= 0) {
5819 Py_XDECREF(new_memo[i]);
5820 }
5821 PyMem_FREE(new_memo);
5822 }
5823 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005824}
5825
5826static PyObject *
5827Unpickler_get_persload(UnpicklerObject *self)
5828{
5829 if (self->pers_func == NULL)
5830 PyErr_SetString(PyExc_AttributeError, "persistent_load");
5831 else
5832 Py_INCREF(self->pers_func);
5833 return self->pers_func;
5834}
5835
5836static int
5837Unpickler_set_persload(UnpicklerObject *self, PyObject *value)
5838{
5839 PyObject *tmp;
5840
5841 if (value == NULL) {
5842 PyErr_SetString(PyExc_TypeError,
5843 "attribute deletion is not supported");
5844 return -1;
5845 }
5846 if (!PyCallable_Check(value)) {
5847 PyErr_SetString(PyExc_TypeError,
5848 "persistent_load must be a callable taking "
5849 "one argument");
5850 return -1;
5851 }
5852
5853 tmp = self->pers_func;
5854 Py_INCREF(value);
5855 self->pers_func = value;
5856 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
5857
5858 return 0;
5859}
5860
5861static PyGetSetDef Unpickler_getsets[] = {
5862 {"memo", (getter)Unpickler_get_memo, (setter)Unpickler_set_memo},
5863 {"persistent_load", (getter)Unpickler_get_persload,
5864 (setter)Unpickler_set_persload},
5865 {NULL}
5866};
5867
5868static PyTypeObject Unpickler_Type = {
5869 PyVarObject_HEAD_INIT(NULL, 0)
5870 "_pickle.Unpickler", /*tp_name*/
5871 sizeof(UnpicklerObject), /*tp_basicsize*/
5872 0, /*tp_itemsize*/
5873 (destructor)Unpickler_dealloc, /*tp_dealloc*/
5874 0, /*tp_print*/
5875 0, /*tp_getattr*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005876 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00005877 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005878 0, /*tp_repr*/
5879 0, /*tp_as_number*/
5880 0, /*tp_as_sequence*/
5881 0, /*tp_as_mapping*/
5882 0, /*tp_hash*/
5883 0, /*tp_call*/
5884 0, /*tp_str*/
5885 0, /*tp_getattro*/
5886 0, /*tp_setattro*/
5887 0, /*tp_as_buffer*/
5888 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5889 Unpickler_doc, /*tp_doc*/
5890 (traverseproc)Unpickler_traverse, /*tp_traverse*/
5891 (inquiry)Unpickler_clear, /*tp_clear*/
5892 0, /*tp_richcompare*/
5893 0, /*tp_weaklistoffset*/
5894 0, /*tp_iter*/
5895 0, /*tp_iternext*/
5896 Unpickler_methods, /*tp_methods*/
5897 0, /*tp_members*/
5898 Unpickler_getsets, /*tp_getset*/
5899 0, /*tp_base*/
5900 0, /*tp_dict*/
5901 0, /*tp_descr_get*/
5902 0, /*tp_descr_set*/
5903 0, /*tp_dictoffset*/
5904 (initproc)Unpickler_init, /*tp_init*/
5905 PyType_GenericAlloc, /*tp_alloc*/
5906 PyType_GenericNew, /*tp_new*/
5907 PyObject_GC_Del, /*tp_free*/
5908 0, /*tp_is_gc*/
5909};
5910
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005911PyDoc_STRVAR(pickle_dump_doc,
5912"dump(obj, file, protocol=None, *, fix_imports=True) -> None\n"
5913"\n"
5914"Write a pickled representation of obj to the open file object file. This\n"
5915"is equivalent to ``Pickler(file, protocol).dump(obj)``, but may be more\n"
5916"efficient.\n"
5917"\n"
5918"The optional protocol argument tells the pickler to use the given protocol;\n"
5919"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
5920"backward-incompatible protocol designed for Python 3.0.\n"
5921"\n"
5922"Specifying a negative protocol version selects the highest protocol version\n"
5923"supported. The higher the protocol used, the more recent the version of\n"
5924"Python needed to read the pickle produced.\n"
5925"\n"
5926"The file argument must have a write() method that accepts a single bytes\n"
5927"argument. It can thus be a file object opened for binary writing, a\n"
5928"io.BytesIO instance, or any other custom object that meets this interface.\n"
5929"\n"
5930"If fix_imports is True and protocol is less than 3, pickle will try to\n"
5931"map the new Python 3.x names to the old module names used in Python 2.x,\n"
5932"so that the pickle data stream is readable with Python 2.x.\n");
5933
5934static PyObject *
5935pickle_dump(PyObject *self, PyObject *args, PyObject *kwds)
5936{
5937 static char *kwlist[] = {"obj", "file", "protocol", "fix_imports", 0};
5938 PyObject *obj;
5939 PyObject *file;
5940 PyObject *proto = NULL;
5941 PyObject *fix_imports = Py_True;
5942 PicklerObject *pickler;
5943
5944 /* fix_imports is a keyword-only argument. */
5945 if (Py_SIZE(args) > 3) {
5946 PyErr_Format(PyExc_TypeError,
5947 "pickle.dump() takes at most 3 positional "
5948 "argument (%zd given)", Py_SIZE(args));
5949 return NULL;
5950 }
5951
5952 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|OO:dump", kwlist,
5953 &obj, &file, &proto, &fix_imports))
5954 return NULL;
5955
5956 pickler = _Pickler_New();
5957 if (pickler == NULL)
5958 return NULL;
5959
5960 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
5961 goto error;
5962
5963 if (_Pickler_SetOutputStream(pickler, file) < 0)
5964 goto error;
5965
5966 if (dump(pickler, obj) < 0)
5967 goto error;
5968
5969 if (_Pickler_FlushToFile(pickler) < 0)
5970 goto error;
5971
5972 Py_DECREF(pickler);
5973 Py_RETURN_NONE;
5974
5975 error:
5976 Py_XDECREF(pickler);
5977 return NULL;
5978}
5979
5980PyDoc_STRVAR(pickle_dumps_doc,
5981"dumps(obj, protocol=None, *, fix_imports=True) -> bytes\n"
5982"\n"
5983"Return the pickled representation of the object as a bytes\n"
5984"object, instead of writing it to a file.\n"
5985"\n"
5986"The optional protocol argument tells the pickler to use the given protocol;\n"
5987"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
5988"backward-incompatible protocol designed for Python 3.0.\n"
5989"\n"
5990"Specifying a negative protocol version selects the highest protocol version\n"
5991"supported. The higher the protocol used, the more recent the version of\n"
5992"Python needed to read the pickle produced.\n"
5993"\n"
5994"If fix_imports is True and *protocol* is less than 3, pickle will try to\n"
5995"map the new Python 3.x names to the old module names used in Python 2.x,\n"
5996"so that the pickle data stream is readable with Python 2.x.\n");
5997
5998static PyObject *
5999pickle_dumps(PyObject *self, PyObject *args, PyObject *kwds)
6000{
6001 static char *kwlist[] = {"obj", "protocol", "fix_imports", 0};
6002 PyObject *obj;
6003 PyObject *proto = NULL;
6004 PyObject *result;
6005 PyObject *fix_imports = Py_True;
6006 PicklerObject *pickler;
6007
6008 /* fix_imports is a keyword-only argument. */
6009 if (Py_SIZE(args) > 2) {
6010 PyErr_Format(PyExc_TypeError,
6011 "pickle.dumps() takes at most 2 positional "
6012 "argument (%zd given)", Py_SIZE(args));
6013 return NULL;
6014 }
6015
6016 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:dumps", kwlist,
6017 &obj, &proto, &fix_imports))
6018 return NULL;
6019
6020 pickler = _Pickler_New();
6021 if (pickler == NULL)
6022 return NULL;
6023
6024 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6025 goto error;
6026
6027 if (dump(pickler, obj) < 0)
6028 goto error;
6029
6030 result = _Pickler_GetString(pickler);
6031 Py_DECREF(pickler);
6032 return result;
6033
6034 error:
6035 Py_XDECREF(pickler);
6036 return NULL;
6037}
6038
6039PyDoc_STRVAR(pickle_load_doc,
6040"load(file, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6041"\n"
6042"Read a pickled object representation from the open file object file and\n"
6043"return the reconstituted object hierarchy specified therein. This is\n"
6044"equivalent to ``Unpickler(file).load()``, but may be more efficient.\n"
6045"\n"
6046"The protocol version of the pickle is detected automatically, so no protocol\n"
6047"argument is needed. Bytes past the pickled object's representation are\n"
6048"ignored.\n"
6049"\n"
6050"The argument file must have two methods, a read() method that takes an\n"
6051"integer argument, and a readline() method that requires no arguments. Both\n"
6052"methods should return bytes. Thus *file* can be a binary file object opened\n"
6053"for reading, a BytesIO object, or any other custom object that meets this\n"
6054"interface.\n"
6055"\n"
6056"Optional keyword arguments are fix_imports, encoding and errors,\n"
6057"which are used to control compatiblity support for pickle stream generated\n"
6058"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6059"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6060"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6061"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6062
6063static PyObject *
6064pickle_load(PyObject *self, PyObject *args, PyObject *kwds)
6065{
6066 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
6067 PyObject *file;
6068 PyObject *fix_imports = Py_True;
6069 PyObject *result;
6070 char *encoding = NULL;
6071 char *errors = NULL;
6072 UnpicklerObject *unpickler;
6073
6074 /* fix_imports, encoding and errors are a keyword-only argument. */
6075 if (Py_SIZE(args) != 1) {
6076 PyErr_Format(PyExc_TypeError,
6077 "pickle.load() takes exactly one positional "
6078 "argument (%zd given)", Py_SIZE(args));
6079 return NULL;
6080 }
6081
6082 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:load", kwlist,
6083 &file, &fix_imports, &encoding, &errors))
6084 return NULL;
6085
6086 unpickler = _Unpickler_New();
6087 if (unpickler == NULL)
6088 return NULL;
6089
6090 if (_Unpickler_SetInputStream(unpickler, file) < 0)
6091 goto error;
6092
6093 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6094 goto error;
6095
6096 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6097 if (unpickler->fix_imports == -1)
6098 goto error;
6099
6100 result = load(unpickler);
6101 Py_DECREF(unpickler);
6102 return result;
6103
6104 error:
6105 Py_XDECREF(unpickler);
6106 return NULL;
6107}
6108
6109PyDoc_STRVAR(pickle_loads_doc,
6110"loads(input, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6111"\n"
6112"Read a pickled object hierarchy from a bytes object and return the\n"
6113"reconstituted object hierarchy specified therein\n"
6114"\n"
6115"The protocol version of the pickle is detected automatically, so no protocol\n"
6116"argument is needed. Bytes past the pickled object's representation are\n"
6117"ignored.\n"
6118"\n"
6119"Optional keyword arguments are fix_imports, encoding and errors, which\n"
6120"are used to control compatiblity support for pickle stream generated\n"
6121"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6122"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6123"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6124"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6125
6126static PyObject *
6127pickle_loads(PyObject *self, PyObject *args, PyObject *kwds)
6128{
6129 static char *kwlist[] = {"input", "fix_imports", "encoding", "errors", 0};
6130 PyObject *input;
6131 PyObject *fix_imports = Py_True;
6132 PyObject *result;
6133 char *encoding = NULL;
6134 char *errors = NULL;
6135 UnpicklerObject *unpickler;
6136
6137 /* fix_imports, encoding and errors are a keyword-only argument. */
6138 if (Py_SIZE(args) != 1) {
6139 PyErr_Format(PyExc_TypeError,
6140 "pickle.loads() takes exactly one positional "
6141 "argument (%zd given)", Py_SIZE(args));
6142 return NULL;
6143 }
6144
6145 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:loads", kwlist,
6146 &input, &fix_imports, &encoding, &errors))
6147 return NULL;
6148
6149 unpickler = _Unpickler_New();
6150 if (unpickler == NULL)
6151 return NULL;
6152
6153 if (_Unpickler_SetStringInput(unpickler, input) < 0)
6154 goto error;
6155
6156 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6157 goto error;
6158
6159 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6160 if (unpickler->fix_imports == -1)
6161 goto error;
6162
6163 result = load(unpickler);
6164 Py_DECREF(unpickler);
6165 return result;
6166
6167 error:
6168 Py_XDECREF(unpickler);
6169 return NULL;
6170}
6171
6172
6173static struct PyMethodDef pickle_methods[] = {
6174 {"dump", (PyCFunction)pickle_dump, METH_VARARGS|METH_KEYWORDS,
6175 pickle_dump_doc},
6176 {"dumps", (PyCFunction)pickle_dumps, METH_VARARGS|METH_KEYWORDS,
6177 pickle_dumps_doc},
6178 {"load", (PyCFunction)pickle_load, METH_VARARGS|METH_KEYWORDS,
6179 pickle_load_doc},
6180 {"loads", (PyCFunction)pickle_loads, METH_VARARGS|METH_KEYWORDS,
6181 pickle_loads_doc},
6182 {NULL, NULL} /* sentinel */
6183};
6184
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006185static int
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006186initmodule(void)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006187{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006188 PyObject *copyreg = NULL;
6189 PyObject *compat_pickle = NULL;
6190
6191 /* XXX: We should ensure that the types of the dictionaries imported are
6192 exactly PyDict objects. Otherwise, it is possible to crash the pickle
6193 since we use the PyDict API directly to access these dictionaries. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006194
6195 copyreg = PyImport_ImportModule("copyreg");
6196 if (!copyreg)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006197 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006198 dispatch_table = PyObject_GetAttrString(copyreg, "dispatch_table");
6199 if (!dispatch_table)
6200 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006201 extension_registry = \
6202 PyObject_GetAttrString(copyreg, "_extension_registry");
6203 if (!extension_registry)
6204 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006205 inverted_registry = PyObject_GetAttrString(copyreg, "_inverted_registry");
6206 if (!inverted_registry)
6207 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006208 extension_cache = PyObject_GetAttrString(copyreg, "_extension_cache");
6209 if (!extension_cache)
6210 goto error;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006211 Py_CLEAR(copyreg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006212
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006213 /* Load the 2.x -> 3.x stdlib module mapping tables */
6214 compat_pickle = PyImport_ImportModule("_compat_pickle");
6215 if (!compat_pickle)
6216 goto error;
6217 name_mapping_2to3 = PyObject_GetAttrString(compat_pickle, "NAME_MAPPING");
6218 if (!name_mapping_2to3)
6219 goto error;
6220 if (!PyDict_CheckExact(name_mapping_2to3)) {
6221 PyErr_Format(PyExc_RuntimeError,
6222 "_compat_pickle.NAME_MAPPING should be a dict, not %.200s",
6223 Py_TYPE(name_mapping_2to3)->tp_name);
6224 goto error;
6225 }
6226 import_mapping_2to3 = PyObject_GetAttrString(compat_pickle,
6227 "IMPORT_MAPPING");
6228 if (!import_mapping_2to3)
6229 goto error;
6230 if (!PyDict_CheckExact(import_mapping_2to3)) {
6231 PyErr_Format(PyExc_RuntimeError,
6232 "_compat_pickle.IMPORT_MAPPING should be a dict, "
6233 "not %.200s", Py_TYPE(import_mapping_2to3)->tp_name);
6234 goto error;
6235 }
6236 /* ... and the 3.x -> 2.x mapping tables */
6237 name_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6238 "REVERSE_NAME_MAPPING");
6239 if (!name_mapping_3to2)
6240 goto error;
6241 if (!PyDict_CheckExact(name_mapping_3to2)) {
6242 PyErr_Format(PyExc_RuntimeError,
6243 "_compat_pickle.REVERSE_NAME_MAPPING shouldbe a dict, "
6244 "not %.200s", Py_TYPE(name_mapping_3to2)->tp_name);
6245 goto error;
6246 }
6247 import_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6248 "REVERSE_IMPORT_MAPPING");
6249 if (!import_mapping_3to2)
6250 goto error;
6251 if (!PyDict_CheckExact(import_mapping_3to2)) {
6252 PyErr_Format(PyExc_RuntimeError,
6253 "_compat_pickle.REVERSE_IMPORT_MAPPING should be a dict, "
6254 "not %.200s", Py_TYPE(import_mapping_3to2)->tp_name);
6255 goto error;
6256 }
6257 Py_CLEAR(compat_pickle);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006258
6259 empty_tuple = PyTuple_New(0);
6260 if (empty_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006261 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006262 two_tuple = PyTuple_New(2);
6263 if (two_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006264 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006265 /* We use this temp container with no regard to refcounts, or to
6266 * keeping containees alive. Exempt from GC, because we don't
6267 * want anything looking at two_tuple() by magic.
6268 */
6269 PyObject_GC_UnTrack(two_tuple);
6270
6271 return 0;
6272
6273 error:
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006274 Py_CLEAR(copyreg);
6275 Py_CLEAR(dispatch_table);
6276 Py_CLEAR(extension_registry);
6277 Py_CLEAR(inverted_registry);
6278 Py_CLEAR(extension_cache);
6279 Py_CLEAR(compat_pickle);
6280 Py_CLEAR(name_mapping_2to3);
6281 Py_CLEAR(import_mapping_2to3);
6282 Py_CLEAR(name_mapping_3to2);
6283 Py_CLEAR(import_mapping_3to2);
6284 Py_CLEAR(empty_tuple);
6285 Py_CLEAR(two_tuple);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006286 return -1;
6287}
6288
6289static struct PyModuleDef _picklemodule = {
6290 PyModuleDef_HEAD_INIT,
6291 "_pickle",
6292 pickle_module_doc,
6293 -1,
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006294 pickle_methods,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006295 NULL,
6296 NULL,
6297 NULL,
6298 NULL
6299};
6300
6301PyMODINIT_FUNC
6302PyInit__pickle(void)
6303{
6304 PyObject *m;
6305
6306 if (PyType_Ready(&Unpickler_Type) < 0)
6307 return NULL;
6308 if (PyType_Ready(&Pickler_Type) < 0)
6309 return NULL;
6310 if (PyType_Ready(&Pdata_Type) < 0)
6311 return NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006312 if (PyType_Ready(&PicklerMemoProxyType) < 0)
6313 return NULL;
6314 if (PyType_Ready(&UnpicklerMemoProxyType) < 0)
6315 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006316
6317 /* Create the module and add the functions. */
6318 m = PyModule_Create(&_picklemodule);
6319 if (m == NULL)
6320 return NULL;
6321
6322 if (PyModule_AddObject(m, "Pickler", (PyObject *)&Pickler_Type) < 0)
6323 return NULL;
6324 if (PyModule_AddObject(m, "Unpickler", (PyObject *)&Unpickler_Type) < 0)
6325 return NULL;
6326
6327 /* Initialize the exceptions. */
6328 PickleError = PyErr_NewException("_pickle.PickleError", NULL, NULL);
6329 if (PickleError == NULL)
6330 return NULL;
6331 PicklingError = \
6332 PyErr_NewException("_pickle.PicklingError", PickleError, NULL);
6333 if (PicklingError == NULL)
6334 return NULL;
6335 UnpicklingError = \
6336 PyErr_NewException("_pickle.UnpicklingError", PickleError, NULL);
6337 if (UnpicklingError == NULL)
6338 return NULL;
6339
6340 if (PyModule_AddObject(m, "PickleError", PickleError) < 0)
6341 return NULL;
6342 if (PyModule_AddObject(m, "PicklingError", PicklingError) < 0)
6343 return NULL;
6344 if (PyModule_AddObject(m, "UnpicklingError", UnpicklingError) < 0)
6345 return NULL;
6346
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006347 if (initmodule() < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006348 return NULL;
6349
6350 return m;
6351}