Issue #8104: socket.recv_into() and socket.recvfrom_into() now support

writing into objects supporting the new buffer API, for example bytearrays
or memoryviews.

diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py
index ee07f87260ab01f17a1912ca385666ba07c357c9..75bd25839ff07d7b9ee9c081bf0f8f8840a9c576 100644 (file)
--- 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

diff --git a/Misc/NEWS b/Misc/NEWS
index ffcdaa7943c024c04612b622d79cf81f53645464..abbcc041b5ee7da657d4e3d28b61fcbaa3ca2393 100644 (file)
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -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 a993e886b078149b668772bb50b19e1544b43111..fbd0239af791af8ea7d130dee9eaf42da85bd829 100644 (file)
--- 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,