Support for buffer protocol for socket and struct.

* Added socket.recv_buf() and socket.recvfrom_buf() methods, that use the buffer
  protocol (send and sendto already did).

* Added struct.pack_to(), that is the corresponding buffer compatible method to
  unpack_from().

* Fixed minor typos in arraymodule.

3 files changed, 97 insertions, 27 deletions

diff --git a/Lib/socket.py b/Lib/socket.py
index 32a92b4535..cc5e65e166 100644
--- a/Lib/socket.py
+++ b/Lib/socket.py
@@ -140,7 +140,9 @@ class _socketobject(object):
 
     __doc__ = _realsocket.__doc__
 
-    __slots__ = ["_sock", "send", "recv", "sendto", "recvfrom",
+    __slots__ = ["_sock",
+                 "recv", "recv_buf", "recvfrom_buf",
+                 "send", "sendto", "recvfrom",
                  "__weakref__"]
 
     def __init__(self, family=AF_INET, type=SOCK_STREAM, proto=0, _sock=None):
@@ -149,8 +151,10 @@ class _socketobject(object):
         self._sock = _sock
         self.send = self._sock.send
         self.recv = self._sock.recv
+        self.recv_buf = self._sock.recv_buf
         self.sendto = self._sock.sendto
         self.recvfrom = self._sock.recvfrom
+        self.recvfrom_buf = self._sock.recvfrom_buf
 
     def close(self):
         self._sock = _closedsocket()
diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py
index 6943080a71..2246fb6e4c 100644
--- a/Lib/test/test_socket.py
+++ b/Lib/test/test_socket.py
@@ -9,6 +9,7 @@ import time
 import thread, threading
 import Queue
 import sys
+import array
 from weakref import proxy
 
 PORT = 50007
@@ -852,8 +853,38 @@ class TestLinuxAbstractNamespace(unittest.TestCase):
         self.assertRaises(socket.error, s.bind, address)
 
 
+class BufferIOTest(SocketConnectedTest):
+    """
+    Test the buffer versions of socket.recv() and socket.send().
+    """
+    def __init__(self, methodName='runTest'):
+        SocketConnectedTest.__init__(self, methodName=methodName)
+
+    def testRecvBuf(self):
+        buf = array.array('c', ' '*1024)
+        nbytes = self.cli_conn.recv_buf(buf)
+        self.assertEqual(nbytes, len(MSG))
+        msg = buf.tostring()[:len(MSG)]
+        self.assertEqual(msg, MSG)
+
+    def _testRecvBuf(self):
+        buf = buffer(MSG)
+        self.serv_conn.send(buf)
+
+    def testRecvFromBuf(self):
+        buf = array.array('c', ' '*1024)
+        nbytes, addr = self.cli_conn.recvfrom_buf(buf)
+        self.assertEqual(nbytes, len(MSG))
+        msg = buf.tostring()[:len(MSG)]
+        self.assertEqual(msg, MSG)
+
+    def _testRecvFromBuf(self):
+        buf = buffer(MSG)
+        self.serv_conn.send(buf)
+
 def test_main():
-    tests = [GeneralModuleTests, BasicTCPTest, TCPTimeoutTest, TestExceptions]
+    tests = [GeneralModuleTests, BasicTCPTest, TCPTimeoutTest, TestExceptions,
+             BufferIOTest]
     if sys.platform != 'mac':
         tests.extend([ BasicUDPTest, UDPTimeoutTest ])
 
diff --git a/Lib/test/test_struct.py b/Lib/test/test_struct.py
index 40fbde1129..1420a08ce2 100644
--- a/Lib/test/test_struct.py
+++ b/Lib/test/test_struct.py
@@ -1,5 +1,8 @@
 from test.test_support import TestFailed, verbose, verify
+import test.test_support
 import struct
+import array
+import unittest
 
 import sys
 ISBIGENDIAN = sys.byteorder == "big"
@@ -438,31 +441,6 @@ def test_705836():
 
 test_705836()
 
-def test_unpack_from():
-    test_string = 'abcd01234'
-    fmt = '4s'
-    s = struct.Struct(fmt)
-    for cls in (str, buffer):
-        data = cls(test_string)
-        assert s.unpack_from(data) == ('abcd',)
-        assert s.unpack_from(data, 2) == ('cd01',)
-        assert s.unpack_from(data, 4) == ('0123',)
-        for i in xrange(6):
-            assert s.unpack_from(data, i) == (data[i:i+4],)
-        for i in xrange(6, len(test_string) + 1):
-            simple_err(s.unpack_from, data, i)
-    for cls in (str, buffer):
-        data = cls(test_string)
-        assert struct.unpack_from(fmt, data) == ('abcd',)
-        assert struct.unpack_from(fmt, data, 2) == ('cd01',)
-        assert struct.unpack_from(fmt, data, 4) == ('0123',)
-        for i in xrange(6):
-            assert struct.unpack_from(fmt, data, i) == (data[i:i+4],)
-        for i in xrange(6, len(test_string) + 1):
-            simple_err(struct.unpack_from, fmt, data, i)
-
-test_unpack_from()
-
 def test_1229380():
     for endian in ('', '>', '<'):
         for cls in (int, long):
@@ -478,3 +456,60 @@ def test_1229380():
 if 0:
     # TODO: bug #1229380
     test_1229380()
+
+class PackBufferTestCase(unittest.TestCase):
+    """
+    Test the packing methods that work on buffers.
+    """
+
+    def test_unpack_from( self ):
+        test_string = 'abcd01234'
+        fmt = '4s'
+        s = struct.Struct(fmt)
+        for cls in (str, buffer):
+            data = cls(test_string)
+            self.assertEquals(s.unpack_from(data), ('abcd',))
+            self.assertEquals(s.unpack_from(data, 2), ('cd01',))
+            self.assertEquals(s.unpack_from(data, 4), ('0123',))
+            for i in xrange(6):
+                self.assertEquals(s.unpack_from(data, i), (data[i:i+4],))
+            for i in xrange(6, len(test_string) + 1):
+                simple_err(s.unpack_from, data, i)
+        for cls in (str, buffer):
+            data = cls(test_string)
+            self.assertEquals(struct.unpack_from(fmt, data), ('abcd',))
+            self.assertEquals(struct.unpack_from(fmt, data, 2), ('cd01',))
+            self.assertEquals(struct.unpack_from(fmt, data, 4), ('0123',))
+            for i in xrange(6):
+                self.assertEquals(struct.unpack_from(fmt, data, i),
+                                  (data[i:i+4],))
+            for i in xrange(6, len(test_string) + 1):
+                simple_err(struct.unpack_from, fmt, data, i)
+
+    def test_pack_to( self ):
+        test_string = 'Reykjavik rocks, eow!'
+        writable_buf = array.array('c', ' '*100)
+        fmt = '21s'
+        s = struct.Struct(fmt)
+
+        # Test without offset
+        s.pack_to(writable_buf, 0, test_string)
+        from_buf = writable_buf.tostring()[:len(test_string)]
+        self.assertEquals(from_buf, test_string)
+
+        # Test with offset.
+        s.pack_to(writable_buf, 10, test_string)
+        from_buf = writable_buf.tostring()[:len(test_string)+10]
+        self.assertEquals(from_buf, (test_string[:10] + test_string))
+
+        # Go beyond boundaries.
+        small_buf = array.array('c', ' '*10)
+        self.assertRaises(struct.error, s.pack_to, small_buf, 0, test_string)
+        self.assertRaises(struct.error, s.pack_to, small_buf, 2, test_string)
+
+def test_main():
+    test.test_support.run_unittest(PackBufferTestCase)
+    
+if __name__ == "__main__":
+    test_main()
+