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 58cd09ca61..12bec64343 100644 --- a/Modules/_pickle.c +++ b/Modules/_pickle.c @@ -613,6 +613,8 @@ 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 */ PyObject *write; /* write() method of the output stream. */ PyObject *output_buffer; /* Write into a local bytearray buffer before @@ -1107,6 +1109,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, @@ -4020,7 +4023,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, obj, + NULL); + if (reduce_value == NULL){ + goto error; + } + + if (reduce_value != PyExc_NotImplementedError){ + status = save_reduce(self, reduce_value, NULL); + Py_DECREF(reduce_value); + if (status < 0) + goto error; + goto done; + } + } + + if (type == &PyType_Type) { status = save_type(self, obj); goto done; } @@ -4662,6 +4686,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} };