00001 """Based on code from timeout_socket.py, with some tweaks for compatibility.
00002 These tweaks should really be rolled back into timeout_socket, but it's
00003 not totally clear who is maintaining it at this point. In the meantime,
00004 we'll use a different module name for our tweaked version to avoid any
00005 confusion.
00006
00007 The original timeout_socket is by:
00008
00009 Scott Cotton <scott@chronis.pobox.com>
00010 Lloyd Zusman <ljz@asfast.com>
00011 Phil Mayes <pmayes@olivebr.com>
00012 Piers Lauder <piers@cs.su.oz.au>
00013 Radovan Garabik <garabik@melkor.dnp.fmph.uniba.sk>
00014 """
00015
00016 ident = "$Id: TimeoutSocket.py 237 2003-05-20 21:10:14Z warnes $"
00017
00018 import string, socket, select, errno
00019
00020 WSAEINVAL = getattr(errno, 'WSAEINVAL', 10022)
00021
00022
00023 class TimeoutSocket:
00024 """A socket imposter that supports timeout limits."""
00025
00026 def __init__(self, timeout=20, sock=None):
00027 self.timeout = float(timeout)
00028 self.inbuf = ''
00029 if sock is None:
00030 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
00031 self.sock = sock
00032 self.sock.setblocking(0)
00033 self._rbuf = ''
00034 self._wbuf = ''
00035
00036 def __getattr__(self, name):
00037
00038 return getattr(self.sock, name)
00039
00040 def connect(self, *addr):
00041 timeout = self.timeout
00042 sock = self.sock
00043 try:
00044
00045 sock.setblocking(0)
00046 apply(sock.connect, addr)
00047 sock.setblocking(timeout != 0)
00048 return 1
00049 except socket.error,why:
00050 if not timeout:
00051 raise
00052 sock.setblocking(1)
00053 if len(why.args) == 1:
00054 code = 0
00055 else:
00056 code, why = why
00057 if code not in (
00058 errno.EINPROGRESS, errno.EALREADY, errno.EWOULDBLOCK
00059 ):
00060 raise
00061 r,w,e = select.select([],[sock],[],timeout)
00062 if w:
00063 try:
00064 apply(sock.connect, addr)
00065 return 1
00066 except socket.error,why:
00067 if len(why.args) == 1:
00068 code = 0
00069 else:
00070 code, why = why
00071 if code in (errno.EISCONN, WSAEINVAL):
00072 return 1
00073 raise
00074 raise TimeoutError('socket connect() timeout.')
00075
00076 def send(self, data, flags=0):
00077 total = len(data)
00078 next = 0
00079 while 1:
00080 r, w, e = select.select([],[self.sock], [], self.timeout)
00081 if w:
00082 buff = data[next:next + 8192]
00083 sent = self.sock.send(buff, flags)
00084 next = next + sent
00085 if next == total:
00086 return total
00087 continue
00088 raise TimeoutError('socket send() timeout.')
00089
00090 def recv(self, amt, flags=0):
00091 if select.select([self.sock], [], [], self.timeout)[0]:
00092 return self.sock.recv(amt, flags)
00093 raise TimeoutError('socket recv() timeout.')
00094
00095 buffsize = 4096
00096 handles = 1
00097
00098 def makefile(self, mode="r", buffsize=-1):
00099 self.handles = self.handles + 1
00100 self.mode = mode
00101 return self
00102
00103 def close(self):
00104 self.handles = self.handles - 1
00105 if self.handles == 0 and self.sock.fileno() >= 0:
00106 self.sock.close()
00107
00108 def read(self, n=-1):
00109 if not isinstance(n, type(1)):
00110 n = -1
00111 if n >= 0:
00112 k = len(self._rbuf)
00113 if n <= k:
00114 data = self._rbuf[:n]
00115 self._rbuf = self._rbuf[n:]
00116 return data
00117 n = n - k
00118 L = [self._rbuf]
00119 self._rbuf = ""
00120 while n > 0:
00121 new = self.recv(max(n, self.buffsize))
00122 if not new: break
00123 k = len(new)
00124 if k > n:
00125 L.append(new[:n])
00126 self._rbuf = new[n:]
00127 break
00128 L.append(new)
00129 n = n - k
00130 return "".join(L)
00131 k = max(4096, self.buffsize)
00132 L = [self._rbuf]
00133 self._rbuf = ""
00134 while 1:
00135 new = self.recv(k)
00136 if not new: break
00137 L.append(new)
00138 k = min(k*2, 1024**2)
00139 return "".join(L)
00140
00141 def readline(self, limit=-1):
00142 data = ""
00143 i = self._rbuf.find('\n')
00144 while i < 0 and not (0 < limit <= len(self._rbuf)):
00145 new = self.recv(self.buffsize)
00146 if not new: break
00147 i = new.find('\n')
00148 if i >= 0: i = i + len(self._rbuf)
00149 self._rbuf = self._rbuf + new
00150 if i < 0: i = len(self._rbuf)
00151 else: i = i+1
00152 if 0 <= limit < len(self._rbuf): i = limit
00153 data, self._rbuf = self._rbuf[:i], self._rbuf[i:]
00154 return data
00155
00156 def readlines(self, sizehint = 0):
00157 total = 0
00158 list = []
00159 while 1:
00160 line = self.readline()
00161 if not line: break
00162 list.append(line)
00163 total += len(line)
00164 if sizehint and total >= sizehint:
00165 break
00166 return list
00167
00168 def writelines(self, list):
00169 self.send(''.join(list))
00170
00171 def write(self, data):
00172 self.send(data)
00173
00174 def flush(self):
00175 pass
00176
00177
00178 class TimeoutError(Exception):
00179 pass