Issue #6218: Make io.BytesIO and io.StringIO picklable.
diff --git a/Modules/_io/bytesio.c b/Modules/_io/bytesio.c
index daef78f..7b56551 100644
--- a/Modules/_io/bytesio.c
+++ b/Modules/_io/bytesio.c
@@ -606,6 +606,120 @@
Py_RETURN_NONE;
}
+/* Pickling support.
+
+ Note that only pickle protocol 2 and onward are supported since we use
+ extended __reduce__ API of PEP 307 to make BytesIO instances picklable.
+
+ Providing support for protocol < 2 would require the __reduce_ex__ method
+ which is notably long-winded when defined properly.
+
+ For BytesIO, the implementation would similar to one coded for
+ object.__reduce_ex__, but slightly less general. To be more specific, we
+ could call bytesio_getstate directly and avoid checking for the presence of
+ a fallback __reduce__ method. However, we would still need a __newobj__
+ function to use the efficient instance representation of PEP 307.
+ */
+
+static PyObject *
+bytesio_getstate(bytesio *self)
+{
+ PyObject *initvalue = bytesio_getvalue(self);
+ PyObject *dict;
+ PyObject *state;
+
+ if (initvalue == NULL)
+ return NULL;
+ if (self->dict == NULL) {
+ Py_INCREF(Py_None);
+ dict = Py_None;
+ }
+ else {
+ dict = PyDict_Copy(self->dict);
+ if (dict == NULL)
+ return NULL;
+ }
+
+ state = Py_BuildValue("(OnN)", initvalue, self->pos, dict);
+ Py_DECREF(initvalue);
+ return state;
+}
+
+static PyObject *
+bytesio_setstate(bytesio *self, PyObject *state)
+{
+ PyObject *result;
+ PyObject *position_obj;
+ PyObject *dict;
+ Py_ssize_t pos;
+
+ assert(state != NULL);
+
+ /* We allow the state tuple to be longer than 3, because we may need
+ someday to extend the object's state without breaking
+ backward-compatibility. */
+ if (!PyTuple_Check(state) || Py_SIZE(state) < 3) {
+ PyErr_Format(PyExc_TypeError,
+ "%.200s.__setstate__ argument should be 3-tuple, got %.200s",
+ Py_TYPE(self)->tp_name, Py_TYPE(state)->tp_name);
+ return NULL;
+ }
+ /* Reset the object to its default state. This is only needed to handle
+ the case of repeated calls to __setstate__. */
+ self->string_size = 0;
+ self->pos = 0;
+
+ /* Set the value of the internal buffer. If state[0] does not support the
+ buffer protocol, bytesio_write will raise the appropriate TypeError. */
+ result = bytesio_write(self, PyTuple_GET_ITEM(state, 0));
+ if (result == NULL)
+ return NULL;
+ Py_DECREF(result);
+
+ /* Set carefully the position value. Alternatively, we could use the seek
+ method instead of modifying self->pos directly to better protect the
+ object internal state against errneous (or malicious) inputs. */
+ position_obj = PyTuple_GET_ITEM(state, 1);
+ if (!PyLong_Check(position_obj)) {
+ PyErr_Format(PyExc_TypeError,
+ "second item of state must be an integer, not %.200s",
+ Py_TYPE(position_obj)->tp_name);
+ return NULL;
+ }
+ pos = PyLong_AsSsize_t(position_obj);
+ if (pos == -1 && PyErr_Occurred())
+ return NULL;
+ if (pos < 0) {
+ PyErr_SetString(PyExc_ValueError,
+ "position value cannot be negative");
+ return NULL;
+ }
+ self->pos = pos;
+
+ /* Set the dictionary of the instance variables. */
+ dict = PyTuple_GET_ITEM(state, 2);
+ if (dict != Py_None) {
+ if (!PyDict_Check(dict)) {
+ PyErr_Format(PyExc_TypeError,
+ "third item of state should be a dict, got a %.200s",
+ Py_TYPE(dict)->tp_name);
+ return NULL;
+ }
+ if (self->dict) {
+ /* Alternatively, we could replace the internal dictionary
+ completely. However, it seems more practical to just update it. */
+ if (PyDict_Update(self->dict, dict) < 0)
+ return NULL;
+ }
+ else {
+ Py_INCREF(dict);
+ self->dict = dict;
+ }
+ }
+
+ Py_RETURN_NONE;
+}
+
static void
bytesio_dealloc(bytesio *self)
{
@@ -630,9 +744,9 @@
if (self == NULL)
return NULL;
- self->string_size = 0;
- self->pos = 0;
- self->buf_size = 0;
+ /* tp_alloc initializes all the fields to zero. So we don't have to
+ initialize them here. */
+
self->buf = (char *)PyMem_Malloc(0);
if (self->buf == NULL) {
Py_DECREF(self);
@@ -705,6 +819,8 @@
{"getvalue", (PyCFunction)bytesio_getvalue, METH_VARARGS, getval_doc},
{"seek", (PyCFunction)bytesio_seek, METH_VARARGS, seek_doc},
{"truncate", (PyCFunction)bytesio_truncate, METH_VARARGS, truncate_doc},
+ {"__getstate__", (PyCFunction)bytesio_getstate, METH_NOARGS, NULL},
+ {"__setstate__", (PyCFunction)bytesio_setstate, METH_O, NULL},
{NULL, NULL} /* sentinel */
};