SF patch 555085 (timeout socket implementation) by Michael Gilfix.

I've made considerable changes to Michael's code, specifically to use
the select() system call directly and to store the timeout as a C
double instead of a Python object; internally, -1.0 (or anything
negative) represents the None from the API.

I'm not 100% sure that all corner cases are covered correctly, so
please keep an eye on this.  Next I'm going to try it Windows before
Tim complains.

No way is this a bugfix candidate. :-)

3 files changed, 267 insertions, 69 deletions

diff --git a/Lib/socket.py b/Lib/socket.py
index 153f602..04e2e33 100644
--- a/Lib/socket.py
+++ b/Lib/socket.py
@@ -134,7 +134,8 @@ def getfqdn(name=''):
 _socketmethods = (
     'bind', 'connect', 'connect_ex', 'fileno', 'listen',
     'getpeername', 'getsockname', 'getsockopt', 'setsockopt',
-    'recv', 'recvfrom', 'send', 'sendall', 'sendto', 'setblocking', 'shutdown')
+    'recv', 'recvfrom', 'send', 'sendall', 'sendto', 'setblocking',
+    'settimeout', 'gettimeout', 'shutdown')
 
 class _socketobject:
 
@@ -168,94 +169,108 @@ class _socketobject:
 
 
 class _fileobject:
+    """Implements a file object on top of a regular socket object."""
 
-    def __init__(self, sock, mode, bufsize):
+    def __init__(self, sock, mode='rb', bufsize=8192):
         self._sock = sock
         self._mode = mode
-        if bufsize < 0:
+        if bufsize <= 0:
             bufsize = 512
-        self._rbufsize = max(1, bufsize)
+        self._rbufsize = bufsize
         self._wbufsize = bufsize
-        self._wbuf = self._rbuf = ""
+        self._rbuf = [ ]
+        self._wbuf = [ ]
 
     def close(self):
         try:
             if self._sock:
                 self.flush()
         finally:
-            self._sock = 0
+            self._sock = None
 
     def __del__(self):
         self.close()
 
     def flush(self):
         if self._wbuf:
-            self._sock.sendall(self._wbuf)
-            self._wbuf = ""
+            buffer = ''.join(self._wbuf)
+            self._sock.sendall(buffer)
+            self._wbuf = [ ]
 
-    def fileno(self):
+    def fileno (self):
         return self._sock.fileno()
 
     def write(self, data):
-        self._wbuf = self._wbuf + data
+        self._wbuf.append (data)
+        # A _wbufsize of 1 means we're doing unbuffered IO.
+        # Flush accordingly.
         if self._wbufsize == 1:
             if '\n' in data:
-                self.flush()
-        else:
-            if len(self._wbuf) >= self._wbufsize:
-                self.flush()
+                self.flush ()
+        elif self.__get_wbuf_len() >= self._wbufsize:
+            self.flush()
 
     def writelines(self, list):
         filter(self._sock.sendall, list)
         self.flush()
 
