diff options
| -rw-r--r-- | Lib/test/test_socket.py | 44 | ||||
| -rw-r--r-- | Misc/NEWS | 4 | ||||
| -rw-r--r-- | Modules/socketmodule.c | 46 |
3 files changed, 73 insertions, 21 deletions
diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py index ee07f87260..75bd25839f 100644 --- a/Lib/test/test_socket.py +++ b/Lib/test/test_socket.py @@ -1226,28 +1226,64 @@ class BufferIOTest(SocketConnectedTest): def __init__(self, methodName='runTest'): SocketConnectedTest.__init__(self, methodName=methodName) - def testRecvInto(self): + def testRecvIntoArray(self): buf = array.array('c', ' '*1024) nbytes = self.cli_conn.recv_into(buf) self.assertEqual(nbytes, len(MSG)) msg = buf.tostring()[:len(MSG)] self.assertEqual(msg, MSG) - def _testRecvInto(self): + def _testRecvIntoArray(self): buf = buffer(MSG) self.serv_conn.send(buf) - def testRecvFromInto(self): + def testRecvIntoBytearray(self): + buf = bytearray(1024) + nbytes = self.cli_conn.recv_into(buf) + self.assertEqual(nbytes, len(MSG)) + msg = buf[:len(MSG)] + self.assertEqual(msg, MSG) + + _testRecvIntoBytearray = _testRecvIntoArray + + def testRecvIntoMemoryview(self): + buf = bytearray(1024) + nbytes = self.cli_conn.recv_into(memoryview(buf)) + self.assertEqual(nbytes, len(MSG)) + msg = buf[:len(MSG)] + self.assertEqual(msg, MSG) + + _testRecvIntoMemoryview = _testRecvIntoArray + + def testRecvFromIntoArray(self): buf = array.array('c', ' '*1024) nbytes, addr = self.cli_conn.recvfrom_into(buf) self.assertEqual(nbytes, len(MSG)) msg = buf.tostring()[:len(MSG)] self.assertEqual(msg, MSG) - def _testRecvFromInto(self): + def _testRecvFromIntoArray(self): buf = buffer(MSG) self.serv_conn.send(buf) + def testRecvFromIntoBytearray(self): + buf = bytearray(1024) + nbytes, addr = self.cli_conn.recvfrom_into(buf) + self.assertEqual(nbytes, len(MSG)) + msg = buf[:len(MSG)] + self.assertEqual(msg, MSG) + + _testRecvFromIntoBytearray = _testRecvFromIntoArray + + def testRecvFromIntoMemoryview(self): + buf = bytearray(1024) + nbytes, addr = self.cli_conn.recvfrom_into(memoryview(buf)) + self.assertEqual(nbytes, len(MSG)) + msg = buf[:len(MSG)] + self.assertEqual(msg, MSG) + + _testRecvFromIntoMemoryview = _testRecvFromIntoArray + TIPC_STYPE = 2000 TIPC_LOWER = 200 @@ -22,6 +22,10 @@ Core and Builtins Library ------- +- Issue #8104: socket.recv_into() and socket.recvfrom_into() now support + writing into objects supporting the new buffer API, for example bytearrays + or memoryviews. + - Issue #4961: Inconsistent/wrong result of askyesno function in tkMessageBox with Tcl/Tk-8.5. diff --git a/Modules/socketmodule.c b/Modules/socketmodule.c index a993e886b0..fbd0239af7 100644 --- a/Modules/socketmodule.c +++ b/Modules/socketmodule.c @@ -2449,19 +2449,20 @@ sock_recv_into(PySocketSockObject *s, PyObject *args, PyObject *kwds) int recvlen = 0, flags = 0; ssize_t readlen; - char *buf; - int buflen; + Py_buffer buf; + Py_ssize_t buflen; /* Get the buffer's memory */ - if (!PyArg_ParseTupleAndKeywords(args, kwds, "w#|ii:recv_into", kwlist, - &buf, &buflen, &recvlen, &flags)) + if (!PyArg_ParseTupleAndKeywords(args, kwds, "w*|ii:recv_into", kwlist, + &buf, &recvlen, &flags)) return NULL; - assert(buf != 0 && buflen > 0); + buflen = buf.len; + assert(buf.buf != 0 && buflen > 0); if (recvlen < 0) { PyErr_SetString(PyExc_ValueError, "negative buffersize in recv_into"); - return NULL; + goto error; } if (recvlen == 0) { /* If nbytes was not specified, use the buffer's length */ @@ -2472,19 +2473,24 @@ sock_recv_into(PySocketSockObject *s, PyObject *args, PyObject *kwds) if (buflen < recvlen) { PyErr_SetString(PyExc_ValueError, "buffer too small for requested bytes"); - return NULL; + goto error; } /* Call the guts */ - readlen = sock_recv_guts(s, buf, recvlen, flags); + readlen = sock_recv_guts(s, buf.buf, recvlen, flags); if (readlen < 0) { /* Return an error. */ - return NULL; + goto error; } + PyBuffer_Release(&buf); /* Return the number of bytes read. Note that we do not do anything special here in the case that readlen < recvlen. */ return PyInt_FromSsize_t(readlen); + +error: + PyBuffer_Release(&buf); + return NULL; } PyDoc_STRVAR(recv_into_doc, @@ -2623,37 +2629,43 @@ sock_recvfrom_into(PySocketSockObject *s, PyObject *args, PyObject* kwds) int recvlen = 0, flags = 0; ssize_t readlen; - char *buf; + Py_buffer buf; int buflen; PyObject *addr = NULL; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "w#|ii:recvfrom_into", - kwlist, &buf, &buflen, + if (!PyArg_ParseTupleAndKeywords(args, kwds, "w*|ii:recvfrom_into", + kwlist, &buf, &recvlen, &flags)) return NULL; - assert(buf != 0 && buflen > 0); + buflen = buf.len; + assert(buf.buf != 0 && buflen > 0); if (recvlen < 0) { PyErr_SetString(PyExc_ValueError, "negative buffersize in recvfrom_into"); - return NULL; + goto error; } if (recvlen == 0) { /* If nbytes was not specified, use the buffer's length */ recvlen = buflen; } - readlen = sock_recvfrom_guts(s, buf, recvlen, flags, &addr); + readlen = sock_recvfrom_guts(s, buf.buf, recvlen, flags, &addr); if (readlen < 0) { /* Return an error */ - Py_XDECREF(addr); - return NULL; + goto error; } + PyBuffer_Release(&buf); /* Return the number of bytes read and the address. Note that we do not do anything special here in the case that readlen < recvlen. */ return Py_BuildValue("lN", readlen, addr); + +error: + Py_XDECREF(addr); + PyBuffer_Release(&buf); + return NULL; } PyDoc_STRVAR(recvfrom_into_doc, |