bpo-35900: Enable custom reduction callback registration in _pickle (GH-12499)

Enable custom reduction callback registration for functions and classes in
_pickle.c, using the new Pickler's attribute ``reducer_override``.
This commit is contained in:
Pierre Glaser 2019-05-08 23:08:25 +02:00 committed by Antoine Pitrou
parent 9a4135e939
commit 289f1f80ee
6 changed files with 227 additions and 24 deletions

View file

@ -616,6 +616,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 *reducer_override; /* 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
@ -1110,6 +1113,7 @@ _Pickler_New(void)
self->fast_memo = NULL;
self->max_output_len = WRITE_BUF_SIZE;
self->output_len = 0;
self->reducer_override = NULL;
self->memo = PyMemoTable_New();
self->output_buffer = PyBytes_FromStringAndSize(NULL,
@ -2220,7 +2224,7 @@ save_bytes(PicklerObject *self, PyObject *obj)
Python 2 *and* the appropriate 'bytes' object when unpickled
using Python 3. Again this is a hack and we don't need to do this
with newer protocols. */
PyObject *reduce_value = NULL;
PyObject *reduce_value;
int status;
if (PyBytes_GET_SIZE(obj) == 0) {
@ -4058,7 +4062,25 @@ save(PicklerObject *self, PyObject *obj, int pers_save)
status = save_tuple(self, obj);
goto done;
}
else if (type == &PyType_Type) {
/* Now, check reducer_override. If it returns NotImplemented,
* fallback to save_type or save_global, and then perhaps to the
* regular reduction mechanism.
*/
if (self->reducer_override != NULL) {
reduce_value = PyObject_CallFunctionObjArgs(self->reducer_override,
obj, NULL);
if (reduce_value == NULL) {
goto error;
}
if (reduce_value != Py_NotImplemented) {
goto reduce;
}
Py_DECREF(reduce_value);
reduce_value = NULL;
}
if (type == &PyType_Type) {
status = save_type(self, obj);
goto done;
}
@ -4149,6 +4171,7 @@ save(PicklerObject *self, PyObject *obj, int pers_save)
if (reduce_value == NULL)
goto error;
reduce:
if (PyUnicode_Check(reduce_value)) {
status = save_global(self, obj, reduce_value);
goto done;
@ -4180,6 +4203,20 @@ static int
dump(PicklerObject *self, PyObject *obj)
{
const char stop_op = STOP;
PyObject *tmp;
_Py_IDENTIFIER(reducer_override);
if (_PyObject_LookupAttrId((PyObject *)self, &PyId_reducer_override,
&tmp) < 0) {
return -1;
}
/* Cache the reducer_override method, if it exists. */
if (tmp != NULL) {
Py_XSETREF(self->reducer_override, tmp);
}
else {
Py_CLEAR(self->reducer_override);
}
if (self->proto >= 2) {
char header[2];
@ -4304,6 +4341,7 @@ Pickler_dealloc(PicklerObject *self)
Py_XDECREF(self->pers_func);
Py_XDECREF(self->dispatch_table);
Py_XDECREF(self->fast_memo);
Py_XDECREF(self->reducer_override);
PyMemoTable_Del(self->memo);
@ -4317,6 +4355,7 @@ Pickler_traverse(PicklerObject *self, visitproc visit, void *arg)
Py_VISIT(self->pers_func);
Py_VISIT(self->dispatch_table);
Py_VISIT(self->fast_memo);
Py_VISIT(self->reducer_override);
return 0;
}
@ -4328,6 +4367,7 @@ Pickler_clear(PicklerObject *self)
Py_CLEAR(self->pers_func);
Py_CLEAR(self->dispatch_table);
Py_CLEAR(self->fast_memo);
Py_CLEAR(self->reducer_override);
if (self->memo != NULL) {
PyMemoTable *memo = self->memo;