-    def read(self, n=-1):
-        if n >= 0:
-            k = len(self._rbuf)
-            if n <= k:
-                data = self._rbuf[:n]
-                self._rbuf = self._rbuf[n:]
-                return data
-            n = n - k
-            L = [self._rbuf]
-            self._rbuf = ""
-            while n > 0:
-                new = self._sock.recv(max(n, self._rbufsize))
-                if not new: break
-                k = len(new)
-                if k > n:
-                    L.append(new[:n])
-                    self._rbuf = new[n:]
-                    break
-                L.append(new)
-                n = n - k
-            return "".join(L)
-        k = max(512, self._rbufsize)
-        L = [self._rbuf]
-        self._rbuf = ""
-        while 1:
-            new = self._sock.recv(k)
-            if not new: break
-            L.append(new)
-            k = min(k*2, 1024**2)
-        return "".join(L)
-
-    def readline(self, limit=-1):
-        data = ""
-        i = self._rbuf.find('\n')
-        while i < 0 and not (0 < limit <= len(self._rbuf)):
-            new = self._sock.recv(self._rbufsize)
-            if not new: break
-            i = new.find('\n')
-            if i >= 0: i = i + len(self._rbuf)
-            self._rbuf = self._rbuf + new
-        if i < 0: i = len(self._rbuf)
-        else: i = i+1
-        if 0 <= limit < len(self._rbuf): i = limit
-        data, self._rbuf = self._rbuf[:i], self._rbuf[i:]
+    def __get_wbuf_len (self):
+        buf_len = 0
+        for i in [len(x) for x in self._wbuf]:
+            buf_len += i
+        return buf_len
+
+    def __get_rbuf_len(self):
+        buf_len = 0
+        for i in [len(x) for x in self._rbuf]:
+            buf_len += i
+        return buf_len
+
+    def read(self, size=-1):
+        buf_len = self.__get_rbuf_len()
+        while size < 0 or buf_len < size:
+            recv_size = max(self._rbufsize, size - buf_len)
+            data = self._sock.recv(recv_size)
+            if not data:
+                break
+            buf_len += len(data)
+            self._rbuf.append(data)
+        data = ''.join(self._rbuf)
+        # Clear the rbuf at the end so we're not affected by
+        # an exception during a recv
+        self._rbuf = [ ]
+        if buf_len > size and size >= 0:
+           self._rbuf.append(data[size:])
+           data = data[:size]
+        return data
+
+    def readline(self, size=-1):
+        index = -1
+        buf_len = self.__get_rbuf_len()
+        if len (self._rbuf):
+            index = min([x.find('\n') for x in self._rbuf])
+        while index < 0 and (size < 0 or buf_len < size):
+            recv_size = max(self._rbufsize, size - buf_len)
+            data = self._sock.recv(recv_size)
+            if not data:
+                break
+            buf_len += len(data)
+            self._rbuf.append(data)
+            index = data.find('\n')
+        data = ''.join(self._rbuf)
+        self._rbuf = [ ]
+        index = data.find('\n')
+        if index >= 0:
+            index += 1
+        elif buf_len > size:
+            index = size
+        else:
+            index = buf_len
+        self._rbuf.append(data[index:])
+        data = data[:index]
         return data
 
-    def readlines(self, sizehint = 0):
+    def readlines(self, sizehint=0):
         total = 0
         list = []
         while 1:
diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py
index 9cef316..56c7fb1 100644
--- a/Lib/test/test_socket.py
+++ b/Lib/test/test_socket.py
@@ -109,6 +109,7 @@ except socket.error:
 canfork = hasattr(os, 'fork')
 try:
     PORT = 50007
+    msg = 'socket test\n'
     if not canfork or os.fork():
         # parent is server
         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -133,13 +134,52 @@ try:
             f = conn.makefile()
             if verbose:
                 print 'file obj:', f
+            data = conn.recv(1024)
+            if verbose:
+                print 'received:', data
+            conn.sendall(data)
+
+            # Perform a few tests on the windows file object
+            if verbose:
+                print "Staring _fileobject tests..."
+            f = socket._fileobject (conn, 'rb', 8192)
+            first_seg = f.read(7)
+            second_seg = f.read(5)
+            if not first_seg == 'socket ' or not second_seg == 'test\n':
+                print "Error performing read with the python _fileobject class"
+                os._exit (1)
+            elif verbose:
+                print "_fileobject buffered read works"
+            f.write (data)
+            f.flush ()
+
+            buf = ''
             while 1:
-                data = conn.recv(1024)
-                if not data:
+                char = f.read(1)
+                if not char:
+                    print "Error performing unbuffered read with the python ", \
+                          "_fileobject class"
+                    os._exit (1)
+                buf += char
+                if buf == msg:
+                    if verbose:
+                        print "__fileobject unbuffered read works"
                     break
-                if verbose:
-                    print 'received:', data
-                conn.sendall(data)
+            if verbose:
+                # If we got this far, write() must work as well
+                print "__fileobject write works"
+            f.write(buf)
+            f.flush()
+
+            line = f.readline()
+            if not line == msg:
+                print "Error perferming readline with the python _fileobject class"
+                os._exit (1)
+            f.write(line)
+            f.flush()
+            if verbose:
+                print "__fileobject readline works"
+
             conn.close()
     else:
         try:
@@ -149,11 +189,18 @@ try:
             if verbose:
                 print 'child connecting'
             s.connect(("127.0.0.1", PORT))
-            msg = 'socket test'
-            s.send(msg)
-            data = s.recv(1024)
-            if msg != data:
-                print 'parent/client mismatch'
+
+            iteration = 0
+            while 1:
+                s.send(msg)
+                data = s.recv(12)
+                if not data:
+                    break
+                if msg != data:
+                    print "parent/client mismatch. Failed in %s iteration. Received: [%s]" \
+                          %(iteration, data)
+                time.sleep (1)
+                iteration += 1
             s.close()
         finally:
             os._exit(1)
