diff options
| -rw-r--r-- | Lib/pickle.py | 13 | ||||
| -rw-r--r-- | Modules/_pickle.c | 16 |
2 files changed, 21 insertions, 8 deletions
diff --git a/Lib/pickle.py b/Lib/pickle.py index 37c3d52..2e55c8a 100644 --- a/Lib/pickle.py +++ b/Lib/pickle.py @@ -345,6 +345,9 @@ class _Pickler: else: self.write(PERSID + str(pid).encode("ascii") + b'\n') + def _isiter(self, obj): + return hasattr(obj, '__next__') and hasattr(obj, '__iter__') + def save_reduce(self, func, args, state=None, listitems=None, dictitems=None, obj=None): # This API is called by some subclasses @@ -357,6 +360,16 @@ class _Pickler: if not hasattr(func, '__call__'): raise PicklingError("func from save_reduce() should be callable") + # Assert that listitems is an iterator + if listitems is not None and not self._isiter(listitems): + raise PicklingError("listitems from save_reduce() should be an " + "iterator") + + # Assert that dictitems is an iterator + if dictitems is not None and not self._isiter(dictitems): + raise PicklingError("dictitems from save_reduce() should be an " + "iterator") + save = self.save write = self.write diff --git a/Modules/_pickle.c b/Modules/_pickle.c index a689c33..b0686af 100644 --- a/Modules/_pickle.c +++ b/Modules/_pickle.c @@ -1963,7 +1963,6 @@ save_reduce(PicklerObject *self, PyObject *args, PyObject *obj) PyObject *state = NULL; PyObject *listitems = Py_None; PyObject *dictitems = Py_None; - Py_ssize_t size; int use_newobj = self->proto >= 2; @@ -1971,13 +1970,6 @@ save_reduce(PicklerObject *self, PyObject *args, PyObject *obj) const char build_op = BUILD; const char newobj_op = NEWOBJ; - size = PyTuple_Size(args); - if (size < 2 || size > 5) { - PyErr_SetString(PicklingError, "tuple returned by " - "__reduce__ must contain 2 through 5 elements"); - return -1; - } - if (!PyArg_UnpackTuple(args, "save_reduce", 2, 5, &callable, &argtup, &state, &listitems, &dictitems)) return -1; @@ -2154,6 +2146,7 @@ save(PicklerObject *self, PyObject *obj, int pers_save) PyObject *reduce_value = NULL; PyObject *memo_key = NULL; int status = 0; + Py_ssize_t size; if (Py_EnterRecursiveCall(" while pickling an object") < 0) return -1; @@ -2332,6 +2325,13 @@ save(PicklerObject *self, PyObject *obj, int pers_save) goto error; } + size = PyTuple_Size(reduce_value); + if (size < 2 || size > 5) { + PyErr_SetString(PicklingError, "tuple returned by " + "__reduce__ must contain 2 through 5 elements"); + goto error; + } + status = save_reduce(self, reduce_value, obj); if (0) { |