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:
parent
39889864c0
commit
65d98d0f53
|
@ -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)
|
||||
|
||||
|
|
|
@ -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
|
||||
|
||||
|
|
|
@ -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}
|
||||
|
|
|
@ -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.
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
Loading…
Reference in New Issue