mirror of
https://github.com/python/cpython.git
synced 2025-08-04 00:48:58 +00:00
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:
parent
9a4135e939
commit
289f1f80ee
6 changed files with 227 additions and 24 deletions
|
@ -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;
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue