bpo-35900: Add a state_setter arg to save_reduce (GH-12588)

Allow reduction methods to return a 6-item tuple where the 6th item specifies a
custom state-setting method that's called instead of the regular
``__setstate__`` method.
This commit is contained in:
Pierre Glaser 2019-05-08 21:40:25 +02:00 committed by Antoine Pitrou
parent 39889864c0
commit 65d98d0f53
5 changed files with 114 additions and 15 deletions

View File

@ -598,7 +598,7 @@ or both.
module; the pickle module searches the module namespace to determine the
object's module. This behaviour is typically useful for singletons.
When a tuple is returned, it must be between two and five items long.
When a tuple is returned, it must be between two and six items long.
Optional items can either be omitted, or ``None`` can be provided as their
value. The semantics of each item are in order:
@ -629,6 +629,15 @@ or both.
value``. This is primarily used for dictionary subclasses, but may be used
by other classes as long as they implement :meth:`__setitem__`.
* Optionally, a callable with a ``(obj, state)`` signature. This
callable allows the user to programatically control the state-updating
behavior of a specific object, instead of using ``obj``'s static
:meth:`__setstate__` method. If not ``None``, this callable will have
priority over ``obj``'s :meth:`__setstate__`.
.. versionadded:: 3.8
The optional sixth tuple item, ``(obj, state)``, was added.
.. method:: object.__reduce_ex__(protocol)

View File

@ -537,9 +537,9 @@ class _Pickler:
# Assert that it returned an appropriately sized tuple
l = len(rv)
if not (2 <= l <= 5):
if not (2 <= l <= 6):
raise PicklingError("Tuple returned by %s must have "
"two to five elements" % reduce)
"two to six elements" % reduce)
# Save the reduce() output and finally memoize the object
self.save_reduce(obj=obj, *rv)
@ -561,7 +561,7 @@ class _Pickler:
"persistent IDs in protocol 0 must be ASCII strings")
def save_reduce(self, func, args, state=None, listitems=None,
dictitems=None, obj=None):
dictitems=None, state_setter=None, obj=None):
# This API is called by some subclasses
if not isinstance(args, tuple):
@ -655,8 +655,25 @@ class _Pickler:
self._batch_setitems(dictitems)
if state is not None:
save(state)
write(BUILD)
if state_setter is None:
save(state)
write(BUILD)
else:
# If a state_setter is specified, call it instead of load_build
# to update obj's with its previous state.
# First, push state_setter and its tuple of expected arguments
# (obj, state) onto the stack.
save(state_setter)
save(obj) # simple BINGET opcode as obj is already memoized.
save(state)
write(TUPLE2)
# Trigger a state_setter(obj, state) function call.
write(REDUCE)
# The purpose of state_setter is to carry-out an
# inplace modification of obj. We do not care about what the
# method might return, so its output is eventually removed from
# the stack.
write(POP)
# Methods below this point are dispatched through the dispatch table

View File

@ -2992,7 +2992,26 @@ class AAA(object):
return str, (REDUCE_A,)
class BBB(object):
pass
def __init__(self):
# Add an instance attribute to enable state-saving routines at pickling
# time.
self.a = "some attribute"
def __setstate__(self, state):
self.a = "BBB.__setstate__"
def setstate_bbb(obj, state):
"""Custom state setter for BBB objects
Such callable may be created by other persons than the ones who created the
BBB class. If passed as the state_setter item of a custom reducer, this
allows for custom state setting behavior of BBB objects. One can think of
it as the analogous of list_setitems or dict_setitems but for foreign
classes/functions.
"""
obj.a = "custom state_setter"
class AbstractDispatchTableTests(unittest.TestCase):
@ -3081,6 +3100,25 @@ class AbstractDispatchTableTests(unittest.TestCase):
self.assertEqual(default_load_dump(a), REDUCE_A)
self.assertIsInstance(default_load_dump(b), BBB)
# End-to-end testing of save_reduce with the state_setter keyword
# argument. This is a dispatch_table test as the primary goal of
# state_setter is to tweak objects reduction behavior.
# In particular, state_setter is useful when the default __setstate__
# behavior is not flexible enough.
# No custom reducer for b has been registered for now, so
# BBB.__setstate__ should be used at unpickling time
self.assertEqual(default_load_dump(b).a, "BBB.__setstate__")
def reduce_bbb(obj):
return BBB, (), obj.__dict__, None, None, setstate_bbb
dispatch_table[BBB] = reduce_bbb
# The custom reducer reduce_bbb includes a state setter, that should
# have priority over BBB.__setstate__
self.assertEqual(custom_load_dump(b).a, "custom state_setter")
if __name__ == "__main__":
# Print some stuff that can be used to rewrite DATA{0,1,2}

View File

@ -0,0 +1,3 @@
Allow reduction methods to return a 6-item tuple where the 6th item specifies a
custom state-setting method that's called instead of the regular
``__setstate__`` method.

View File

@ -3662,6 +3662,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;
@ -3672,14 +3673,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");
"__reduce__ must contain 2 through 6 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)) {
@ -3714,6 +3716,15 @@ save_reduce(PicklerObject *self, PyObject *args, PyObject *obj)
return -1;
}
if (state_setter == Py_None)
state_setter = NULL;
else if (!PyCallable_Check(state_setter)) {
PyErr_Format(st->PicklingError, "sixth element of the tuple "
"returned by __reduce__ must be a function, not %s",
Py_TYPE(state_setter)->tp_name);
return -1;
}
if (self->proto >= 2) {
PyObject *name;
_Py_IDENTIFIER(__name__);
@ -3933,11 +3944,32 @@ 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 {
/* If a state_setter is specified, call it instead of load_build to
* update obj's with its previous state.
* The first 4 save/write instructions push state_setter and its
* tuple of expected arguments (obj, state) onto the stack. The
* REDUCE opcode triggers the state_setter(obj, state) function
* call. Finally, because state-updating routines only do in-place
* modification, the whole operation has to be stack-transparent.
* Thus, we finally pop the call's output from the stack.*/
const char tupletwo_op = TUPLE2;
const char pop_op = POP;
if (save(self, state_setter, 0) < 0 ||
save(self, obj, 0) < 0 || save(self, state, 0) < 0 ||
_Pickler_Write(self, &tupletwo_op, 1) < 0 ||
_Pickler_Write(self, &reduce_op, 1) < 0 ||
_Pickler_Write(self, &pop_op, 1) < 0)
return -1;
}
}
return 0;
}