blob: fca9f2ce26c93ab62bf98b052259e7901a783427 [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',
74 SHORT_BINBYTES = 'C',
75};
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) */
106 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
Antoine Pitrou7eecffd2010-10-22 19:43:59 +00002843static PyObject *
2844get_class(PyObject *obj)
2845{
2846 PyObject *cls;
2847 static PyObject *str_class;
2848
2849 if (str_class == NULL) {
2850 str_class = PyUnicode_InternFromString("__class__");
2851 if (str_class == NULL)
2852 return NULL;
2853 }
2854 cls = PyObject_GetAttr(obj, str_class);
2855 if (cls == NULL) {
2856 if (PyErr_ExceptionMatches(PyExc_AttributeError)) {
2857 PyErr_Clear();
2858 cls = (PyObject *) Py_TYPE(obj);
2859 Py_INCREF(cls);
2860 }
2861 }
2862 return cls;
2863}
2864
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002865/* We're saving obj, and args is the 2-thru-5 tuple returned by the
2866 * appropriate __reduce__ method for obj.
2867 */
2868static int
2869save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
2870{
2871 PyObject *callable;
2872 PyObject *argtup;
2873 PyObject *state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002874 PyObject *listitems = Py_None;
2875 PyObject *dictitems = Py_None;
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002876 Py_ssize_t size;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002877
2878 int use_newobj = self->proto >= 2;
2879
2880 const char reduce_op = REDUCE;
2881 const char build_op = BUILD;
2882 const char newobj_op = NEWOBJ;
2883
Hirokazu Yamamotob46a6332008-11-04 00:35:10 +00002884 size = PyTuple_Size(args);
2885 if (size < 2 || size > 5) {
2886 PyErr_SetString(PicklingError, "tuple returned by "
2887 "__reduce__ must contain 2 through 5 elements");
2888 return -1;
2889 }
2890
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002891 if (!PyArg_UnpackTuple(args, "save_reduce", 2, 5,
2892 &callable, &argtup, &state, &listitems, &dictitems))
2893 return -1;
2894
2895 if (!PyCallable_Check(callable)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002896 PyErr_SetString(PicklingError, "first item of the tuple "
2897 "returned by __reduce__ must be callable");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002898 return -1;
2899 }
2900 if (!PyTuple_Check(argtup)) {
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002901 PyErr_SetString(PicklingError, "second item of the tuple "
2902 "returned by __reduce__ must be a tuple");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002903 return -1;
2904 }
2905
2906 if (state == Py_None)
2907 state = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002908
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002909 if (listitems == Py_None)
2910 listitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002911 else if (!PyIter_Check(listitems)) {
2912 PyErr_Format(PicklingError, "Fourth element of tuple"
2913 "returned by __reduce__ must be an iterator, not %s",
2914 Py_TYPE(listitems)->tp_name);
2915 return -1;
2916 }
2917
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002918 if (dictitems == Py_None)
2919 dictitems = NULL;
Amaury Forgeot d'Arc424b4812008-10-30 22:25:31 +00002920 else if (!PyIter_Check(dictitems)) {
2921 PyErr_Format(PicklingError, "Fifth element of tuple"
2922 "returned by __reduce__ must be an iterator, not %s",
2923 Py_TYPE(dictitems)->tp_name);
2924 return -1;
2925 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002926
2927 /* Protocol 2 special case: if callable's name is __newobj__, use
2928 NEWOBJ. */
2929 if (use_newobj) {
Antoine Pitrou7eecffd2010-10-22 19:43:59 +00002930 static PyObject *newobj_str = NULL, *name_str = NULL;
2931 PyObject *name;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002932
2933 if (newobj_str == NULL) {
2934 newobj_str = PyUnicode_InternFromString("__newobj__");
Antoine Pitrou7eecffd2010-10-22 19:43:59 +00002935 name_str = PyUnicode_InternFromString("__name__");
2936 if (newobj_str == NULL || name_str == NULL)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00002937 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002938 }
2939
Antoine Pitrou7eecffd2010-10-22 19:43:59 +00002940 name = PyObject_GetAttr(callable, name_str);
2941 if (name == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002942 if (PyErr_ExceptionMatches(PyExc_AttributeError))
2943 PyErr_Clear();
2944 else
2945 return -1;
2946 use_newobj = 0;
2947 }
2948 else {
Antoine Pitrou7eecffd2010-10-22 19:43:59 +00002949 use_newobj = PyUnicode_Check(name) &&
2950 PyUnicode_Compare(name, newobj_str) == 0;
2951 Py_DECREF(name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002952 }
2953 }
2954 if (use_newobj) {
2955 PyObject *cls;
2956 PyObject *newargtup;
2957 PyObject *obj_class;
2958 int p;
2959
2960 /* Sanity checks. */
2961 if (Py_SIZE(argtup) < 1) {
2962 PyErr_SetString(PicklingError, "__newobj__ arglist is empty");
2963 return -1;
2964 }
2965
2966 cls = PyTuple_GET_ITEM(argtup, 0);
Antoine Pitrou7eecffd2010-10-22 19:43:59 +00002967 if (!PyType_Check(cls)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002968 PyErr_SetString(PicklingError, "args[0] from "
Antoine Pitrou7eecffd2010-10-22 19:43:59 +00002969 "__newobj__ args is not a type");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002970 return -1;
2971 }
2972
2973 if (obj != NULL) {
Antoine Pitrou7eecffd2010-10-22 19:43:59 +00002974 obj_class = get_class(obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00002975 p = obj_class != cls; /* true iff a problem */
2976 Py_DECREF(obj_class);
2977 if (p) {
2978 PyErr_SetString(PicklingError, "args[0] from "
2979 "__newobj__ args has the wrong class");
2980 return -1;
2981 }
2982 }
2983 /* XXX: These calls save() are prone to infinite recursion. Imagine
2984 what happen if the value returned by the __reduce__() method of
2985 some extension type contains another object of the same type. Ouch!
2986
2987 Here is a quick example, that I ran into, to illustrate what I
2988 mean:
2989
2990 >>> import pickle, copyreg
2991 >>> copyreg.dispatch_table.pop(complex)
2992 >>> pickle.dumps(1+2j)
2993 Traceback (most recent call last):
2994 ...
2995 RuntimeError: maximum recursion depth exceeded
2996
2997 Removing the complex class from copyreg.dispatch_table made the
2998 __reduce_ex__() method emit another complex object:
2999
3000 >>> (1+1j).__reduce_ex__(2)
3001 (<function __newobj__ at 0xb7b71c3c>,
3002 (<class 'complex'>, (1+1j)), None, None, None)
3003
3004 Thus when save() was called on newargstup (the 2nd item) recursion
3005 ensued. Of course, the bug was in the complex class which had a
3006 broken __getnewargs__() that emitted another complex object. But,
3007 the point, here, is it is quite easy to end up with a broken reduce
3008 function. */
3009
3010 /* Save the class and its __new__ arguments. */
3011 if (save(self, cls, 0) < 0)
3012 return -1;
3013
3014 newargtup = PyTuple_GetSlice(argtup, 1, Py_SIZE(argtup));
3015 if (newargtup == NULL)
3016 return -1;
3017
3018 p = save(self, newargtup, 0);
3019 Py_DECREF(newargtup);
3020 if (p < 0)
3021 return -1;
3022
3023 /* Add NEWOBJ opcode. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003024 if (_Pickler_Write(self, &newobj_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003025 return -1;
3026 }
3027 else { /* Not using NEWOBJ. */
3028 if (save(self, callable, 0) < 0 ||
3029 save(self, argtup, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003030 _Pickler_Write(self, &reduce_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003031 return -1;
3032 }
3033
3034 /* obj can be NULL when save_reduce() is used directly. A NULL obj means
3035 the caller do not want to memoize the object. Not particularly useful,
3036 but that is to mimic the behavior save_reduce() in pickle.py when
3037 obj is None. */
3038 if (obj && memo_put(self, obj) < 0)
3039 return -1;
3040
3041 if (listitems && batch_list(self, listitems) < 0)
3042 return -1;
3043
3044 if (dictitems && batch_dict(self, dictitems) < 0)
3045 return -1;
3046
3047 if (state) {
3048 if (save(self, state, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003049 _Pickler_Write(self, &build_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003050 return -1;
3051 }
3052
3053 return 0;
3054}
3055
3056static int
3057save(PicklerObject *self, PyObject *obj, int pers_save)
3058{
3059 PyTypeObject *type;
3060 PyObject *reduce_func = NULL;
3061 PyObject *reduce_value = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003062 int status = 0;
3063
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003064 if (Py_EnterRecursiveCall(" while pickling an object") < 0)
3065 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003066
3067 /* The extra pers_save argument is necessary to avoid calling save_pers()
3068 on its returned object. */
3069 if (!pers_save && self->pers_func) {
3070 /* save_pers() returns:
3071 -1 to signal an error;
3072 0 if it did nothing successfully;
3073 1 if a persistent id was saved.
3074 */
3075 if ((status = save_pers(self, obj, self->pers_func)) != 0)
3076 goto done;
3077 }
3078
3079 type = Py_TYPE(obj);
3080
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003081 /* The old cPickle had an optimization that used switch-case statement
3082 dispatching on the first letter of the type name. This has was removed
3083 since benchmarks shown that this optimization was actually slowing
3084 things down. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003085
3086 /* Atom types; these aren't memoized, so don't check the memo. */
3087
3088 if (obj == Py_None) {
3089 status = save_none(self, obj);
3090 goto done;
3091 }
3092 else if (obj == Py_False || obj == Py_True) {
3093 status = save_bool(self, obj);
3094 goto done;
3095 }
3096 else if (type == &PyLong_Type) {
3097 status = save_long(self, obj);
3098 goto done;
3099 }
3100 else if (type == &PyFloat_Type) {
3101 status = save_float(self, obj);
3102 goto done;
3103 }
3104
3105 /* Check the memo to see if it has the object. If so, generate
3106 a GET (or BINGET) opcode, instead of pickling the object
3107 once again. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003108 if (PyMemoTable_Get(self->memo, obj)) {
3109 if (memo_get(self, obj) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003110 goto error;
3111 goto done;
3112 }
3113
3114 if (type == &PyBytes_Type) {
3115 status = save_bytes(self, obj);
3116 goto done;
3117 }
3118 else if (type == &PyUnicode_Type) {
3119 status = save_unicode(self, obj);
3120 goto done;
3121 }
3122 else if (type == &PyDict_Type) {
3123 status = save_dict(self, obj);
3124 goto done;
3125 }
3126 else if (type == &PyList_Type) {
3127 status = save_list(self, obj);
3128 goto done;
3129 }
3130 else if (type == &PyTuple_Type) {
3131 status = save_tuple(self, obj);
3132 goto done;
3133 }
3134 else if (type == &PyType_Type) {
3135 status = save_global(self, obj, NULL);
3136 goto done;
3137 }
3138 else if (type == &PyFunction_Type) {
3139 status = save_global(self, obj, NULL);
3140 if (status < 0 && PyErr_ExceptionMatches(PickleError)) {
3141 /* fall back to reduce */
3142 PyErr_Clear();
3143 }
3144 else {
3145 goto done;
3146 }
3147 }
3148 else if (type == &PyCFunction_Type) {
3149 status = save_global(self, obj, NULL);
3150 goto done;
3151 }
3152 else if (PyType_IsSubtype(type, &PyType_Type)) {
3153 status = save_global(self, obj, NULL);
3154 goto done;
3155 }
3156
3157 /* XXX: This part needs some unit tests. */
3158
3159 /* Get a reduction callable, and call it. This may come from
3160 * copyreg.dispatch_table, the object's __reduce_ex__ method,
3161 * or the object's __reduce__ method.
3162 */
3163 reduce_func = PyDict_GetItem(dispatch_table, (PyObject *)type);
3164 if (reduce_func != NULL) {
3165 /* Here, the reference count of the reduce_func object returned by
3166 PyDict_GetItem needs to be increased to be consistent with the one
3167 returned by PyObject_GetAttr. This is allow us to blindly DECREF
3168 reduce_func at the end of the save() routine.
3169 */
3170 Py_INCREF(reduce_func);
3171 Py_INCREF(obj);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003172 reduce_value = _Pickler_FastCall(self, reduce_func, obj);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003173 }
3174 else {
3175 static PyObject *reduce_str = NULL;
3176 static PyObject *reduce_ex_str = NULL;
3177
3178 /* Cache the name of the reduce methods. */
3179 if (reduce_str == NULL) {
3180 reduce_str = PyUnicode_InternFromString("__reduce__");
3181 if (reduce_str == NULL)
3182 goto error;
3183 reduce_ex_str = PyUnicode_InternFromString("__reduce_ex__");
3184 if (reduce_ex_str == NULL)
3185 goto error;
3186 }
3187
3188 /* XXX: If the __reduce__ method is defined, __reduce_ex__ is
3189 automatically defined as __reduce__. While this is convenient, this
3190 make it impossible to know which method was actually called. Of
3191 course, this is not a big deal. But still, it would be nice to let
3192 the user know which method was called when something go
3193 wrong. Incidentally, this means if __reduce_ex__ is not defined, we
3194 don't actually have to check for a __reduce__ method. */
3195
3196 /* Check for a __reduce_ex__ method. */
3197 reduce_func = PyObject_GetAttr(obj, reduce_ex_str);
3198 if (reduce_func != NULL) {
3199 PyObject *proto;
3200 proto = PyLong_FromLong(self->proto);
3201 if (proto != NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003202 reduce_value = _Pickler_FastCall(self, reduce_func, proto);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003203 }
3204 }
3205 else {
3206 if (PyErr_ExceptionMatches(PyExc_AttributeError))
3207 PyErr_Clear();
3208 else
3209 goto error;
3210 /* Check for a __reduce__ method. */
3211 reduce_func = PyObject_GetAttr(obj, reduce_str);
3212 if (reduce_func != NULL) {
3213 reduce_value = PyObject_Call(reduce_func, empty_tuple, NULL);
3214 }
3215 else {
3216 PyErr_Format(PicklingError, "can't pickle '%.200s' object: %R",
3217 type->tp_name, obj);
3218 goto error;
3219 }
3220 }
3221 }
3222
3223 if (reduce_value == NULL)
3224 goto error;
3225
3226 if (PyUnicode_Check(reduce_value)) {
3227 status = save_global(self, obj, reduce_value);
3228 goto done;
3229 }
3230
3231 if (!PyTuple_Check(reduce_value)) {
3232 PyErr_SetString(PicklingError,
3233 "__reduce__ must return a string or tuple");
3234 goto error;
3235 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003236
3237 status = save_reduce(self, reduce_value, obj);
3238
3239 if (0) {
3240 error:
3241 status = -1;
3242 }
3243 done:
Alexandre Vassalottidff18342008-07-13 18:48:30 +00003244 Py_LeaveRecursiveCall();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003245 Py_XDECREF(reduce_func);
3246 Py_XDECREF(reduce_value);
3247
3248 return status;
3249}
3250
3251static int
3252dump(PicklerObject *self, PyObject *obj)
3253{
3254 const char stop_op = STOP;
3255
3256 if (self->proto >= 2) {
3257 char header[2];
3258
3259 header[0] = PROTO;
3260 assert(self->proto >= 0 && self->proto < 256);
3261 header[1] = (unsigned char)self->proto;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003262 if (_Pickler_Write(self, header, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003263 return -1;
3264 }
3265
3266 if (save(self, obj, 0) < 0 ||
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003267 _Pickler_Write(self, &stop_op, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003268 return -1;
3269
3270 return 0;
3271}
3272
3273PyDoc_STRVAR(Pickler_clear_memo_doc,
3274"clear_memo() -> None. Clears the pickler's \"memo\"."
3275"\n"
3276"The memo is the data structure that remembers which objects the\n"
3277"pickler has already seen, so that shared or recursive objects are\n"
3278"pickled by reference and not by value. This method is useful when\n"
3279"re-using picklers.");
3280
3281static PyObject *
3282Pickler_clear_memo(PicklerObject *self)
3283{
3284 if (self->memo)
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003285 PyMemoTable_Clear(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003286
3287 Py_RETURN_NONE;
3288}
3289
3290PyDoc_STRVAR(Pickler_dump_doc,
3291"dump(obj) -> None. Write a pickled representation of obj to the open file.");
3292
3293static PyObject *
3294Pickler_dump(PicklerObject *self, PyObject *args)
3295{
3296 PyObject *obj;
3297
Amaury Forgeot d'Arc87eee632008-10-17 20:15:53 +00003298 /* Check whether the Pickler was initialized correctly (issue3664).
3299 Developers often forget to call __init__() in their subclasses, which
3300 would trigger a segfault without this check. */
3301 if (self->write == NULL) {
3302 PyErr_Format(PicklingError,
3303 "Pickler.__init__() was not called by %s.__init__()",
3304 Py_TYPE(self)->tp_name);
3305 return NULL;
3306 }
3307
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003308 if (!PyArg_ParseTuple(args, "O:dump", &obj))
3309 return NULL;
3310
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003311 if (_Pickler_ClearBuffer(self) < 0)
3312 return NULL;
3313
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003314 if (dump(self, obj) < 0)
3315 return NULL;
3316
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003317 if (_Pickler_FlushToFile(self) < 0)
3318 return NULL;
3319
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003320 Py_RETURN_NONE;
3321}
3322
3323static struct PyMethodDef Pickler_methods[] = {
3324 {"dump", (PyCFunction)Pickler_dump, METH_VARARGS,
3325 Pickler_dump_doc},
3326 {"clear_memo", (PyCFunction)Pickler_clear_memo, METH_NOARGS,
3327 Pickler_clear_memo_doc},
3328 {NULL, NULL} /* sentinel */
3329};
3330
3331static void
3332Pickler_dealloc(PicklerObject *self)
3333{
3334 PyObject_GC_UnTrack(self);
3335
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003336 Py_XDECREF(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003337 Py_XDECREF(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003338 Py_XDECREF(self->pers_func);
3339 Py_XDECREF(self->arg);
3340 Py_XDECREF(self->fast_memo);
3341
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003342 PyMemoTable_Del(self->memo);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003343
3344 Py_TYPE(self)->tp_free((PyObject *)self);
3345}
3346
3347static int
3348Pickler_traverse(PicklerObject *self, visitproc visit, void *arg)
3349{
3350 Py_VISIT(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003351 Py_VISIT(self->pers_func);
3352 Py_VISIT(self->arg);
3353 Py_VISIT(self->fast_memo);
3354 return 0;
3355}
3356
3357static int
3358Pickler_clear(PicklerObject *self)
3359{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003360 Py_CLEAR(self->output_buffer);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003361 Py_CLEAR(self->write);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003362 Py_CLEAR(self->pers_func);
3363 Py_CLEAR(self->arg);
3364 Py_CLEAR(self->fast_memo);
3365
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003366 if (self->memo != NULL) {
3367 PyMemoTable *memo = self->memo;
3368 self->memo = NULL;
3369 PyMemoTable_Del(memo);
3370 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003371 return 0;
3372}
3373
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003374
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003375PyDoc_STRVAR(Pickler_doc,
3376"Pickler(file, protocol=None)"
3377"\n"
3378"This takes a binary file for writing a pickle data stream.\n"
3379"\n"
3380"The optional protocol argument tells the pickler to use the\n"
3381"given protocol; supported protocols are 0, 1, 2, 3. The default\n"
3382"protocol is 3; a backward-incompatible protocol designed for\n"
3383"Python 3.0.\n"
3384"\n"
3385"Specifying a negative protocol version selects the highest\n"
3386"protocol version supported. The higher the protocol used, the\n"
3387"more recent the version of Python needed to read the pickle\n"
3388"produced.\n"
3389"\n"
3390"The file argument must have a write() method that accepts a single\n"
3391"bytes argument. It can thus be a file object opened for binary\n"
3392"writing, a io.BytesIO instance, or any other custom object that\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003393"meets this interface.\n"
3394"\n"
3395"If fix_imports is True and protocol is less than 3, pickle will try to\n"
3396"map the new Python 3.x names to the old module names used in Python\n"
3397"2.x, so that the pickle data stream is readable with Python 2.x.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003398
3399static int
3400Pickler_init(PicklerObject *self, PyObject *args, PyObject *kwds)
3401{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003402 static char *kwlist[] = {"file", "protocol", "fix_imports", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003403 PyObject *file;
3404 PyObject *proto_obj = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003405 PyObject *fix_imports = Py_True;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003406
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003407 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:Pickler",
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003408 kwlist, &file, &proto_obj, &fix_imports))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003409 return -1;
3410
3411 /* In case of multiple __init__() calls, clear previous content. */
3412 if (self->write != NULL)
3413 (void)Pickler_clear(self);
3414
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003415 if (_Pickler_SetProtocol(self, proto_obj, fix_imports) < 0)
3416 return -1;
3417
3418 if (_Pickler_SetOutputStream(self, file) < 0)
3419 return -1;
3420
3421 /* memo and output_buffer may have already been created in _Pickler_New */
3422 if (self->memo == NULL) {
3423 self->memo = PyMemoTable_New();
3424 if (self->memo == NULL)
3425 return -1;
3426 }
3427 self->output_len = 0;
3428 if (self->output_buffer == NULL) {
3429 self->max_output_len = WRITE_BUF_SIZE;
3430 self->output_buffer = PyBytes_FromStringAndSize(NULL,
3431 self->max_output_len);
3432 if (self->output_buffer == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003433 return -1;
3434 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003435
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00003436 self->arg = NULL;
3437 self->fast = 0;
3438 self->fast_nesting = 0;
3439 self->fast_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003440 self->pers_func = NULL;
3441 if (PyObject_HasAttrString((PyObject *)self, "persistent_id")) {
3442 self->pers_func = PyObject_GetAttrString((PyObject *)self,
3443 "persistent_id");
3444 if (self->pers_func == NULL)
3445 return -1;
3446 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003447 return 0;
3448}
3449
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003450/* Define a proxy object for the Pickler's internal memo object. This is to
3451 * avoid breaking code like:
3452 * pickler.memo.clear()
3453 * and
3454 * pickler.memo = saved_memo
3455 * Is this a good idea? Not really, but we don't want to break code that uses
3456 * it. Note that we don't implement the entire mapping API here. This is
3457 * intentional, as these should be treated as black-box implementation details.
3458 */
3459
3460typedef struct {
3461 PyObject_HEAD
3462 PicklerObject *pickler; /* Pickler whose memo table we're proxying. */
3463} PicklerMemoProxyObject;
3464
3465PyDoc_STRVAR(pmp_clear_doc,
3466"memo.clear() -> None. Remove all items from memo.");
3467
3468static PyObject *
3469pmp_clear(PicklerMemoProxyObject *self)
3470{
3471 if (self->pickler->memo)
3472 PyMemoTable_Clear(self->pickler->memo);
3473 Py_RETURN_NONE;
3474}
3475
3476PyDoc_STRVAR(pmp_copy_doc,
3477"memo.copy() -> new_memo. Copy the memo to a new object.");
3478
3479static PyObject *
3480pmp_copy(PicklerMemoProxyObject *self)
3481{
3482 Py_ssize_t i;
3483 PyMemoTable *memo;
3484 PyObject *new_memo = PyDict_New();
3485 if (new_memo == NULL)
3486 return NULL;
3487
3488 memo = self->pickler->memo;
3489 for (i = 0; i < memo->mt_allocated; ++i) {
3490 PyMemoEntry entry = memo->mt_table[i];
3491 if (entry.me_key != NULL) {
3492 int status;
3493 PyObject *key, *value;
3494
3495 key = PyLong_FromVoidPtr(entry.me_key);
3496 value = Py_BuildValue("lO", entry.me_value, entry.me_key);
3497
3498 if (key == NULL || value == NULL) {
3499 Py_XDECREF(key);
3500 Py_XDECREF(value);
3501 goto error;
3502 }
3503 status = PyDict_SetItem(new_memo, key, value);
3504 Py_DECREF(key);
3505 Py_DECREF(value);
3506 if (status < 0)
3507 goto error;
3508 }
3509 }
3510 return new_memo;
3511
3512 error:
3513 Py_XDECREF(new_memo);
3514 return NULL;
3515}
3516
3517PyDoc_STRVAR(pmp_reduce_doc,
3518"memo.__reduce__(). Pickling support.");
3519
3520static PyObject *
3521pmp_reduce(PicklerMemoProxyObject *self, PyObject *args)
3522{
3523 PyObject *reduce_value, *dict_args;
3524 PyObject *contents = pmp_copy(self);
3525 if (contents == NULL)
3526 return NULL;
3527
3528 reduce_value = PyTuple_New(2);
3529 if (reduce_value == NULL) {
3530 Py_DECREF(contents);
3531 return NULL;
3532 }
3533 dict_args = PyTuple_New(1);
3534 if (dict_args == NULL) {
3535 Py_DECREF(contents);
3536 Py_DECREF(reduce_value);
3537 return NULL;
3538 }
3539 PyTuple_SET_ITEM(dict_args, 0, contents);
3540 Py_INCREF((PyObject *)&PyDict_Type);
3541 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
3542 PyTuple_SET_ITEM(reduce_value, 1, dict_args);
3543 return reduce_value;
3544}
3545
3546static PyMethodDef picklerproxy_methods[] = {
3547 {"clear", (PyCFunction)pmp_clear, METH_NOARGS, pmp_clear_doc},
3548 {"copy", (PyCFunction)pmp_copy, METH_NOARGS, pmp_copy_doc},
3549 {"__reduce__", (PyCFunction)pmp_reduce, METH_VARARGS, pmp_reduce_doc},
3550 {NULL, NULL} /* sentinel */
3551};
3552
3553static void
3554PicklerMemoProxy_dealloc(PicklerMemoProxyObject *self)
3555{
3556 PyObject_GC_UnTrack(self);
3557 Py_XDECREF(self->pickler);
3558 PyObject_GC_Del((PyObject *)self);
3559}
3560
3561static int
3562PicklerMemoProxy_traverse(PicklerMemoProxyObject *self,
3563 visitproc visit, void *arg)
3564{
3565 Py_VISIT(self->pickler);
3566 return 0;
3567}
3568
3569static int
3570PicklerMemoProxy_clear(PicklerMemoProxyObject *self)
3571{
3572 Py_CLEAR(self->pickler);
3573 return 0;
3574}
3575
3576static PyTypeObject PicklerMemoProxyType = {
3577 PyVarObject_HEAD_INIT(NULL, 0)
3578 "_pickle.PicklerMemoProxy", /*tp_name*/
3579 sizeof(PicklerMemoProxyObject), /*tp_basicsize*/
3580 0,
3581 (destructor)PicklerMemoProxy_dealloc, /* tp_dealloc */
3582 0, /* tp_print */
3583 0, /* tp_getattr */
3584 0, /* tp_setattr */
3585 0, /* tp_compare */
3586 0, /* tp_repr */
3587 0, /* tp_as_number */
3588 0, /* tp_as_sequence */
3589 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00003590 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003591 0, /* tp_call */
3592 0, /* tp_str */
3593 PyObject_GenericGetAttr, /* tp_getattro */
3594 PyObject_GenericSetAttr, /* tp_setattro */
3595 0, /* tp_as_buffer */
3596 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3597 0, /* tp_doc */
3598 (traverseproc)PicklerMemoProxy_traverse, /* tp_traverse */
3599 (inquiry)PicklerMemoProxy_clear, /* tp_clear */
3600 0, /* tp_richcompare */
3601 0, /* tp_weaklistoffset */
3602 0, /* tp_iter */
3603 0, /* tp_iternext */
3604 picklerproxy_methods, /* tp_methods */
3605};
3606
3607static PyObject *
3608PicklerMemoProxy_New(PicklerObject *pickler)
3609{
3610 PicklerMemoProxyObject *self;
3611
3612 self = PyObject_GC_New(PicklerMemoProxyObject, &PicklerMemoProxyType);
3613 if (self == NULL)
3614 return NULL;
3615 Py_INCREF(pickler);
3616 self->pickler = pickler;
3617 PyObject_GC_Track(self);
3618 return (PyObject *)self;
3619}
3620
3621/*****************************************************************************/
3622
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003623static PyObject *
3624Pickler_get_memo(PicklerObject *self)
3625{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003626 return PicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003627}
3628
3629static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003630Pickler_set_memo(PicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003631{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003632 PyMemoTable *new_memo = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003633
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003634 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003635 PyErr_SetString(PyExc_TypeError,
3636 "attribute deletion is not supported");
3637 return -1;
3638 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003639
3640 if (Py_TYPE(obj) == &PicklerMemoProxyType) {
3641 PicklerObject *pickler =
3642 ((PicklerMemoProxyObject *)obj)->pickler;
3643
3644 new_memo = PyMemoTable_Copy(pickler->memo);
3645 if (new_memo == NULL)
3646 return -1;
3647 }
3648 else if (PyDict_Check(obj)) {
3649 Py_ssize_t i = 0;
3650 PyObject *key, *value;
3651
3652 new_memo = PyMemoTable_New();
3653 if (new_memo == NULL)
3654 return -1;
3655
3656 while (PyDict_Next(obj, &i, &key, &value)) {
3657 long memo_id;
3658 PyObject *memo_obj;
3659
3660 if (!PyTuple_Check(value) || Py_SIZE(value) != 2) {
3661 PyErr_SetString(PyExc_TypeError,
3662 "'memo' values must be 2-item tuples");
3663 goto error;
3664 }
3665 memo_id = PyLong_AsLong(PyTuple_GET_ITEM(value, 0));
3666 if (memo_id == -1 && PyErr_Occurred())
3667 goto error;
3668 memo_obj = PyTuple_GET_ITEM(value, 1);
3669 if (PyMemoTable_Set(new_memo, memo_obj, memo_id) < 0)
3670 goto error;
3671 }
3672 }
3673 else {
3674 PyErr_Format(PyExc_TypeError,
3675 "'memo' attribute must be an PicklerMemoProxy object"
3676 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003677 return -1;
3678 }
3679
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003680 PyMemoTable_Del(self->memo);
3681 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003682
3683 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003684
3685 error:
3686 if (new_memo)
3687 PyMemoTable_Del(new_memo);
3688 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003689}
3690
3691static PyObject *
3692Pickler_get_persid(PicklerObject *self)
3693{
3694 if (self->pers_func == NULL)
3695 PyErr_SetString(PyExc_AttributeError, "persistent_id");
3696 else
3697 Py_INCREF(self->pers_func);
3698 return self->pers_func;
3699}
3700
3701static int
3702Pickler_set_persid(PicklerObject *self, PyObject *value)
3703{
3704 PyObject *tmp;
3705
3706 if (value == NULL) {
3707 PyErr_SetString(PyExc_TypeError,
3708 "attribute deletion is not supported");
3709 return -1;
3710 }
3711 if (!PyCallable_Check(value)) {
3712 PyErr_SetString(PyExc_TypeError,
3713 "persistent_id must be a callable taking one argument");
3714 return -1;
3715 }
3716
3717 tmp = self->pers_func;
3718 Py_INCREF(value);
3719 self->pers_func = value;
3720 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
3721
3722 return 0;
3723}
3724
3725static PyMemberDef Pickler_members[] = {
3726 {"bin", T_INT, offsetof(PicklerObject, bin)},
3727 {"fast", T_INT, offsetof(PicklerObject, fast)},
3728 {NULL}
3729};
3730
3731static PyGetSetDef Pickler_getsets[] = {
3732 {"memo", (getter)Pickler_get_memo,
3733 (setter)Pickler_set_memo},
3734 {"persistent_id", (getter)Pickler_get_persid,
3735 (setter)Pickler_set_persid},
3736 {NULL}
3737};
3738
3739static PyTypeObject Pickler_Type = {
3740 PyVarObject_HEAD_INIT(NULL, 0)
3741 "_pickle.Pickler" , /*tp_name*/
3742 sizeof(PicklerObject), /*tp_basicsize*/
3743 0, /*tp_itemsize*/
3744 (destructor)Pickler_dealloc, /*tp_dealloc*/
3745 0, /*tp_print*/
3746 0, /*tp_getattr*/
3747 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00003748 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003749 0, /*tp_repr*/
3750 0, /*tp_as_number*/
3751 0, /*tp_as_sequence*/
3752 0, /*tp_as_mapping*/
3753 0, /*tp_hash*/
3754 0, /*tp_call*/
3755 0, /*tp_str*/
3756 0, /*tp_getattro*/
3757 0, /*tp_setattro*/
3758 0, /*tp_as_buffer*/
3759 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
3760 Pickler_doc, /*tp_doc*/
3761 (traverseproc)Pickler_traverse, /*tp_traverse*/
3762 (inquiry)Pickler_clear, /*tp_clear*/
3763 0, /*tp_richcompare*/
3764 0, /*tp_weaklistoffset*/
3765 0, /*tp_iter*/
3766 0, /*tp_iternext*/
3767 Pickler_methods, /*tp_methods*/
3768 Pickler_members, /*tp_members*/
3769 Pickler_getsets, /*tp_getset*/
3770 0, /*tp_base*/
3771 0, /*tp_dict*/
3772 0, /*tp_descr_get*/
3773 0, /*tp_descr_set*/
3774 0, /*tp_dictoffset*/
3775 (initproc)Pickler_init, /*tp_init*/
3776 PyType_GenericAlloc, /*tp_alloc*/
3777 PyType_GenericNew, /*tp_new*/
3778 PyObject_GC_Del, /*tp_free*/
3779 0, /*tp_is_gc*/
3780};
3781
3782/* Temporary helper for calling self.find_class().
3783
3784 XXX: It would be nice to able to avoid Python function call overhead, by
3785 using directly the C version of find_class(), when find_class() is not
3786 overridden by a subclass. Although, this could become rather hackish. A
3787 simpler optimization would be to call the C function when self is not a
3788 subclass instance. */
3789static PyObject *
3790find_class(UnpicklerObject *self, PyObject *module_name, PyObject *global_name)
3791{
3792 return PyObject_CallMethod((PyObject *)self, "find_class", "OO",
3793 module_name, global_name);
3794}
3795
3796static int
3797marker(UnpicklerObject *self)
3798{
3799 if (self->num_marks < 1) {
3800 PyErr_SetString(UnpicklingError, "could not find MARK");
3801 return -1;
3802 }
3803
3804 return self->marks[--self->num_marks];
3805}
3806
3807static int
3808load_none(UnpicklerObject *self)
3809{
3810 PDATA_APPEND(self->stack, Py_None, -1);
3811 return 0;
3812}
3813
3814static int
3815bad_readline(void)
3816{
3817 PyErr_SetString(UnpicklingError, "pickle data was truncated");
3818 return -1;
3819}
3820
3821static int
3822load_int(UnpicklerObject *self)
3823{
3824 PyObject *value;
3825 char *endptr, *s;
3826 Py_ssize_t len;
3827 long x;
3828
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003829 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003830 return -1;
3831 if (len < 2)
3832 return bad_readline();
3833
3834 errno = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003835 /* XXX: Should the base argument of strtol() be explicitly set to 10?
3836 XXX(avassalotti): Should this uses PyOS_strtol()? */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003837 x = strtol(s, &endptr, 0);
3838
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003839 if (errno || (*endptr != '\n' && *endptr != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003840 /* Hm, maybe we've got something long. Let's try reading
3841 * it as a Python long object. */
3842 errno = 0;
3843 /* XXX: Same thing about the base here. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003844 value = PyLong_FromString(s, NULL, 0);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003845 if (value == NULL) {
3846 PyErr_SetString(PyExc_ValueError,
3847 "could not convert string to int");
3848 return -1;
3849 }
3850 }
3851 else {
3852 if (len == 3 && (x == 0 || x == 1)) {
3853 if ((value = PyBool_FromLong(x)) == NULL)
3854 return -1;
3855 }
3856 else {
3857 if ((value = PyLong_FromLong(x)) == NULL)
3858 return -1;
3859 }
3860 }
3861
3862 PDATA_PUSH(self->stack, value, -1);
3863 return 0;
3864}
3865
3866static int
3867load_bool(UnpicklerObject *self, PyObject *boolean)
3868{
3869 assert(boolean == Py_True || boolean == Py_False);
3870 PDATA_APPEND(self->stack, boolean, -1);
3871 return 0;
3872}
3873
3874/* s contains x bytes of a little-endian integer. Return its value as a
3875 * C int. Obscure: when x is 1 or 2, this is an unsigned little-endian
3876 * int, but when x is 4 it's a signed one. This is an historical source
3877 * of x-platform bugs.
3878 */
3879static long
3880calc_binint(char *bytes, int size)
3881{
3882 unsigned char *s = (unsigned char *)bytes;
3883 int i = size;
3884 long x = 0;
3885
3886 for (i = 0; i < size; i++) {
3887 x |= (long)s[i] << (i * 8);
3888 }
3889
3890 /* Unlike BININT1 and BININT2, BININT (more accurately BININT4)
3891 * is signed, so on a box with longs bigger than 4 bytes we need
3892 * to extend a BININT's sign bit to the full width.
3893 */
3894 if (SIZEOF_LONG > 4 && size == 4) {
3895 x |= -(x & (1L << 31));
3896 }
3897
3898 return x;
3899}
3900
3901static int
3902load_binintx(UnpicklerObject *self, char *s, int size)
3903{
3904 PyObject *value;
3905 long x;
3906
3907 x = calc_binint(s, size);
3908
3909 if ((value = PyLong_FromLong(x)) == NULL)
3910 return -1;
3911
3912 PDATA_PUSH(self->stack, value, -1);
3913 return 0;
3914}
3915
3916static int
3917load_binint(UnpicklerObject *self)
3918{
3919 char *s;
3920
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003921 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003922 return -1;
3923
3924 return load_binintx(self, s, 4);
3925}
3926
3927static int
3928load_binint1(UnpicklerObject *self)
3929{
3930 char *s;
3931
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003932 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003933 return -1;
3934
3935 return load_binintx(self, s, 1);
3936}
3937
3938static int
3939load_binint2(UnpicklerObject *self)
3940{
3941 char *s;
3942
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003943 if (_Unpickler_Read(self, &s, 2) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003944 return -1;
3945
3946 return load_binintx(self, s, 2);
3947}
3948
3949static int
3950load_long(UnpicklerObject *self)
3951{
3952 PyObject *value;
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00003953 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003954 Py_ssize_t len;
3955
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003956 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003957 return -1;
3958 if (len < 2)
3959 return bad_readline();
3960
Mark Dickinson8dd05142009-01-20 20:43:58 +00003961 /* s[len-2] will usually be 'L' (and s[len-1] is '\n'); we need to remove
3962 the 'L' before calling PyLong_FromString. In order to maintain
3963 compatibility with Python 3.0.0, we don't actually *require*
3964 the 'L' to be present. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003965 if (s[len-2] == 'L')
Alexandre Vassalotti446f7ff2009-01-23 04:43:46 +00003966 s[len-2] = '\0';
Alexandre Vassalottie4bccb72009-01-24 01:47:57 +00003967 /* XXX: Should the base argument explicitly set to 10? */
3968 value = PyLong_FromString(s, NULL, 0);
Mark Dickinson8dd05142009-01-20 20:43:58 +00003969 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003970 return -1;
3971
3972 PDATA_PUSH(self->stack, value, -1);
3973 return 0;
3974}
3975
3976/* 'size' bytes contain the # of bytes of little-endian 256's-complement
3977 * data following.
3978 */
3979static int
3980load_counted_long(UnpicklerObject *self, int size)
3981{
3982 PyObject *value;
3983 char *nbytes;
3984 char *pdata;
3985
3986 assert(size == 1 || size == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00003987 if (_Unpickler_Read(self, &nbytes, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00003988 return -1;
3989
3990 size = calc_binint(nbytes, size);
3991 if (size < 0) {
3992 /* Corrupt or hostile pickle -- we never write one like this */
3993 PyErr_SetString(UnpicklingError,
3994 "LONG pickle has negative byte count");
3995 return -1;
3996 }
3997
3998 if (size == 0)
3999 value = PyLong_FromLong(0L);
4000 else {
4001 /* Read the raw little-endian bytes and convert. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004002 if (_Unpickler_Read(self, &pdata, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004003 return -1;
4004 value = _PyLong_FromByteArray((unsigned char *)pdata, (size_t)size,
4005 1 /* little endian */ , 1 /* signed */ );
4006 }
4007 if (value == NULL)
4008 return -1;
4009 PDATA_PUSH(self->stack, value, -1);
4010 return 0;
4011}
4012
4013static int
4014load_float(UnpicklerObject *self)
4015{
4016 PyObject *value;
4017 char *endptr, *s;
4018 Py_ssize_t len;
4019 double d;
4020
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004021 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004022 return -1;
4023 if (len < 2)
4024 return bad_readline();
4025
4026 errno = 0;
Mark Dickinson725bfd82009-05-03 20:33:40 +00004027 d = PyOS_string_to_double(s, &endptr, PyExc_OverflowError);
4028 if (d == -1.0 && PyErr_Occurred())
4029 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004030 if ((endptr[0] != '\n') && (endptr[0] != '\0')) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004031 PyErr_SetString(PyExc_ValueError, "could not convert string to float");
4032 return -1;
4033 }
Mark Dickinson725bfd82009-05-03 20:33:40 +00004034 value = PyFloat_FromDouble(d);
4035 if (value == NULL)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004036 return -1;
4037
4038 PDATA_PUSH(self->stack, value, -1);
4039 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004040}
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004041
4042static int
4043load_binfloat(UnpicklerObject *self)
4044{
4045 PyObject *value;
4046 double x;
4047 char *s;
4048
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004049 if (_Unpickler_Read(self, &s, 8) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004050 return -1;
4051
4052 x = _PyFloat_Unpack8((unsigned char *)s, 0);
4053 if (x == -1.0 && PyErr_Occurred())
4054 return -1;
4055
4056 if ((value = PyFloat_FromDouble(x)) == NULL)
4057 return -1;
4058
4059 PDATA_PUSH(self->stack, value, -1);
4060 return 0;
4061}
4062
4063static int
4064load_string(UnpicklerObject *self)
4065{
4066 PyObject *bytes;
4067 PyObject *str = NULL;
4068 Py_ssize_t len;
4069 char *s, *p;
4070
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004071 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004072 return -1;
4073 if (len < 3)
4074 return bad_readline();
4075 if ((s = strdup(s)) == NULL) {
4076 PyErr_NoMemory();
4077 return -1;
4078 }
4079
4080 /* Strip outermost quotes */
4081 while (s[len - 1] <= ' ')
4082 len--;
4083 if (s[0] == '"' && s[len - 1] == '"') {
4084 s[len - 1] = '\0';
4085 p = s + 1;
4086 len -= 2;
4087 }
4088 else if (s[0] == '\'' && s[len - 1] == '\'') {
4089 s[len - 1] = '\0';
4090 p = s + 1;
4091 len -= 2;
4092 }
4093 else {
4094 free(s);
4095 PyErr_SetString(PyExc_ValueError, "insecure string pickle");
4096 return -1;
4097 }
4098
4099 /* Use the PyBytes API to decode the string, since that is what is used
4100 to encode, and then coerce the result to Unicode. */
4101 bytes = PyBytes_DecodeEscape(p, len, NULL, 0, NULL);
4102 free(s);
4103 if (bytes == NULL)
4104 return -1;
4105 str = PyUnicode_FromEncodedObject(bytes, self->encoding, self->errors);
4106 Py_DECREF(bytes);
4107 if (str == NULL)
4108 return -1;
4109
4110 PDATA_PUSH(self->stack, str, -1);
4111 return 0;
4112}
4113
4114static int
4115load_binbytes(UnpicklerObject *self)
4116{
4117 PyObject *bytes;
4118 long x;
4119 char *s;
4120
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004121 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004122 return -1;
4123
4124 x = calc_binint(s, 4);
4125 if (x < 0) {
4126 PyErr_SetString(UnpicklingError,
4127 "BINBYTES pickle has negative byte count");
4128 return -1;
4129 }
4130
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004131 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004132 return -1;
4133 bytes = PyBytes_FromStringAndSize(s, x);
4134 if (bytes == NULL)
4135 return -1;
4136
4137 PDATA_PUSH(self->stack, bytes, -1);
4138 return 0;
4139}
4140
4141static int
4142load_short_binbytes(UnpicklerObject *self)
4143{
4144 PyObject *bytes;
4145 unsigned char x;
4146 char *s;
4147
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004148 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004149 return -1;
4150
4151 x = (unsigned char)s[0];
4152
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004153 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004154 return -1;
4155
4156 bytes = PyBytes_FromStringAndSize(s, x);
4157 if (bytes == NULL)
4158 return -1;
4159
4160 PDATA_PUSH(self->stack, bytes, -1);
4161 return 0;
4162}
4163
4164static int
4165load_binstring(UnpicklerObject *self)
4166{
4167 PyObject *str;
4168 long x;
4169 char *s;
4170
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004171 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004172 return -1;
4173
4174 x = calc_binint(s, 4);
4175 if (x < 0) {
4176 PyErr_SetString(UnpicklingError,
4177 "BINSTRING pickle has negative byte count");
4178 return -1;
4179 }
4180
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004181 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004182 return -1;
4183
4184 /* Convert Python 2.x strings to unicode. */
4185 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4186 if (str == NULL)
4187 return -1;
4188
4189 PDATA_PUSH(self->stack, str, -1);
4190 return 0;
4191}
4192
4193static int
4194load_short_binstring(UnpicklerObject *self)
4195{
4196 PyObject *str;
4197 unsigned char x;
4198 char *s;
4199
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004200 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004201 return -1;
4202
4203 x = (unsigned char)s[0];
4204
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004205 if (_Unpickler_Read(self, &s, x) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004206 return -1;
4207
4208 /* Convert Python 2.x strings to unicode. */
4209 str = PyUnicode_Decode(s, x, self->encoding, self->errors);
4210 if (str == NULL)
4211 return -1;
4212
4213 PDATA_PUSH(self->stack, str, -1);
4214 return 0;
4215}
4216
4217static int
4218load_unicode(UnpicklerObject *self)
4219{
4220 PyObject *str;
4221 Py_ssize_t len;
4222 char *s;
4223
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004224 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004225 return -1;
4226 if (len < 1)
4227 return bad_readline();
4228
4229 str = PyUnicode_DecodeRawUnicodeEscape(s, len - 1, NULL);
4230 if (str == NULL)
4231 return -1;
4232
4233 PDATA_PUSH(self->stack, str, -1);
4234 return 0;
4235}
4236
4237static int
4238load_binunicode(UnpicklerObject *self)
4239{
4240 PyObject *str;
4241 long size;
4242 char *s;
4243
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004244 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004245 return -1;
4246
4247 size = calc_binint(s, 4);
4248 if (size < 0) {
4249 PyErr_SetString(UnpicklingError,
4250 "BINUNICODE pickle has negative byte count");
4251 return -1;
4252 }
4253
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004254 if (_Unpickler_Read(self, &s, size) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004255 return -1;
4256
Victor Stinner485fb562010-04-13 11:07:24 +00004257 str = PyUnicode_DecodeUTF8(s, size, "surrogatepass");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004258 if (str == NULL)
4259 return -1;
4260
4261 PDATA_PUSH(self->stack, str, -1);
4262 return 0;
4263}
4264
4265static int
4266load_tuple(UnpicklerObject *self)
4267{
4268 PyObject *tuple;
4269 int i;
4270
4271 if ((i = marker(self)) < 0)
4272 return -1;
4273
4274 tuple = Pdata_poptuple(self->stack, i);
4275 if (tuple == NULL)
4276 return -1;
4277 PDATA_PUSH(self->stack, tuple, -1);
4278 return 0;
4279}
4280
4281static int
4282load_counted_tuple(UnpicklerObject *self, int len)
4283{
4284 PyObject *tuple;
4285
4286 tuple = PyTuple_New(len);
4287 if (tuple == NULL)
4288 return -1;
4289
4290 while (--len >= 0) {
4291 PyObject *item;
4292
4293 PDATA_POP(self->stack, item);
4294 if (item == NULL)
4295 return -1;
4296 PyTuple_SET_ITEM(tuple, len, item);
4297 }
4298 PDATA_PUSH(self->stack, tuple, -1);
4299 return 0;
4300}
4301
4302static int
4303load_empty_list(UnpicklerObject *self)
4304{
4305 PyObject *list;
4306
4307 if ((list = PyList_New(0)) == NULL)
4308 return -1;
4309 PDATA_PUSH(self->stack, list, -1);
4310 return 0;
4311}
4312
4313static int
4314load_empty_dict(UnpicklerObject *self)
4315{
4316 PyObject *dict;
4317
4318 if ((dict = PyDict_New()) == NULL)
4319 return -1;
4320 PDATA_PUSH(self->stack, dict, -1);
4321 return 0;
4322}
4323
4324static int
4325load_list(UnpicklerObject *self)
4326{
4327 PyObject *list;
4328 int i;
4329
4330 if ((i = marker(self)) < 0)
4331 return -1;
4332
4333 list = Pdata_poplist(self->stack, i);
4334 if (list == NULL)
4335 return -1;
4336 PDATA_PUSH(self->stack, list, -1);
4337 return 0;
4338}
4339
4340static int
4341load_dict(UnpicklerObject *self)
4342{
4343 PyObject *dict, *key, *value;
4344 int i, j, k;
4345
4346 if ((i = marker(self)) < 0)
4347 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004348 j = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004349
4350 if ((dict = PyDict_New()) == NULL)
4351 return -1;
4352
4353 for (k = i + 1; k < j; k += 2) {
4354 key = self->stack->data[k - 1];
4355 value = self->stack->data[k];
4356 if (PyDict_SetItem(dict, key, value) < 0) {
4357 Py_DECREF(dict);
4358 return -1;
4359 }
4360 }
4361 Pdata_clear(self->stack, i);
4362 PDATA_PUSH(self->stack, dict, -1);
4363 return 0;
4364}
4365
4366static PyObject *
4367instantiate(PyObject *cls, PyObject *args)
4368{
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004369 PyObject *result = NULL;
4370 /* Caller must assure args are a tuple. Normally, args come from
4371 Pdata_poptuple which packs objects from the top of the stack
4372 into a newly created tuple. */
4373 assert(PyTuple_Check(args));
4374 if (Py_SIZE(args) > 0 || !PyType_Check(cls) ||
4375 PyObject_HasAttrString(cls, "__getinitargs__")) {
4376 result = PyObject_CallObject(cls, args);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004377 }
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004378 else {
4379 result = PyObject_CallMethod(cls, "__new__", "O", cls);
4380 }
4381 return result;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004382}
4383
4384static int
4385load_obj(UnpicklerObject *self)
4386{
4387 PyObject *cls, *args, *obj = NULL;
4388 int i;
4389
4390 if ((i = marker(self)) < 0)
4391 return -1;
4392
4393 args = Pdata_poptuple(self->stack, i + 1);
4394 if (args == NULL)
4395 return -1;
4396
4397 PDATA_POP(self->stack, cls);
4398 if (cls) {
4399 obj = instantiate(cls, args);
4400 Py_DECREF(cls);
4401 }
4402 Py_DECREF(args);
4403 if (obj == NULL)
4404 return -1;
4405
4406 PDATA_PUSH(self->stack, obj, -1);
4407 return 0;
4408}
4409
4410static int
4411load_inst(UnpicklerObject *self)
4412{
4413 PyObject *cls = NULL;
4414 PyObject *args = NULL;
4415 PyObject *obj = NULL;
4416 PyObject *module_name;
4417 PyObject *class_name;
4418 Py_ssize_t len;
4419 int i;
4420 char *s;
4421
4422 if ((i = marker(self)) < 0)
4423 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004424 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004425 return -1;
4426 if (len < 2)
4427 return bad_readline();
4428
4429 /* Here it is safe to use PyUnicode_DecodeASCII(), even though non-ASCII
4430 identifiers are permitted in Python 3.0, since the INST opcode is only
4431 supported by older protocols on Python 2.x. */
4432 module_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
4433 if (module_name == NULL)
4434 return -1;
4435
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004436 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004437 if (len < 2)
4438 return bad_readline();
4439 class_name = PyUnicode_DecodeASCII(s, len - 1, "strict");
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00004440 if (class_name != NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004441 cls = find_class(self, module_name, class_name);
4442 Py_DECREF(class_name);
4443 }
4444 }
4445 Py_DECREF(module_name);
4446
4447 if (cls == NULL)
4448 return -1;
4449
4450 if ((args = Pdata_poptuple(self->stack, i)) != NULL) {
4451 obj = instantiate(cls, args);
4452 Py_DECREF(args);
4453 }
4454 Py_DECREF(cls);
4455
4456 if (obj == NULL)
4457 return -1;
4458
4459 PDATA_PUSH(self->stack, obj, -1);
4460 return 0;
4461}
4462
4463static int
4464load_newobj(UnpicklerObject *self)
4465{
4466 PyObject *args = NULL;
4467 PyObject *clsraw = NULL;
4468 PyTypeObject *cls; /* clsraw cast to its true type */
4469 PyObject *obj;
4470
4471 /* Stack is ... cls argtuple, and we want to call
4472 * cls.__new__(cls, *argtuple).
4473 */
4474 PDATA_POP(self->stack, args);
4475 if (args == NULL)
4476 goto error;
4477 if (!PyTuple_Check(args)) {
4478 PyErr_SetString(UnpicklingError, "NEWOBJ expected an arg " "tuple.");
4479 goto error;
4480 }
4481
4482 PDATA_POP(self->stack, clsraw);
4483 cls = (PyTypeObject *)clsraw;
4484 if (cls == NULL)
4485 goto error;
4486 if (!PyType_Check(cls)) {
4487 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4488 "isn't a type object");
4489 goto error;
4490 }
4491 if (cls->tp_new == NULL) {
4492 PyErr_SetString(UnpicklingError, "NEWOBJ class argument "
4493 "has NULL tp_new");
4494 goto error;
4495 }
4496
4497 /* Call __new__. */
4498 obj = cls->tp_new(cls, args, NULL);
4499 if (obj == NULL)
4500 goto error;
4501
4502 Py_DECREF(args);
4503 Py_DECREF(clsraw);
4504 PDATA_PUSH(self->stack, obj, -1);
4505 return 0;
4506
4507 error:
4508 Py_XDECREF(args);
4509 Py_XDECREF(clsraw);
4510 return -1;
4511}
4512
4513static int
4514load_global(UnpicklerObject *self)
4515{
4516 PyObject *global = NULL;
4517 PyObject *module_name;
4518 PyObject *global_name;
4519 Py_ssize_t len;
4520 char *s;
4521
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004522 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004523 return -1;
4524 if (len < 2)
4525 return bad_readline();
4526 module_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4527 if (!module_name)
4528 return -1;
4529
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004530 if ((len = _Unpickler_Readline(self, &s)) >= 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004531 if (len < 2) {
4532 Py_DECREF(module_name);
4533 return bad_readline();
4534 }
4535 global_name = PyUnicode_DecodeUTF8(s, len - 1, "strict");
4536 if (global_name) {
4537 global = find_class(self, module_name, global_name);
4538 Py_DECREF(global_name);
4539 }
4540 }
4541 Py_DECREF(module_name);
4542
4543 if (global == NULL)
4544 return -1;
4545 PDATA_PUSH(self->stack, global, -1);
4546 return 0;
4547}
4548
4549static int
4550load_persid(UnpicklerObject *self)
4551{
4552 PyObject *pid;
4553 Py_ssize_t len;
4554 char *s;
4555
4556 if (self->pers_func) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004557 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004558 return -1;
4559 if (len < 2)
4560 return bad_readline();
4561
4562 pid = PyBytes_FromStringAndSize(s, len - 1);
4563 if (pid == NULL)
4564 return -1;
4565
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004566 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004567 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004568 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004569 if (pid == NULL)
4570 return -1;
4571
4572 PDATA_PUSH(self->stack, pid, -1);
4573 return 0;
4574 }
4575 else {
4576 PyErr_SetString(UnpicklingError,
4577 "A load persistent id instruction was encountered,\n"
4578 "but no persistent_load function was specified.");
4579 return -1;
4580 }
4581}
4582
4583static int
4584load_binpersid(UnpicklerObject *self)
4585{
4586 PyObject *pid;
4587
4588 if (self->pers_func) {
4589 PDATA_POP(self->stack, pid);
4590 if (pid == NULL)
4591 return -1;
4592
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004593 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004594 reference to pid first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004595 pid = _Unpickler_FastCall(self, self->pers_func, pid);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004596 if (pid == NULL)
4597 return -1;
4598
4599 PDATA_PUSH(self->stack, pid, -1);
4600 return 0;
4601 }
4602 else {
4603 PyErr_SetString(UnpicklingError,
4604 "A load persistent id instruction was encountered,\n"
4605 "but no persistent_load function was specified.");
4606 return -1;
4607 }
4608}
4609
4610static int
4611load_pop(UnpicklerObject *self)
4612{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004613 int len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004614
4615 /* Note that we split the (pickle.py) stack into two stacks,
4616 * an object stack and a mark stack. We have to be clever and
4617 * pop the right one. We do this by looking at the top of the
Collin Winter8ca69de2009-05-26 16:53:41 +00004618 * mark stack first, and only signalling a stack underflow if
4619 * the object stack is empty and the mark stack doesn't match
4620 * our expectations.
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004621 */
Collin Winter8ca69de2009-05-26 16:53:41 +00004622 if (self->num_marks > 0 && self->marks[self->num_marks - 1] == len) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004623 self->num_marks--;
Antoine Pitrou01a15ea2010-01-07 17:57:31 +00004624 } else if (len > 0) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004625 len--;
4626 Py_DECREF(self->stack->data[len]);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004627 Py_SIZE(self->stack) = len;
Collin Winter8ca69de2009-05-26 16:53:41 +00004628 } else {
4629 return stack_underflow();
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004630 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004631 return 0;
4632}
4633
4634static int
4635load_pop_mark(UnpicklerObject *self)
4636{
4637 int i;
4638
4639 if ((i = marker(self)) < 0)
4640 return -1;
4641
4642 Pdata_clear(self->stack, i);
4643
4644 return 0;
4645}
4646
4647static int
4648load_dup(UnpicklerObject *self)
4649{
4650 PyObject *last;
4651 int len;
4652
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004653 if ((len = Py_SIZE(self->stack)) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004654 return stack_underflow();
4655 last = self->stack->data[len - 1];
4656 PDATA_APPEND(self->stack, last, -1);
4657 return 0;
4658}
4659
4660static int
4661load_get(UnpicklerObject *self)
4662{
4663 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004664 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004665 Py_ssize_t len;
4666 char *s;
4667
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004668 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004669 return -1;
4670 if (len < 2)
4671 return bad_readline();
4672
4673 key = PyLong_FromString(s, NULL, 10);
4674 if (key == NULL)
4675 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004676 idx = PyLong_AsSsize_t(key);
4677 if (idx == -1 && PyErr_Occurred()) {
4678 Py_DECREF(key);
4679 return -1;
4680 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004681
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004682 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004683 if (value == NULL) {
4684 if (!PyErr_Occurred())
4685 PyErr_SetObject(PyExc_KeyError, key);
4686 Py_DECREF(key);
4687 return -1;
4688 }
4689 Py_DECREF(key);
4690
4691 PDATA_APPEND(self->stack, value, -1);
4692 return 0;
4693}
4694
4695static int
4696load_binget(UnpicklerObject *self)
4697{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004698 PyObject *value;
4699 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004700 char *s;
4701
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004702 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004703 return -1;
4704
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004705 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004706
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004707 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004708 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004709 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004710 if (!PyErr_Occurred())
4711 PyErr_SetObject(PyExc_KeyError, key);
4712 Py_DECREF(key);
4713 return -1;
4714 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004715
4716 PDATA_APPEND(self->stack, value, -1);
4717 return 0;
4718}
4719
4720static int
4721load_long_binget(UnpicklerObject *self)
4722{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004723 PyObject *value;
4724 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004725 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004726
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004727 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004728 return -1;
4729
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004730 idx = (long)Py_CHARMASK(s[0]);
4731 idx |= (long)Py_CHARMASK(s[1]) << 8;
4732 idx |= (long)Py_CHARMASK(s[2]) << 16;
4733 idx |= (long)Py_CHARMASK(s[3]) << 24;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004734
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004735 value = _Unpickler_MemoGet(self, idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004736 if (value == NULL) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004737 PyObject *key = PyLong_FromSsize_t(idx);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004738 if (!PyErr_Occurred())
4739 PyErr_SetObject(PyExc_KeyError, key);
4740 Py_DECREF(key);
4741 return -1;
4742 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004743
4744 PDATA_APPEND(self->stack, value, -1);
4745 return 0;
4746}
4747
4748/* Push an object from the extension registry (EXT[124]). nbytes is
4749 * the number of bytes following the opcode, holding the index (code) value.
4750 */
4751static int
4752load_extension(UnpicklerObject *self, int nbytes)
4753{
4754 char *codebytes; /* the nbytes bytes after the opcode */
4755 long code; /* calc_binint returns long */
4756 PyObject *py_code; /* code as a Python int */
4757 PyObject *obj; /* the object to push */
4758 PyObject *pair; /* (module_name, class_name) */
4759 PyObject *module_name, *class_name;
4760
4761 assert(nbytes == 1 || nbytes == 2 || nbytes == 4);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004762 if (_Unpickler_Read(self, &codebytes, nbytes) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004763 return -1;
4764 code = calc_binint(codebytes, nbytes);
4765 if (code <= 0) { /* note that 0 is forbidden */
4766 /* Corrupt or hostile pickle. */
4767 PyErr_SetString(UnpicklingError, "EXT specifies code <= 0");
4768 return -1;
4769 }
4770
4771 /* Look for the code in the cache. */
4772 py_code = PyLong_FromLong(code);
4773 if (py_code == NULL)
4774 return -1;
4775 obj = PyDict_GetItem(extension_cache, py_code);
4776 if (obj != NULL) {
4777 /* Bingo. */
4778 Py_DECREF(py_code);
4779 PDATA_APPEND(self->stack, obj, -1);
4780 return 0;
4781 }
4782
4783 /* Look up the (module_name, class_name) pair. */
4784 pair = PyDict_GetItem(inverted_registry, py_code);
4785 if (pair == NULL) {
4786 Py_DECREF(py_code);
4787 PyErr_Format(PyExc_ValueError, "unregistered extension "
4788 "code %ld", code);
4789 return -1;
4790 }
4791 /* Since the extension registry is manipulable via Python code,
4792 * confirm that pair is really a 2-tuple of strings.
4793 */
4794 if (!PyTuple_Check(pair) || PyTuple_Size(pair) != 2 ||
4795 !PyUnicode_Check(module_name = PyTuple_GET_ITEM(pair, 0)) ||
4796 !PyUnicode_Check(class_name = PyTuple_GET_ITEM(pair, 1))) {
4797 Py_DECREF(py_code);
4798 PyErr_Format(PyExc_ValueError, "_inverted_registry[%ld] "
4799 "isn't a 2-tuple of strings", code);
4800 return -1;
4801 }
4802 /* Load the object. */
4803 obj = find_class(self, module_name, class_name);
4804 if (obj == NULL) {
4805 Py_DECREF(py_code);
4806 return -1;
4807 }
4808 /* Cache code -> obj. */
4809 code = PyDict_SetItem(extension_cache, py_code, obj);
4810 Py_DECREF(py_code);
4811 if (code < 0) {
4812 Py_DECREF(obj);
4813 return -1;
4814 }
4815 PDATA_PUSH(self->stack, obj, -1);
4816 return 0;
4817}
4818
4819static int
4820load_put(UnpicklerObject *self)
4821{
4822 PyObject *key, *value;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004823 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004824 Py_ssize_t len;
4825 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004826
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004827 if ((len = _Unpickler_Readline(self, &s)) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004828 return -1;
4829 if (len < 2)
4830 return bad_readline();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004831 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004832 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004833 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004834
4835 key = PyLong_FromString(s, NULL, 10);
4836 if (key == NULL)
4837 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004838 idx = PyLong_AsSsize_t(key);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004839 Py_DECREF(key);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004840 if (idx == -1 && PyErr_Occurred())
4841 return -1;
4842
4843 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004844}
4845
4846static int
4847load_binput(UnpicklerObject *self)
4848{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004849 PyObject *value;
4850 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004851 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004852
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004853 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004854 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004855
4856 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004857 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004858 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004859
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004860 idx = Py_CHARMASK(s[0]);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004861
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004862 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004863}
4864
4865static int
4866load_long_binput(UnpicklerObject *self)
4867{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004868 PyObject *value;
4869 Py_ssize_t idx;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004870 char *s;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004871
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004872 if (_Unpickler_Read(self, &s, 4) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004873 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004874
4875 if (Py_SIZE(self->stack) <= 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004876 return stack_underflow();
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004877 value = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004878
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004879 idx = (long)Py_CHARMASK(s[0]);
4880 idx |= (long)Py_CHARMASK(s[1]) << 8;
4881 idx |= (long)Py_CHARMASK(s[2]) << 16;
4882 idx |= (long)Py_CHARMASK(s[3]) << 24;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004883
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004884 return _Unpickler_MemoPut(self, idx, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004885}
4886
4887static int
4888do_append(UnpicklerObject *self, int x)
4889{
4890 PyObject *value;
4891 PyObject *list;
4892 int len, i;
4893
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004894 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004895 if (x > len || x <= 0)
4896 return stack_underflow();
4897 if (len == x) /* nothing to do */
4898 return 0;
4899
4900 list = self->stack->data[x - 1];
4901
4902 if (PyList_Check(list)) {
4903 PyObject *slice;
4904 Py_ssize_t list_len;
4905
4906 slice = Pdata_poplist(self->stack, x);
4907 if (!slice)
4908 return -1;
4909 list_len = PyList_GET_SIZE(list);
4910 i = PyList_SetSlice(list, list_len, list_len, slice);
4911 Py_DECREF(slice);
4912 return i;
4913 }
4914 else {
4915 PyObject *append_func;
4916
4917 append_func = PyObject_GetAttrString(list, "append");
4918 if (append_func == NULL)
4919 return -1;
4920 for (i = x; i < len; i++) {
4921 PyObject *result;
4922
4923 value = self->stack->data[i];
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004924 result = _Unpickler_FastCall(self, append_func, value);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004925 if (result == NULL) {
4926 Pdata_clear(self->stack, i + 1);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004927 Py_SIZE(self->stack) = x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004928 return -1;
4929 }
4930 Py_DECREF(result);
4931 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004932 Py_SIZE(self->stack) = x;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004933 }
4934
4935 return 0;
4936}
4937
4938static int
4939load_append(UnpicklerObject *self)
4940{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004941 return do_append(self, Py_SIZE(self->stack) - 1);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004942}
4943
4944static int
4945load_appends(UnpicklerObject *self)
4946{
4947 return do_append(self, marker(self));
4948}
4949
4950static int
4951do_setitems(UnpicklerObject *self, int x)
4952{
4953 PyObject *value, *key;
4954 PyObject *dict;
4955 int len, i;
4956 int status = 0;
4957
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004958 len = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004959 if (x > len || x <= 0)
4960 return stack_underflow();
4961 if (len == x) /* nothing to do */
4962 return 0;
4963 if ((len - x) % 2 != 0) {
4964 /* Currupt or hostile pickle -- we never write one like this. */
4965 PyErr_SetString(UnpicklingError, "odd number of items for SETITEMS");
4966 return -1;
4967 }
4968
4969 /* Here, dict does not actually need to be a PyDict; it could be anything
4970 that supports the __setitem__ attribute. */
4971 dict = self->stack->data[x - 1];
4972
4973 for (i = x + 1; i < len; i += 2) {
4974 key = self->stack->data[i - 1];
4975 value = self->stack->data[i];
4976 if (PyObject_SetItem(dict, key, value) < 0) {
4977 status = -1;
4978 break;
4979 }
4980 }
4981
4982 Pdata_clear(self->stack, x);
4983 return status;
4984}
4985
4986static int
4987load_setitem(UnpicklerObject *self)
4988{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00004989 return do_setitems(self, Py_SIZE(self->stack) - 2);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00004990}
4991
4992static int
4993load_setitems(UnpicklerObject *self)
4994{
4995 return do_setitems(self, marker(self));
4996}
4997
4998static int
4999load_build(UnpicklerObject *self)
5000{
5001 PyObject *state, *inst, *slotstate;
5002 PyObject *setstate;
5003 int status = 0;
5004
5005 /* Stack is ... instance, state. We want to leave instance at
5006 * the stack top, possibly mutated via instance.__setstate__(state).
5007 */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005008 if (Py_SIZE(self->stack) < 2)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005009 return stack_underflow();
5010
5011 PDATA_POP(self->stack, state);
5012 if (state == NULL)
5013 return -1;
5014
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005015 inst = self->stack->data[Py_SIZE(self->stack) - 1];
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005016
5017 setstate = PyObject_GetAttrString(inst, "__setstate__");
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005018 if (setstate == NULL) {
5019 if (PyErr_ExceptionMatches(PyExc_AttributeError))
5020 PyErr_Clear();
Antoine Pitroud79dc622008-09-05 00:03:33 +00005021 else {
5022 Py_DECREF(state);
Alexandre Vassalotti1f9d9072008-08-15 03:07:47 +00005023 return -1;
Antoine Pitroud79dc622008-09-05 00:03:33 +00005024 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005025 }
5026 else {
5027 PyObject *result;
5028
5029 /* The explicit __setstate__ is responsible for everything. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005030 /* Ugh... this does not leak since _Unpickler_FastCall() steals the
Antoine Pitroud79dc622008-09-05 00:03:33 +00005031 reference to state first. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005032 result = _Unpickler_FastCall(self, setstate, state);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005033 Py_DECREF(setstate);
5034 if (result == NULL)
5035 return -1;
5036 Py_DECREF(result);
5037 return 0;
5038 }
5039
5040 /* A default __setstate__. First see whether state embeds a
5041 * slot state dict too (a proto 2 addition).
5042 */
5043 if (PyTuple_Check(state) && Py_SIZE(state) == 2) {
5044 PyObject *tmp = state;
5045
5046 state = PyTuple_GET_ITEM(tmp, 0);
5047 slotstate = PyTuple_GET_ITEM(tmp, 1);
5048 Py_INCREF(state);
5049 Py_INCREF(slotstate);
5050 Py_DECREF(tmp);
5051 }
5052 else
5053 slotstate = NULL;
5054
5055 /* Set inst.__dict__ from the state dict (if any). */
5056 if (state != Py_None) {
5057 PyObject *dict;
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005058 PyObject *d_key, *d_value;
5059 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005060
5061 if (!PyDict_Check(state)) {
5062 PyErr_SetString(UnpicklingError, "state is not a dictionary");
5063 goto error;
5064 }
5065 dict = PyObject_GetAttrString(inst, "__dict__");
5066 if (dict == NULL)
5067 goto error;
5068
Antoine Pitroua9f48a02009-05-02 21:41:14 +00005069 i = 0;
5070 while (PyDict_Next(state, &i, &d_key, &d_value)) {
5071 /* normally the keys for instance attributes are
5072 interned. we should try to do that here. */
5073 Py_INCREF(d_key);
5074 if (PyUnicode_CheckExact(d_key))
5075 PyUnicode_InternInPlace(&d_key);
5076 if (PyObject_SetItem(dict, d_key, d_value) < 0) {
5077 Py_DECREF(d_key);
5078 goto error;
5079 }
5080 Py_DECREF(d_key);
5081 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005082 Py_DECREF(dict);
5083 }
5084
5085 /* Also set instance attributes from the slotstate dict (if any). */
5086 if (slotstate != NULL) {
5087 PyObject *d_key, *d_value;
5088 Py_ssize_t i;
5089
5090 if (!PyDict_Check(slotstate)) {
5091 PyErr_SetString(UnpicklingError,
5092 "slot state is not a dictionary");
5093 goto error;
5094 }
5095 i = 0;
5096 while (PyDict_Next(slotstate, &i, &d_key, &d_value)) {
5097 if (PyObject_SetAttr(inst, d_key, d_value) < 0)
5098 goto error;
5099 }
5100 }
5101
5102 if (0) {
5103 error:
5104 status = -1;
5105 }
5106
5107 Py_DECREF(state);
5108 Py_XDECREF(slotstate);
5109 return status;
5110}
5111
5112static int
5113load_mark(UnpicklerObject *self)
5114{
5115
5116 /* Note that we split the (pickle.py) stack into two stacks, an
5117 * object stack and a mark stack. Here we push a mark onto the
5118 * mark stack.
5119 */
5120
5121 if ((self->num_marks + 1) >= self->marks_size) {
5122 size_t alloc;
5123 int *marks;
5124
5125 /* Use the size_t type to check for overflow. */
5126 alloc = ((size_t)self->num_marks << 1) + 20;
Alexandre Vassalotti7634ff52008-06-13 02:16:06 +00005127 if (alloc > PY_SSIZE_T_MAX ||
5128 alloc <= ((size_t)self->num_marks + 1)) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005129 PyErr_NoMemory();
5130 return -1;
5131 }
5132
5133 if (self->marks == NULL)
5134 marks = (int *)PyMem_Malloc(alloc * sizeof(int));
5135 else
5136 marks = (int *)PyMem_Realloc(self->marks, alloc * sizeof(int));
5137 if (marks == NULL) {
5138 PyErr_NoMemory();
5139 return -1;
5140 }
5141 self->marks = marks;
5142 self->marks_size = (Py_ssize_t)alloc;
5143 }
5144
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005145 self->marks[self->num_marks++] = Py_SIZE(self->stack);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005146
5147 return 0;
5148}
5149
5150static int
5151load_reduce(UnpicklerObject *self)
5152{
5153 PyObject *callable = NULL;
5154 PyObject *argtup = NULL;
5155 PyObject *obj = NULL;
5156
5157 PDATA_POP(self->stack, argtup);
5158 if (argtup == NULL)
5159 return -1;
5160 PDATA_POP(self->stack, callable);
5161 if (callable) {
Alexander Belopolskyd92f0402010-07-17 22:50:45 +00005162 obj = PyObject_CallObject(callable, argtup);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005163 Py_DECREF(callable);
5164 }
5165 Py_DECREF(argtup);
5166
5167 if (obj == NULL)
5168 return -1;
5169
5170 PDATA_PUSH(self->stack, obj, -1);
5171 return 0;
5172}
5173
5174/* Just raises an error if we don't know the protocol specified. PROTO
5175 * is the first opcode for protocols >= 2.
5176 */
5177static int
5178load_proto(UnpicklerObject *self)
5179{
5180 char *s;
5181 int i;
5182
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005183 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005184 return -1;
5185
5186 i = (unsigned char)s[0];
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005187 if (i <= HIGHEST_PROTOCOL) {
5188 self->proto = i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005189 return 0;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005190 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005191
5192 PyErr_Format(PyExc_ValueError, "unsupported pickle protocol: %d", i);
5193 return -1;
5194}
5195
5196static PyObject *
5197load(UnpicklerObject *self)
5198{
5199 PyObject *err;
5200 PyObject *value = NULL;
5201 char *s;
5202
5203 self->num_marks = 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005204 if (Py_SIZE(self->stack))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005205 Pdata_clear(self->stack, 0);
5206
5207 /* Convenient macros for the dispatch while-switch loop just below. */
5208#define OP(opcode, load_func) \
5209 case opcode: if (load_func(self) < 0) break; continue;
5210
5211#define OP_ARG(opcode, load_func, arg) \
5212 case opcode: if (load_func(self, (arg)) < 0) break; continue;
5213
5214 while (1) {
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005215 if (_Unpickler_Read(self, &s, 1) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005216 break;
5217
5218 switch ((enum opcode)s[0]) {
5219 OP(NONE, load_none)
5220 OP(BININT, load_binint)
5221 OP(BININT1, load_binint1)
5222 OP(BININT2, load_binint2)
5223 OP(INT, load_int)
5224 OP(LONG, load_long)
5225 OP_ARG(LONG1, load_counted_long, 1)
5226 OP_ARG(LONG4, load_counted_long, 4)
5227 OP(FLOAT, load_float)
5228 OP(BINFLOAT, load_binfloat)
5229 OP(BINBYTES, load_binbytes)
5230 OP(SHORT_BINBYTES, load_short_binbytes)
5231 OP(BINSTRING, load_binstring)
5232 OP(SHORT_BINSTRING, load_short_binstring)
5233 OP(STRING, load_string)
5234 OP(UNICODE, load_unicode)
5235 OP(BINUNICODE, load_binunicode)
5236 OP_ARG(EMPTY_TUPLE, load_counted_tuple, 0)
5237 OP_ARG(TUPLE1, load_counted_tuple, 1)
5238 OP_ARG(TUPLE2, load_counted_tuple, 2)
5239 OP_ARG(TUPLE3, load_counted_tuple, 3)
5240 OP(TUPLE, load_tuple)
5241 OP(EMPTY_LIST, load_empty_list)
5242 OP(LIST, load_list)
5243 OP(EMPTY_DICT, load_empty_dict)
5244 OP(DICT, load_dict)
5245 OP(OBJ, load_obj)
5246 OP(INST, load_inst)
5247 OP(NEWOBJ, load_newobj)
5248 OP(GLOBAL, load_global)
5249 OP(APPEND, load_append)
5250 OP(APPENDS, load_appends)
5251 OP(BUILD, load_build)
5252 OP(DUP, load_dup)
5253 OP(BINGET, load_binget)
5254 OP(LONG_BINGET, load_long_binget)
5255 OP(GET, load_get)
5256 OP(MARK, load_mark)
5257 OP(BINPUT, load_binput)
5258 OP(LONG_BINPUT, load_long_binput)
5259 OP(PUT, load_put)
5260 OP(POP, load_pop)
5261 OP(POP_MARK, load_pop_mark)
5262 OP(SETITEM, load_setitem)
5263 OP(SETITEMS, load_setitems)
5264 OP(PERSID, load_persid)
5265 OP(BINPERSID, load_binpersid)
5266 OP(REDUCE, load_reduce)
5267 OP(PROTO, load_proto)
5268 OP_ARG(EXT1, load_extension, 1)
5269 OP_ARG(EXT2, load_extension, 2)
5270 OP_ARG(EXT4, load_extension, 4)
5271 OP_ARG(NEWTRUE, load_bool, Py_True)
5272 OP_ARG(NEWFALSE, load_bool, Py_False)
5273
5274 case STOP:
5275 break;
5276
5277 case '\0':
5278 PyErr_SetNone(PyExc_EOFError);
5279 return NULL;
5280
5281 default:
5282 PyErr_Format(UnpicklingError,
5283 "invalid load key, '%c'.", s[0]);
5284 return NULL;
5285 }
5286
5287 break; /* and we are done! */
5288 }
5289
Antoine Pitrou04248a82010-10-12 20:51:21 +00005290 if (_Unpickler_SkipConsumed(self) < 0)
5291 return NULL;
5292
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005293 /* XXX: It is not clear what this is actually for. */
5294 if ((err = PyErr_Occurred())) {
5295 if (err == PyExc_EOFError) {
5296 PyErr_SetNone(PyExc_EOFError);
5297 }
5298 return NULL;
5299 }
5300
5301 PDATA_POP(self->stack, value);
5302 return value;
5303}
5304
5305PyDoc_STRVAR(Unpickler_load_doc,
5306"load() -> object. Load a pickle."
5307"\n"
5308"Read a pickled object representation from the open file object given in\n"
5309"the constructor, and return the reconstituted object hierarchy specified\n"
5310"therein.\n");
5311
5312static PyObject *
5313Unpickler_load(UnpicklerObject *self)
5314{
5315 /* Check whether the Unpickler was initialized correctly. This prevents
5316 segfaulting if a subclass overridden __init__ with a function that does
5317 not call Unpickler.__init__(). Here, we simply ensure that self->read
5318 is not NULL. */
5319 if (self->read == NULL) {
5320 PyErr_Format(UnpicklingError,
5321 "Unpickler.__init__() was not called by %s.__init__()",
5322 Py_TYPE(self)->tp_name);
5323 return NULL;
5324 }
5325
5326 return load(self);
5327}
5328
5329/* The name of find_class() is misleading. In newer pickle protocols, this
5330 function is used for loading any global (i.e., functions), not just
5331 classes. The name is kept only for backward compatibility. */
5332
5333PyDoc_STRVAR(Unpickler_find_class_doc,
5334"find_class(module_name, global_name) -> object.\n"
5335"\n"
5336"Return an object from a specified module, importing the module if\n"
5337"necessary. Subclasses may override this method (e.g. to restrict\n"
5338"unpickling of arbitrary classes and functions).\n"
5339"\n"
5340"This method is called whenever a class or a function object is\n"
5341"needed. Both arguments passed are str objects.\n");
5342
5343static PyObject *
5344Unpickler_find_class(UnpicklerObject *self, PyObject *args)
5345{
5346 PyObject *global;
5347 PyObject *modules_dict;
5348 PyObject *module;
5349 PyObject *module_name, *global_name;
5350
5351 if (!PyArg_UnpackTuple(args, "find_class", 2, 2,
5352 &module_name, &global_name))
5353 return NULL;
5354
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005355 /* Try to map the old names used in Python 2.x to the new ones used in
5356 Python 3.x. We do this only with old pickle protocols and when the
5357 user has not disabled the feature. */
5358 if (self->proto < 3 && self->fix_imports) {
5359 PyObject *key;
5360 PyObject *item;
5361
5362 /* Check if the global (i.e., a function or a class) was renamed
5363 or moved to another module. */
5364 key = PyTuple_Pack(2, module_name, global_name);
5365 if (key == NULL)
5366 return NULL;
5367 item = PyDict_GetItemWithError(name_mapping_2to3, key);
5368 Py_DECREF(key);
5369 if (item) {
5370 if (!PyTuple_Check(item) || PyTuple_GET_SIZE(item) != 2) {
5371 PyErr_Format(PyExc_RuntimeError,
5372 "_compat_pickle.NAME_MAPPING values should be "
5373 "2-tuples, not %.200s", Py_TYPE(item)->tp_name);
5374 return NULL;
5375 }
5376 module_name = PyTuple_GET_ITEM(item, 0);
5377 global_name = PyTuple_GET_ITEM(item, 1);
5378 if (!PyUnicode_Check(module_name) ||
5379 !PyUnicode_Check(global_name)) {
5380 PyErr_Format(PyExc_RuntimeError,
5381 "_compat_pickle.NAME_MAPPING values should be "
5382 "pairs of str, not (%.200s, %.200s)",
5383 Py_TYPE(module_name)->tp_name,
5384 Py_TYPE(global_name)->tp_name);
5385 return NULL;
5386 }
5387 }
5388 else if (PyErr_Occurred()) {
5389 return NULL;
5390 }
5391
5392 /* Check if the module was renamed. */
5393 item = PyDict_GetItemWithError(import_mapping_2to3, module_name);
5394 if (item) {
5395 if (!PyUnicode_Check(item)) {
5396 PyErr_Format(PyExc_RuntimeError,
5397 "_compat_pickle.IMPORT_MAPPING values should be "
5398 "strings, not %.200s", Py_TYPE(item)->tp_name);
5399 return NULL;
5400 }
5401 module_name = item;
5402 }
5403 else if (PyErr_Occurred()) {
5404 return NULL;
5405 }
5406 }
5407
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005408 modules_dict = PySys_GetObject("modules");
5409 if (modules_dict == NULL)
5410 return NULL;
5411
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005412 module = PyDict_GetItemWithError(modules_dict, module_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005413 if (module == NULL) {
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005414 if (PyErr_Occurred())
5415 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005416 module = PyImport_Import(module_name);
5417 if (module == NULL)
5418 return NULL;
5419 global = PyObject_GetAttr(module, global_name);
5420 Py_DECREF(module);
5421 }
5422 else {
5423 global = PyObject_GetAttr(module, global_name);
5424 }
5425 return global;
5426}
5427
5428static struct PyMethodDef Unpickler_methods[] = {
5429 {"load", (PyCFunction)Unpickler_load, METH_NOARGS,
5430 Unpickler_load_doc},
5431 {"find_class", (PyCFunction)Unpickler_find_class, METH_VARARGS,
5432 Unpickler_find_class_doc},
5433 {NULL, NULL} /* sentinel */
5434};
5435
5436static void
5437Unpickler_dealloc(UnpicklerObject *self)
5438{
5439 PyObject_GC_UnTrack((PyObject *)self);
5440 Py_XDECREF(self->readline);
5441 Py_XDECREF(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005442 Py_XDECREF(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005443 Py_XDECREF(self->stack);
5444 Py_XDECREF(self->pers_func);
5445 Py_XDECREF(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005446 if (self->buffer.buf != NULL) {
5447 PyBuffer_Release(&self->buffer);
5448 self->buffer.buf = NULL;
5449 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005450
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005451 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005452 PyMem_Free(self->marks);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005453 PyMem_Free(self->input_line);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005454 free(self->encoding);
5455 free(self->errors);
5456
5457 Py_TYPE(self)->tp_free((PyObject *)self);
5458}
5459
5460static int
5461Unpickler_traverse(UnpicklerObject *self, visitproc visit, void *arg)
5462{
5463 Py_VISIT(self->readline);
5464 Py_VISIT(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005465 Py_VISIT(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005466 Py_VISIT(self->stack);
5467 Py_VISIT(self->pers_func);
5468 Py_VISIT(self->arg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005469 return 0;
5470}
5471
5472static int
5473Unpickler_clear(UnpicklerObject *self)
5474{
5475 Py_CLEAR(self->readline);
5476 Py_CLEAR(self->read);
Antoine Pitrou04248a82010-10-12 20:51:21 +00005477 Py_CLEAR(self->peek);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005478 Py_CLEAR(self->stack);
5479 Py_CLEAR(self->pers_func);
5480 Py_CLEAR(self->arg);
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005481 if (self->buffer.buf != NULL) {
5482 PyBuffer_Release(&self->buffer);
5483 self->buffer.buf = NULL;
5484 }
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005485
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005486 _Unpickler_MemoCleanup(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005487 PyMem_Free(self->marks);
5488 self->marks = NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005489 PyMem_Free(self->input_line);
5490 self->input_line = NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005491 free(self->encoding);
5492 self->encoding = NULL;
5493 free(self->errors);
5494 self->errors = NULL;
5495
5496 return 0;
5497}
5498
5499PyDoc_STRVAR(Unpickler_doc,
5500"Unpickler(file, *, encoding='ASCII', errors='strict')"
5501"\n"
5502"This takes a binary file for reading a pickle data stream.\n"
5503"\n"
5504"The protocol version of the pickle is detected automatically, so no\n"
5505"proto argument is needed.\n"
5506"\n"
5507"The file-like object must have two methods, a read() method\n"
5508"that takes an integer argument, and a readline() method that\n"
5509"requires no arguments. Both methods should return bytes.\n"
5510"Thus file-like object can be a binary file object opened for\n"
5511"reading, a BytesIO object, or any other custom object that\n"
5512"meets this interface.\n"
5513"\n"
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005514"Optional keyword arguments are *fix_imports*, *encoding* and *errors*,\n"
5515"which are used to control compatiblity support for pickle stream\n"
5516"generated by Python 2.x. If *fix_imports* is True, pickle will try to\n"
5517"map the old Python 2.x names to the new names used in Python 3.x. The\n"
5518"*encoding* and *errors* tell pickle how to decode 8-bit string\n"
5519"instances pickled by Python 2.x; these default to 'ASCII' and\n"
5520"'strict', respectively.\n");
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005521
5522static int
5523Unpickler_init(UnpicklerObject *self, PyObject *args, PyObject *kwds)
5524{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005525 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005526 PyObject *file;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005527 PyObject *fix_imports = Py_True;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005528 char *encoding = NULL;
5529 char *errors = NULL;
5530
5531 /* XXX: That is an horrible error message. But, I don't know how to do
5532 better... */
5533 if (Py_SIZE(args) != 1) {
5534 PyErr_Format(PyExc_TypeError,
5535 "%s takes exactly one positional argument (%zd given)",
5536 Py_TYPE(self)->tp_name, Py_SIZE(args));
5537 return -1;
5538 }
5539
5540 /* Arguments parsing needs to be done in the __init__() method to allow
5541 subclasses to define their own __init__() method, which may (or may
5542 not) support Unpickler arguments. However, this means we need to be
5543 extra careful in the other Unpickler methods, since a subclass could
5544 forget to call Unpickler.__init__() thus breaking our internal
5545 invariants. */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005546 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:Unpickler", kwlist,
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005547 &file, &fix_imports, &encoding, &errors))
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005548 return -1;
5549
5550 /* In case of multiple __init__() calls, clear previous content. */
5551 if (self->read != NULL)
5552 (void)Unpickler_clear(self);
5553
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005554 if (_Unpickler_SetInputStream(self, file) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005555 return -1;
5556
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005557 if (_Unpickler_SetInputEncoding(self, encoding, errors) < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005558 return -1;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005559
5560 self->fix_imports = PyObject_IsTrue(fix_imports);
5561 if (self->fix_imports == -1)
5562 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005563
5564 if (PyObject_HasAttrString((PyObject *)self, "persistent_load")) {
5565 self->pers_func = PyObject_GetAttrString((PyObject *)self,
5566 "persistent_load");
5567 if (self->pers_func == NULL)
5568 return -1;
5569 }
5570 else {
5571 self->pers_func = NULL;
5572 }
5573
5574 self->stack = (Pdata *)Pdata_New();
5575 if (self->stack == NULL)
5576 return -1;
5577
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005578 self->memo_size = 32;
5579 self->memo = _Unpickler_NewMemo(self->memo_size);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005580 if (self->memo == NULL)
5581 return -1;
5582
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005583 self->arg = NULL;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00005584 self->proto = 0;
Alexandre Vassalotti0e7aa8c2009-04-03 04:17:41 +00005585
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005586 return 0;
5587}
5588
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005589/* Define a proxy object for the Unpickler's internal memo object. This is to
5590 * avoid breaking code like:
5591 * unpickler.memo.clear()
5592 * and
5593 * unpickler.memo = saved_memo
5594 * Is this a good idea? Not really, but we don't want to break code that uses
5595 * it. Note that we don't implement the entire mapping API here. This is
5596 * intentional, as these should be treated as black-box implementation details.
5597 *
5598 * We do, however, have to implement pickling/unpickling support because of
5599 * real-world code like cvs2svn.
5600 */
5601
5602typedef struct {
5603 PyObject_HEAD
5604 UnpicklerObject *unpickler;
5605} UnpicklerMemoProxyObject;
5606
5607PyDoc_STRVAR(ump_clear_doc,
5608"memo.clear() -> None. Remove all items from memo.");
5609
5610static PyObject *
5611ump_clear(UnpicklerMemoProxyObject *self)
5612{
5613 _Unpickler_MemoCleanup(self->unpickler);
5614 self->unpickler->memo = _Unpickler_NewMemo(self->unpickler->memo_size);
5615 if (self->unpickler->memo == NULL)
5616 return NULL;
5617 Py_RETURN_NONE;
5618}
5619
5620PyDoc_STRVAR(ump_copy_doc,
5621"memo.copy() -> new_memo. Copy the memo to a new object.");
5622
5623static PyObject *
5624ump_copy(UnpicklerMemoProxyObject *self)
5625{
5626 Py_ssize_t i;
5627 PyObject *new_memo = PyDict_New();
5628 if (new_memo == NULL)
5629 return NULL;
5630
5631 for (i = 0; i < self->unpickler->memo_size; i++) {
5632 int status;
5633 PyObject *key, *value;
5634
5635 value = self->unpickler->memo[i];
5636 if (value == NULL)
5637 continue;
5638
5639 key = PyLong_FromSsize_t(i);
5640 if (key == NULL)
5641 goto error;
5642 status = PyDict_SetItem(new_memo, key, value);
5643 Py_DECREF(key);
5644 if (status < 0)
5645 goto error;
5646 }
5647 return new_memo;
5648
5649error:
5650 Py_DECREF(new_memo);
5651 return NULL;
5652}
5653
5654PyDoc_STRVAR(ump_reduce_doc,
5655"memo.__reduce__(). Pickling support.");
5656
5657static PyObject *
5658ump_reduce(UnpicklerMemoProxyObject *self, PyObject *args)
5659{
5660 PyObject *reduce_value;
5661 PyObject *constructor_args;
5662 PyObject *contents = ump_copy(self);
5663 if (contents == NULL)
5664 return NULL;
5665
5666 reduce_value = PyTuple_New(2);
5667 if (reduce_value == NULL) {
5668 Py_DECREF(contents);
5669 return NULL;
5670 }
5671 constructor_args = PyTuple_New(1);
5672 if (constructor_args == NULL) {
5673 Py_DECREF(contents);
5674 Py_DECREF(reduce_value);
5675 return NULL;
5676 }
5677 PyTuple_SET_ITEM(constructor_args, 0, contents);
5678 Py_INCREF((PyObject *)&PyDict_Type);
5679 PyTuple_SET_ITEM(reduce_value, 0, (PyObject *)&PyDict_Type);
5680 PyTuple_SET_ITEM(reduce_value, 1, constructor_args);
5681 return reduce_value;
5682}
5683
5684static PyMethodDef unpicklerproxy_methods[] = {
5685 {"clear", (PyCFunction)ump_clear, METH_NOARGS, ump_clear_doc},
5686 {"copy", (PyCFunction)ump_copy, METH_NOARGS, ump_copy_doc},
5687 {"__reduce__", (PyCFunction)ump_reduce, METH_VARARGS, ump_reduce_doc},
5688 {NULL, NULL} /* sentinel */
5689};
5690
5691static void
5692UnpicklerMemoProxy_dealloc(UnpicklerMemoProxyObject *self)
5693{
5694 PyObject_GC_UnTrack(self);
5695 Py_XDECREF(self->unpickler);
5696 PyObject_GC_Del((PyObject *)self);
5697}
5698
5699static int
5700UnpicklerMemoProxy_traverse(UnpicklerMemoProxyObject *self,
5701 visitproc visit, void *arg)
5702{
5703 Py_VISIT(self->unpickler);
5704 return 0;
5705}
5706
5707static int
5708UnpicklerMemoProxy_clear(UnpicklerMemoProxyObject *self)
5709{
5710 Py_CLEAR(self->unpickler);
5711 return 0;
5712}
5713
5714static PyTypeObject UnpicklerMemoProxyType = {
5715 PyVarObject_HEAD_INIT(NULL, 0)
5716 "_pickle.UnpicklerMemoProxy", /*tp_name*/
5717 sizeof(UnpicklerMemoProxyObject), /*tp_basicsize*/
5718 0,
5719 (destructor)UnpicklerMemoProxy_dealloc, /* tp_dealloc */
5720 0, /* tp_print */
5721 0, /* tp_getattr */
5722 0, /* tp_setattr */
5723 0, /* tp_compare */
5724 0, /* tp_repr */
5725 0, /* tp_as_number */
5726 0, /* tp_as_sequence */
5727 0, /* tp_as_mapping */
Georg Brandlf038b322010-10-18 07:35:09 +00005728 PyObject_HashNotImplemented, /* tp_hash */
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005729 0, /* tp_call */
5730 0, /* tp_str */
5731 PyObject_GenericGetAttr, /* tp_getattro */
5732 PyObject_GenericSetAttr, /* tp_setattro */
5733 0, /* tp_as_buffer */
5734 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5735 0, /* tp_doc */
5736 (traverseproc)UnpicklerMemoProxy_traverse, /* tp_traverse */
5737 (inquiry)UnpicklerMemoProxy_clear, /* tp_clear */
5738 0, /* tp_richcompare */
5739 0, /* tp_weaklistoffset */
5740 0, /* tp_iter */
5741 0, /* tp_iternext */
5742 unpicklerproxy_methods, /* tp_methods */
5743};
5744
5745static PyObject *
5746UnpicklerMemoProxy_New(UnpicklerObject *unpickler)
5747{
5748 UnpicklerMemoProxyObject *self;
5749
5750 self = PyObject_GC_New(UnpicklerMemoProxyObject,
5751 &UnpicklerMemoProxyType);
5752 if (self == NULL)
5753 return NULL;
5754 Py_INCREF(unpickler);
5755 self->unpickler = unpickler;
5756 PyObject_GC_Track(self);
5757 return (PyObject *)self;
5758}
5759
5760/*****************************************************************************/
5761
5762
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005763static PyObject *
5764Unpickler_get_memo(UnpicklerObject *self)
5765{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005766 return UnpicklerMemoProxy_New(self);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005767}
5768
5769static int
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005770Unpickler_set_memo(UnpicklerObject *self, PyObject *obj)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005771{
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005772 PyObject **new_memo;
5773 Py_ssize_t new_memo_size = 0;
5774 Py_ssize_t i;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005775
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005776 if (obj == NULL) {
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005777 PyErr_SetString(PyExc_TypeError,
5778 "attribute deletion is not supported");
5779 return -1;
5780 }
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005781
5782 if (Py_TYPE(obj) == &UnpicklerMemoProxyType) {
5783 UnpicklerObject *unpickler =
5784 ((UnpicklerMemoProxyObject *)obj)->unpickler;
5785
5786 new_memo_size = unpickler->memo_size;
5787 new_memo = _Unpickler_NewMemo(new_memo_size);
5788 if (new_memo == NULL)
5789 return -1;
5790
5791 for (i = 0; i < new_memo_size; i++) {
5792 Py_XINCREF(unpickler->memo[i]);
5793 new_memo[i] = unpickler->memo[i];
5794 }
5795 }
5796 else if (PyDict_Check(obj)) {
5797 Py_ssize_t i = 0;
5798 PyObject *key, *value;
5799
5800 new_memo_size = PyDict_Size(obj);
5801 new_memo = _Unpickler_NewMemo(new_memo_size);
5802 if (new_memo == NULL)
5803 return -1;
5804
5805 while (PyDict_Next(obj, &i, &key, &value)) {
5806 Py_ssize_t idx;
5807 if (!PyLong_Check(key)) {
5808 PyErr_SetString(PyExc_TypeError,
5809 "memo key must be integers");
5810 goto error;
5811 }
5812 idx = PyLong_AsSsize_t(key);
5813 if (idx == -1 && PyErr_Occurred())
5814 goto error;
5815 if (_Unpickler_MemoPut(self, idx, value) < 0)
5816 goto error;
5817 }
5818 }
5819 else {
5820 PyErr_Format(PyExc_TypeError,
5821 "'memo' attribute must be an UnpicklerMemoProxy object"
5822 "or dict, not %.200s", Py_TYPE(obj)->tp_name);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005823 return -1;
5824 }
5825
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005826 _Unpickler_MemoCleanup(self);
5827 self->memo_size = new_memo_size;
5828 self->memo = new_memo;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005829
5830 return 0;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005831
5832 error:
5833 if (new_memo_size) {
5834 i = new_memo_size;
5835 while (--i >= 0) {
5836 Py_XDECREF(new_memo[i]);
5837 }
5838 PyMem_FREE(new_memo);
5839 }
5840 return -1;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005841}
5842
5843static PyObject *
5844Unpickler_get_persload(UnpicklerObject *self)
5845{
5846 if (self->pers_func == NULL)
5847 PyErr_SetString(PyExc_AttributeError, "persistent_load");
5848 else
5849 Py_INCREF(self->pers_func);
5850 return self->pers_func;
5851}
5852
5853static int
5854Unpickler_set_persload(UnpicklerObject *self, PyObject *value)
5855{
5856 PyObject *tmp;
5857
5858 if (value == NULL) {
5859 PyErr_SetString(PyExc_TypeError,
5860 "attribute deletion is not supported");
5861 return -1;
5862 }
5863 if (!PyCallable_Check(value)) {
5864 PyErr_SetString(PyExc_TypeError,
5865 "persistent_load must be a callable taking "
5866 "one argument");
5867 return -1;
5868 }
5869
5870 tmp = self->pers_func;
5871 Py_INCREF(value);
5872 self->pers_func = value;
5873 Py_XDECREF(tmp); /* self->pers_func can be NULL, so be careful. */
5874
5875 return 0;
5876}
5877
5878static PyGetSetDef Unpickler_getsets[] = {
5879 {"memo", (getter)Unpickler_get_memo, (setter)Unpickler_set_memo},
5880 {"persistent_load", (getter)Unpickler_get_persload,
5881 (setter)Unpickler_set_persload},
5882 {NULL}
5883};
5884
5885static PyTypeObject Unpickler_Type = {
5886 PyVarObject_HEAD_INIT(NULL, 0)
5887 "_pickle.Unpickler", /*tp_name*/
5888 sizeof(UnpicklerObject), /*tp_basicsize*/
5889 0, /*tp_itemsize*/
5890 (destructor)Unpickler_dealloc, /*tp_dealloc*/
5891 0, /*tp_print*/
5892 0, /*tp_getattr*/
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005893 0, /*tp_setattr*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00005894 0, /*tp_reserved*/
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00005895 0, /*tp_repr*/
5896 0, /*tp_as_number*/
5897 0, /*tp_as_sequence*/
5898 0, /*tp_as_mapping*/
5899 0, /*tp_hash*/
5900 0, /*tp_call*/
5901 0, /*tp_str*/
5902 0, /*tp_getattro*/
5903 0, /*tp_setattro*/
5904 0, /*tp_as_buffer*/
5905 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
5906 Unpickler_doc, /*tp_doc*/
5907 (traverseproc)Unpickler_traverse, /*tp_traverse*/
5908 (inquiry)Unpickler_clear, /*tp_clear*/
5909 0, /*tp_richcompare*/
5910 0, /*tp_weaklistoffset*/
5911 0, /*tp_iter*/
5912 0, /*tp_iternext*/
5913 Unpickler_methods, /*tp_methods*/
5914 0, /*tp_members*/
5915 Unpickler_getsets, /*tp_getset*/
5916 0, /*tp_base*/
5917 0, /*tp_dict*/
5918 0, /*tp_descr_get*/
5919 0, /*tp_descr_set*/
5920 0, /*tp_dictoffset*/
5921 (initproc)Unpickler_init, /*tp_init*/
5922 PyType_GenericAlloc, /*tp_alloc*/
5923 PyType_GenericNew, /*tp_new*/
5924 PyObject_GC_Del, /*tp_free*/
5925 0, /*tp_is_gc*/
5926};
5927
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00005928PyDoc_STRVAR(pickle_dump_doc,
5929"dump(obj, file, protocol=None, *, fix_imports=True) -> None\n"
5930"\n"
5931"Write a pickled representation of obj to the open file object file. This\n"
5932"is equivalent to ``Pickler(file, protocol).dump(obj)``, but may be more\n"
5933"efficient.\n"
5934"\n"
5935"The optional protocol argument tells the pickler to use the given protocol;\n"
5936"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
5937"backward-incompatible protocol designed for Python 3.0.\n"
5938"\n"
5939"Specifying a negative protocol version selects the highest protocol version\n"
5940"supported. The higher the protocol used, the more recent the version of\n"
5941"Python needed to read the pickle produced.\n"
5942"\n"
5943"The file argument must have a write() method that accepts a single bytes\n"
5944"argument. It can thus be a file object opened for binary writing, a\n"
5945"io.BytesIO instance, or any other custom object that meets this interface.\n"
5946"\n"
5947"If fix_imports is True and protocol is less than 3, pickle will try to\n"
5948"map the new Python 3.x names to the old module names used in Python 2.x,\n"
5949"so that the pickle data stream is readable with Python 2.x.\n");
5950
5951static PyObject *
5952pickle_dump(PyObject *self, PyObject *args, PyObject *kwds)
5953{
5954 static char *kwlist[] = {"obj", "file", "protocol", "fix_imports", 0};
5955 PyObject *obj;
5956 PyObject *file;
5957 PyObject *proto = NULL;
5958 PyObject *fix_imports = Py_True;
5959 PicklerObject *pickler;
5960
5961 /* fix_imports is a keyword-only argument. */
5962 if (Py_SIZE(args) > 3) {
5963 PyErr_Format(PyExc_TypeError,
5964 "pickle.dump() takes at most 3 positional "
5965 "argument (%zd given)", Py_SIZE(args));
5966 return NULL;
5967 }
5968
5969 if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|OO:dump", kwlist,
5970 &obj, &file, &proto, &fix_imports))
5971 return NULL;
5972
5973 pickler = _Pickler_New();
5974 if (pickler == NULL)
5975 return NULL;
5976
5977 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
5978 goto error;
5979
5980 if (_Pickler_SetOutputStream(pickler, file) < 0)
5981 goto error;
5982
5983 if (dump(pickler, obj) < 0)
5984 goto error;
5985
5986 if (_Pickler_FlushToFile(pickler) < 0)
5987 goto error;
5988
5989 Py_DECREF(pickler);
5990 Py_RETURN_NONE;
5991
5992 error:
5993 Py_XDECREF(pickler);
5994 return NULL;
5995}
5996
5997PyDoc_STRVAR(pickle_dumps_doc,
5998"dumps(obj, protocol=None, *, fix_imports=True) -> bytes\n"
5999"\n"
6000"Return the pickled representation of the object as a bytes\n"
6001"object, instead of writing it to a file.\n"
6002"\n"
6003"The optional protocol argument tells the pickler to use the given protocol;\n"
6004"supported protocols are 0, 1, 2, 3. The default protocol is 3; a\n"
6005"backward-incompatible protocol designed for Python 3.0.\n"
6006"\n"
6007"Specifying a negative protocol version selects the highest protocol version\n"
6008"supported. The higher the protocol used, the more recent the version of\n"
6009"Python needed to read the pickle produced.\n"
6010"\n"
6011"If fix_imports is True and *protocol* is less than 3, pickle will try to\n"
6012"map the new Python 3.x names to the old module names used in Python 2.x,\n"
6013"so that the pickle data stream is readable with Python 2.x.\n");
6014
6015static PyObject *
6016pickle_dumps(PyObject *self, PyObject *args, PyObject *kwds)
6017{
6018 static char *kwlist[] = {"obj", "protocol", "fix_imports", 0};
6019 PyObject *obj;
6020 PyObject *proto = NULL;
6021 PyObject *result;
6022 PyObject *fix_imports = Py_True;
6023 PicklerObject *pickler;
6024
6025 /* fix_imports is a keyword-only argument. */
6026 if (Py_SIZE(args) > 2) {
6027 PyErr_Format(PyExc_TypeError,
6028 "pickle.dumps() takes at most 2 positional "
6029 "argument (%zd given)", Py_SIZE(args));
6030 return NULL;
6031 }
6032
6033 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:dumps", kwlist,
6034 &obj, &proto, &fix_imports))
6035 return NULL;
6036
6037 pickler = _Pickler_New();
6038 if (pickler == NULL)
6039 return NULL;
6040
6041 if (_Pickler_SetProtocol(pickler, proto, fix_imports) < 0)
6042 goto error;
6043
6044 if (dump(pickler, obj) < 0)
6045 goto error;
6046
6047 result = _Pickler_GetString(pickler);
6048 Py_DECREF(pickler);
6049 return result;
6050
6051 error:
6052 Py_XDECREF(pickler);
6053 return NULL;
6054}
6055
6056PyDoc_STRVAR(pickle_load_doc,
6057"load(file, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6058"\n"
6059"Read a pickled object representation from the open file object file and\n"
6060"return the reconstituted object hierarchy specified therein. This is\n"
6061"equivalent to ``Unpickler(file).load()``, but may be more efficient.\n"
6062"\n"
6063"The protocol version of the pickle is detected automatically, so no protocol\n"
6064"argument is needed. Bytes past the pickled object's representation are\n"
6065"ignored.\n"
6066"\n"
6067"The argument file must have two methods, a read() method that takes an\n"
6068"integer argument, and a readline() method that requires no arguments. Both\n"
6069"methods should return bytes. Thus *file* can be a binary file object opened\n"
6070"for reading, a BytesIO object, or any other custom object that meets this\n"
6071"interface.\n"
6072"\n"
6073"Optional keyword arguments are fix_imports, encoding and errors,\n"
6074"which are used to control compatiblity support for pickle stream generated\n"
6075"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6076"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6077"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6078"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6079
6080static PyObject *
6081pickle_load(PyObject *self, PyObject *args, PyObject *kwds)
6082{
6083 static char *kwlist[] = {"file", "fix_imports", "encoding", "errors", 0};
6084 PyObject *file;
6085 PyObject *fix_imports = Py_True;
6086 PyObject *result;
6087 char *encoding = NULL;
6088 char *errors = NULL;
6089 UnpicklerObject *unpickler;
6090
6091 /* fix_imports, encoding and errors are a keyword-only argument. */
6092 if (Py_SIZE(args) != 1) {
6093 PyErr_Format(PyExc_TypeError,
6094 "pickle.load() takes exactly one positional "
6095 "argument (%zd given)", Py_SIZE(args));
6096 return NULL;
6097 }
6098
6099 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:load", kwlist,
6100 &file, &fix_imports, &encoding, &errors))
6101 return NULL;
6102
6103 unpickler = _Unpickler_New();
6104 if (unpickler == NULL)
6105 return NULL;
6106
6107 if (_Unpickler_SetInputStream(unpickler, file) < 0)
6108 goto error;
6109
6110 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6111 goto error;
6112
6113 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6114 if (unpickler->fix_imports == -1)
6115 goto error;
6116
6117 result = load(unpickler);
6118 Py_DECREF(unpickler);
6119 return result;
6120
6121 error:
6122 Py_XDECREF(unpickler);
6123 return NULL;
6124}
6125
6126PyDoc_STRVAR(pickle_loads_doc,
6127"loads(input, *, fix_imports=True, encoding='ASCII', errors='strict') -> object\n"
6128"\n"
6129"Read a pickled object hierarchy from a bytes object and return the\n"
6130"reconstituted object hierarchy specified therein\n"
6131"\n"
6132"The protocol version of the pickle is detected automatically, so no protocol\n"
6133"argument is needed. Bytes past the pickled object's representation are\n"
6134"ignored.\n"
6135"\n"
6136"Optional keyword arguments are fix_imports, encoding and errors, which\n"
6137"are used to control compatiblity support for pickle stream generated\n"
6138"by Python 2.x. If fix_imports is True, pickle will try to map the old\n"
6139"Python 2.x names to the new names used in Python 3.x. The encoding and\n"
6140"errors tell pickle how to decode 8-bit string instances pickled by Python\n"
6141"2.x; these default to 'ASCII' and 'strict', respectively.\n");
6142
6143static PyObject *
6144pickle_loads(PyObject *self, PyObject *args, PyObject *kwds)
6145{
6146 static char *kwlist[] = {"input", "fix_imports", "encoding", "errors", 0};
6147 PyObject *input;
6148 PyObject *fix_imports = Py_True;
6149 PyObject *result;
6150 char *encoding = NULL;
6151 char *errors = NULL;
6152 UnpicklerObject *unpickler;
6153
6154 /* fix_imports, encoding and errors are a keyword-only argument. */
6155 if (Py_SIZE(args) != 1) {
6156 PyErr_Format(PyExc_TypeError,
6157 "pickle.loads() takes exactly one positional "
6158 "argument (%zd given)", Py_SIZE(args));
6159 return NULL;
6160 }
6161
6162 if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|Oss:loads", kwlist,
6163 &input, &fix_imports, &encoding, &errors))
6164 return NULL;
6165
6166 unpickler = _Unpickler_New();
6167 if (unpickler == NULL)
6168 return NULL;
6169
6170 if (_Unpickler_SetStringInput(unpickler, input) < 0)
6171 goto error;
6172
6173 if (_Unpickler_SetInputEncoding(unpickler, encoding, errors) < 0)
6174 goto error;
6175
6176 unpickler->fix_imports = PyObject_IsTrue(fix_imports);
6177 if (unpickler->fix_imports == -1)
6178 goto error;
6179
6180 result = load(unpickler);
6181 Py_DECREF(unpickler);
6182 return result;
6183
6184 error:
6185 Py_XDECREF(unpickler);
6186 return NULL;
6187}
6188
6189
6190static struct PyMethodDef pickle_methods[] = {
6191 {"dump", (PyCFunction)pickle_dump, METH_VARARGS|METH_KEYWORDS,
6192 pickle_dump_doc},
6193 {"dumps", (PyCFunction)pickle_dumps, METH_VARARGS|METH_KEYWORDS,
6194 pickle_dumps_doc},
6195 {"load", (PyCFunction)pickle_load, METH_VARARGS|METH_KEYWORDS,
6196 pickle_load_doc},
6197 {"loads", (PyCFunction)pickle_loads, METH_VARARGS|METH_KEYWORDS,
6198 pickle_loads_doc},
6199 {NULL, NULL} /* sentinel */
6200};
6201
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006202static int
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006203initmodule(void)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006204{
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006205 PyObject *copyreg = NULL;
6206 PyObject *compat_pickle = NULL;
6207
6208 /* XXX: We should ensure that the types of the dictionaries imported are
6209 exactly PyDict objects. Otherwise, it is possible to crash the pickle
6210 since we use the PyDict API directly to access these dictionaries. */
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006211
6212 copyreg = PyImport_ImportModule("copyreg");
6213 if (!copyreg)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006214 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006215 dispatch_table = PyObject_GetAttrString(copyreg, "dispatch_table");
6216 if (!dispatch_table)
6217 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006218 extension_registry = \
6219 PyObject_GetAttrString(copyreg, "_extension_registry");
6220 if (!extension_registry)
6221 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006222 inverted_registry = PyObject_GetAttrString(copyreg, "_inverted_registry");
6223 if (!inverted_registry)
6224 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006225 extension_cache = PyObject_GetAttrString(copyreg, "_extension_cache");
6226 if (!extension_cache)
6227 goto error;
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006228 Py_CLEAR(copyreg);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006229
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006230 /* Load the 2.x -> 3.x stdlib module mapping tables */
6231 compat_pickle = PyImport_ImportModule("_compat_pickle");
6232 if (!compat_pickle)
6233 goto error;
6234 name_mapping_2to3 = PyObject_GetAttrString(compat_pickle, "NAME_MAPPING");
6235 if (!name_mapping_2to3)
6236 goto error;
6237 if (!PyDict_CheckExact(name_mapping_2to3)) {
6238 PyErr_Format(PyExc_RuntimeError,
6239 "_compat_pickle.NAME_MAPPING should be a dict, not %.200s",
6240 Py_TYPE(name_mapping_2to3)->tp_name);
6241 goto error;
6242 }
6243 import_mapping_2to3 = PyObject_GetAttrString(compat_pickle,
6244 "IMPORT_MAPPING");
6245 if (!import_mapping_2to3)
6246 goto error;
6247 if (!PyDict_CheckExact(import_mapping_2to3)) {
6248 PyErr_Format(PyExc_RuntimeError,
6249 "_compat_pickle.IMPORT_MAPPING should be a dict, "
6250 "not %.200s", Py_TYPE(import_mapping_2to3)->tp_name);
6251 goto error;
6252 }
6253 /* ... and the 3.x -> 2.x mapping tables */
6254 name_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6255 "REVERSE_NAME_MAPPING");
6256 if (!name_mapping_3to2)
6257 goto error;
6258 if (!PyDict_CheckExact(name_mapping_3to2)) {
6259 PyErr_Format(PyExc_RuntimeError,
6260 "_compat_pickle.REVERSE_NAME_MAPPING shouldbe a dict, "
6261 "not %.200s", Py_TYPE(name_mapping_3to2)->tp_name);
6262 goto error;
6263 }
6264 import_mapping_3to2 = PyObject_GetAttrString(compat_pickle,
6265 "REVERSE_IMPORT_MAPPING");
6266 if (!import_mapping_3to2)
6267 goto error;
6268 if (!PyDict_CheckExact(import_mapping_3to2)) {
6269 PyErr_Format(PyExc_RuntimeError,
6270 "_compat_pickle.REVERSE_IMPORT_MAPPING should be a dict, "
6271 "not %.200s", Py_TYPE(import_mapping_3to2)->tp_name);
6272 goto error;
6273 }
6274 Py_CLEAR(compat_pickle);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006275
6276 empty_tuple = PyTuple_New(0);
6277 if (empty_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006278 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006279 two_tuple = PyTuple_New(2);
6280 if (two_tuple == NULL)
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006281 goto error;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006282 /* We use this temp container with no regard to refcounts, or to
6283 * keeping containees alive. Exempt from GC, because we don't
6284 * want anything looking at two_tuple() by magic.
6285 */
6286 PyObject_GC_UnTrack(two_tuple);
6287
6288 return 0;
6289
6290 error:
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006291 Py_CLEAR(copyreg);
6292 Py_CLEAR(dispatch_table);
6293 Py_CLEAR(extension_registry);
6294 Py_CLEAR(inverted_registry);
6295 Py_CLEAR(extension_cache);
6296 Py_CLEAR(compat_pickle);
6297 Py_CLEAR(name_mapping_2to3);
6298 Py_CLEAR(import_mapping_2to3);
6299 Py_CLEAR(name_mapping_3to2);
6300 Py_CLEAR(import_mapping_3to2);
6301 Py_CLEAR(empty_tuple);
6302 Py_CLEAR(two_tuple);
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006303 return -1;
6304}
6305
6306static struct PyModuleDef _picklemodule = {
6307 PyModuleDef_HEAD_INIT,
6308 "_pickle",
6309 pickle_module_doc,
6310 -1,
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006311 pickle_methods,
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006312 NULL,
6313 NULL,
6314 NULL,
6315 NULL
6316};
6317
6318PyMODINIT_FUNC
6319PyInit__pickle(void)
6320{
6321 PyObject *m;
6322
6323 if (PyType_Ready(&Unpickler_Type) < 0)
6324 return NULL;
6325 if (PyType_Ready(&Pickler_Type) < 0)
6326 return NULL;
6327 if (PyType_Ready(&Pdata_Type) < 0)
6328 return NULL;
Antoine Pitrouea99c5c2010-09-09 18:33:21 +00006329 if (PyType_Ready(&PicklerMemoProxyType) < 0)
6330 return NULL;
6331 if (PyType_Ready(&UnpicklerMemoProxyType) < 0)
6332 return NULL;
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006333
6334 /* Create the module and add the functions. */
6335 m = PyModule_Create(&_picklemodule);
6336 if (m == NULL)
6337 return NULL;
6338
6339 if (PyModule_AddObject(m, "Pickler", (PyObject *)&Pickler_Type) < 0)
6340 return NULL;
6341 if (PyModule_AddObject(m, "Unpickler", (PyObject *)&Unpickler_Type) < 0)
6342 return NULL;
6343
6344 /* Initialize the exceptions. */
6345 PickleError = PyErr_NewException("_pickle.PickleError", NULL, NULL);
6346 if (PickleError == NULL)
6347 return NULL;
6348 PicklingError = \
6349 PyErr_NewException("_pickle.PicklingError", PickleError, NULL);
6350 if (PicklingError == NULL)
6351 return NULL;
6352 UnpicklingError = \
6353 PyErr_NewException("_pickle.UnpicklingError", PickleError, NULL);
6354 if (UnpicklingError == NULL)
6355 return NULL;
6356
6357 if (PyModule_AddObject(m, "PickleError", PickleError) < 0)
6358 return NULL;
6359 if (PyModule_AddObject(m, "PicklingError", PicklingError) < 0)
6360 return NULL;
6361 if (PyModule_AddObject(m, "UnpicklingError", UnpicklingError) < 0)
6362 return NULL;
6363
Antoine Pitroud9dfaa92009-06-04 20:32:06 +00006364 if (initmodule() < 0)
Alexandre Vassalottica2d6102008-06-12 18:26:05 +00006365 return NULL;
6366
6367 return m;
6368}