diff --git a/Lib/test/test_timeout.py b/Lib/test/test_timeout.py
new file mode 100644
index 0000000..c71efa7
--- /dev/null
+++ b/Lib/test/test_timeout.py
@@ -0,0 +1,136 @@
+#!/home/bernie/src/python23/dist/src/python
+
+import unittest
+
+import time
+import socket
+
+class creationTestCase(unittest.TestCase):
+    """Test Case for socket.gettimeout() and socket.settimeout()"""
+    def setUp(self):
+        self.__s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
+    def tearDown(self):
+        self.__s.close()
+
+    def testObjectCreation(self):
+        "Test Socket creation"
+        self.assertEqual(self.__s.gettimeout(), None, 
+            "Timeout socket not default to disable (None)")
+
+    def testFloatReturnValue(self):
+        "Test return value of getter/setter"
+        self.__s.settimeout(7.345)
+        self.assertEqual(self.__s.gettimeout(), 7.345, 
+            "settimeout() and gettimeout() return different result")
+
+        self.__s.settimeout(3)
+        self.assertEqual(self.__s.gettimeout(), 3, 
+            "settimeout() and gettimeout() return different result")
+
+    def testReturnType(self):
+        "Test return type of getter/setter"
+        self.__s.settimeout(1)
+        self.assertEqual(type(self.__s.gettimeout()), type(1.0),
+            "return type of gettimeout() is not FloatType")
+
+        self.__s.settimeout(3.9)
+        self.assertEqual(type(self.__s.gettimeout()), type(1.0),
+            "return type of gettimeout() is not FloatType")
+
+
+class timeoutTestCase(unittest.TestCase):
+    """Test Case for socket.socket() timeout functions"""
+    def setUp(self):
+        self.__s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.__addr_remote = ('www.google.com', 80)
+        self.__addr_local  = ('127.0.0.1', 25339)
+
+    def tearDown(self):
+        self.__s.close()
+
+    def testConnectTimeout(self):
+        "Test connect() timeout"
+        _timeout = 0.02
+        self.__s.settimeout(_timeout)
+
+        _t1 = time.time()
+        self.failUnlessRaises(socket.error, self.__s.connect,
+                self.__addr_remote)
+        _t2 = time.time()
+
+        _delta = abs(_t1 - _t2)
+        self.assert_(_delta < _timeout + 0.5,
+                "timeout (%f) is 0.5 seconds more than required (%f)"
+                %(_delta, _timeout))
+
+    def testRecvTimeout(self):
+        "Test recv() timeout"
+        _timeout = 0.02
+        self.__s.connect(self.__addr_remote)
+        self.__s.settimeout(_timeout)
+
+        _t1 = time.time()
+        self.failUnlessRaises(socket.error, self.__s.recv, 1024)
+        _t2 = time.time()
+
+        _delta = abs(_t1 - _t2)
+        self.assert_(_delta < _timeout + 0.5,
+                "timeout (%f) is 0.5 seconds more than required (%f)" 
+                %(_delta, _timeout))
+
+    def testAcceptTimeout(self):
+        "Test accept() timeout()"
+        _timeout = 2
+        self.__s.settimeout(_timeout)
+        self.__s.bind(self.__addr_local)
+        self.__s.listen(5)
+
+        _t1 = time.time()
+        self.failUnlessRaises(socket.error, self.__s.accept)
+        _t2 = time.time()
+
+        _delta = abs(_t1 - _t2)
+        self.assert_(_delta < _timeout + 0.5,
+                "timeout (%f) is 0.5 seconds more than required (%f)" 
+                %(_delta, _timeout))
+
+    def testRecvfromTimeout(self):
+        "Test recvfrom() timeout()"
+        _timeout = 2
+        self.__s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.__s.settimeout(_timeout)
+        self.__s.bind(self.__addr_local)
+
+        _t1 = time.time()
+        self.failUnlessRaises(socket.error, self.__s.recvfrom, 8192)
+        _t2 = time.time()
+
+        _delta = abs(_t1 - _t2)
+        self.assert_(_delta < _timeout + 0.5,
+                "timeout (%f) is 0.5 seconds more than required (%f)" 
+                %(_delta, _timeout))
+
+    def testSend(self):
+        "Test send() timeout"
+        # couldn't figure out how to test it
+        pass
+
+    def testSendto(self):
+        "Test sendto() timeout"
+        # couldn't figure out how to test it
+        pass
+
+    def testSendall(self):
+        "Test sendall() timeout"
+        # couldn't figure out how to test it
+        pass
+
+
+def suite():
+    suite = unittest.TestSuite()
+
+    return suite
+
+if __name__ == "__main__":
+    unittest.main()