diff --git a/Lib/pickle.py b/Lib/pickle.py index d533e660af..849b54ad49 100644 --- a/Lib/pickle.py +++ b/Lib/pickle.py @@ -1547,6 +1547,11 @@ class _Unpickler: slotstate = None if isinstance(state, tuple) and len(state) == 2: state, slotstate = state + + elif isinstance(state, tuple) and len(state) == 3: + setstate, state, slostate = state + setstate(state, slotstate) + return if state: inst_dict = inst.__dict__ intern = sys.intern diff --git a/Lib/test/test_pickle.py b/Lib/test/test_pickle.py index b4bce7e6ac..2cbf52d866 100644 --- a/Lib/test/test_pickle.py +++ b/Lib/test/test_pickle.py @@ -258,7 +258,7 @@ if has_c_implementation: check_sizeof = support.check_sizeof def test_pickler(self): - basesize = support.calcobjsize('6P2n3i2n3iP') + basesize = support.calcobjsize('6P2n3i2n3i2P') p = _pickle.Pickler(io.BytesIO()) self.assertEqual(object.__sizeof__(p), basesize) MT_size = struct.calcsize('3nP0n') diff --git a/Modules/_pickle.c b/Modules/_pickle.c index 2b97294e1e..6978012aee 100644 --- a/Modules/_pickle.c +++ b/Modules/_pickle.c @@ -613,6 +613,9 @@ typedef struct PicklerObject { PyObject *pers_func_self; /* borrowed reference to self if pers_func is an unbound method, NULL otherwise */ PyObject *dispatch_table; /* private dispatch_table, can be NULL */ + PyObject *global_hook; /* hook for invoking user-defined callbacks + instead of save_global when pickling + functions and classes*/ PyObject *write; /* write() method of the output stream. */ PyObject *output_buffer; /* Write into a local bytearray buffer before @@ -1107,6 +1110,7 @@ _Pickler_New(void) self->fast_memo = NULL; self->max_output_len = WRITE_BUF_SIZE; self->output_len = 0; + self->global_hook = NULL; self->memo = PyMemoTable_New(); self->output_buffer = PyBytes_FromStringAndSize(NULL, @@ -3656,6 +3660,7 @@ save_reduce(PicklerObject *self, PyObject *args, PyObject *obj) PyObject *state = NULL; PyObject *listitems = Py_None; PyObject *dictitems = Py_None; + PyObject *state_setter = Py_None; PickleState *st = _Pickle_GetGlobalState(); Py_ssize_t size; int use_newobj = 0, use_newobj_ex = 0; @@ -3666,14 +3671,15 @@ save_reduce(PicklerObject *self, PyObject *args, PyObject *obj) const char newobj_ex_op = NEWOBJ_EX; size = PyTuple_Size(args); - if (size < 2 || size > 5) { + if (size < 2 || size > 6) { PyErr_SetString(st->PicklingError, "tuple returned by " "__reduce__ must contain 2 through 5 elements"); return -1; } - if (!PyArg_UnpackTuple(args, "save_reduce", 2, 5, - &callable, &argtup, &state, &listitems, &dictitems)) + if (!PyArg_UnpackTuple(args, "save_reduce", 2, 6, + &callable, &argtup, &state, &listitems, &dictitems, + &state_setter)) return -1; if (!PyCallable_Check(callable)) { @@ -3708,6 +3714,15 @@ save_reduce(PicklerObject *self, PyObject *args, PyObject *obj) return -1; } + if (state_setter == Py_None) + state_setter = NULL; + else if (!PyFunction_Check(state_setter)) { + PyErr_Format(st->PicklingError, "sixth element of the tuple " + "returned by __reduce__ must be an function, not %s", + Py_TYPE(state_setter)->tp_name); + return -1; + } + if (self->proto >= 2) { PyObject *name; _Py_IDENTIFIER(__name__); @@ -3927,11 +3942,48 @@ save_reduce(PicklerObject *self, PyObject *args, PyObject *obj) return -1; if (state) { - if (save(self, state, 0) < 0 || - _Pickler_Write(self, &build_op, 1) < 0) - return -1; - } + if (state_setter == NULL) { + if (save(self, state, 0) < 0 || + _Pickler_Write(self, &build_op, 1) < 0) + return -1; + } + else { + PyObject *statetup = NULL; + + /* If a state_setter is specified, and state is a dict, we could be + * tempted to save a (state_setter, state) as state, but this would + * collide with load_build's (state, slotstate) special handling. + * Therefore, create a new format for state saving: (state_setter, + * state, slotstate) + */ + if PyDict_Check(state) + statetup = Py_BuildValue("(OOO)", state_setter, state, + Py_None); + else if PyTuple_Check(state) { + if (PyTuple_GET_SIZE(state) == 2){ + statetup = Py_BuildValue("(OOO)", state_setter, + PyTuple_GetItem(state, 0), + PyTuple_GetItem(state, 1)); + } + } + + if (statetup == NULL) { + PyErr_SetString(st->PicklingError, + "state must be either a dict or a tuple of" + " length 2"); + return -1; + } + + if (save(self, statetup, 0) < 0 || + _Pickler_Write(self, &build_op, 1) < 0){ + Py_DECREF(statetup); + return -1; + } + Py_DECREF(statetup); + } + + } return 0; } @@ -4020,7 +4072,28 @@ save(PicklerObject *self, PyObject *obj, int pers_save) status = save_tuple(self, obj); goto done; } - else if (type == &PyType_Type) { + /* The switch-on-type statement ends here because the next three + * conditions are not exclusive anymore. If global_hook returns + * NotImplementedError, then we must fallback to save_type or save_global + * */ + if (self->global_hook != NULL){ + PyObject *reduce_value = NULL; + reduce_value = PyObject_CallFunctionObjArgs(self->global_hook, self, obj, + NULL); + if (reduce_value == NULL){ + goto error; + } + + if (reduce_value != PyExc_NotImplementedError){ + status = save_reduce(self, reduce_value, obj); + Py_DECREF(reduce_value); + if (status < 0) + goto error; + goto done; + } + } + + if (type == &PyType_Type) { status = save_type(self, obj); goto done; } @@ -4662,6 +4735,7 @@ static PyMemberDef Pickler_members[] = { {"bin", T_INT, offsetof(PicklerObject, bin)}, {"fast", T_INT, offsetof(PicklerObject, fast)}, {"dispatch_table", T_OBJECT_EX, offsetof(PicklerObject, dispatch_table)}, + {"global_hook", T_OBJECT_EX, offsetof(PicklerObject, global_hook)}, {NULL} }; @@ -6229,6 +6303,30 @@ load_build(UnpicklerObject *self) Py_INCREF(slotstate); Py_DECREF(tmp); } + /* proto 5 addition: state can embed a callable state setter */ + else if (PyTuple_Check(state) && PyTuple_GET_SIZE(state) == 3) { + /* borrowed reference*/ + PyObject *tmp = state; + + setstate = PyTuple_GET_ITEM(tmp, 0); + state = PyTuple_GET_ITEM(tmp, 1); + slotstate = PyTuple_GET_ITEM(tmp, 2); + + Py_INCREF(setstate); + Py_INCREF(state); + Py_INCREF(slotstate); + Py_DECREF(tmp); + /* call the setstate function */ + if (PyObject_CallFunctionObjArgs(setstate, inst, state, + slotstate, NULL) == NULL){ + return -1; + } + + Py_DECREF(setstate); + Py_DECREF(state); + Py_DECREF(slotstate); + return 0; + } else slotstate = NULL;