306 lines
11 KiB
Python
306 lines
11 KiB
Python
from io import BytesIO
|
|
from socket import error
|
|
from errno import EINTR
|
|
|
|
# Based on socket._fileobject from python2.7
|
|
class fileobject(object):
|
|
"""Faux file object attached to a socket object."""
|
|
|
|
default_bufsize = 8192
|
|
name = "<socket>"
|
|
|
|
__slots__ = ["mode", "bufsize", "softspace",
|
|
# "closed" is a property, see below
|
|
"_sock", "_rbufsize", "_wbufsize", "_rbuf", "_wbuf", "_wbuf_len",
|
|
"_close"]
|
|
|
|
def __init__(self, sock, mode='rb', bufsize=-1, close=False):
|
|
self._sock = sock
|
|
self.mode = mode # Not actually used in this version
|
|
if bufsize < 0:
|
|
bufsize = self.default_bufsize
|
|
self.bufsize = bufsize
|
|
self.softspace = False
|
|
# _rbufsize is the suggested recv buffer size. It is *strictly*
|
|
# obeyed within readline() for recv calls. If it is larger than
|
|
# default_bufsize it will be used for recv calls within read().
|
|
if bufsize == 0:
|
|
self._rbufsize = 1
|
|
elif bufsize == 1:
|
|
self._rbufsize = self.default_bufsize
|
|
else:
|
|
self._rbufsize = bufsize
|
|
self._wbufsize = bufsize
|
|
# We use BytesIO for the read buffer to avoid holding a list
|
|
# of variously sized string objects which have been known to
|
|
# fragment the heap due to how they are malloc()ed and often
|
|
# realloc()ed down much smaller than their original allocation.
|
|
self._rbuf = BytesIO()
|
|
self._wbuf = [] # A list of strings
|
|
self._wbuf_len = 0
|
|
self._close = close
|
|
|
|
def _getclosed(self):
|
|
return self._sock is None
|
|
closed = property(_getclosed, doc="True if the file is closed")
|
|
|
|
def close(self):
|
|
try:
|
|
if self._sock:
|
|
self.flush()
|
|
finally:
|
|
if self._close:
|
|
self._sock.close()
|
|
self._sock = None
|
|
|
|
def __del__(self):
|
|
try:
|
|
self.close()
|
|
except:
|
|
# close() may fail if __init__ didn't complete
|
|
pass
|
|
|
|
def flush(self):
|
|
if self._wbuf:
|
|
data = b"".join(self._wbuf)
|
|
self._wbuf = []
|
|
self._wbuf_len = 0
|
|
buffer_size = max(self._rbufsize, self.default_bufsize)
|
|
data_size = len(data)
|
|
write_offset = 0
|
|
view = memoryview(data)
|
|
try:
|
|
while write_offset < data_size:
|
|
self._sock.sendall(view[write_offset:write_offset+buffer_size])
|
|
write_offset += buffer_size
|
|
finally:
|
|
if write_offset < data_size:
|
|
remainder = data[write_offset:]
|
|
del view, data # explicit free
|
|
self._wbuf.append(remainder)
|
|
self._wbuf_len = len(remainder)
|
|
|
|
def fileno(self):
|
|
return self._sock.fileno()
|
|
|
|
def write(self, data):
|
|
data = bytes(data) # XXX Should really reject non-string non-buffers
|
|
if not data:
|
|
return
|
|
self._wbuf.append(data)
|
|
self._wbuf_len += len(data)
|
|
if (self._wbufsize == 0 or
|
|
(self._wbufsize == 1 and b'\n' in data) or
|
|
(self._wbufsize > 1 and self._wbuf_len >= self._wbufsize)):
|
|
self.flush()
|
|
|
|
def writelines(self, list):
|
|
# XXX We could do better here for very long lists
|
|
# XXX Should really reject non-string non-buffers
|
|
lines = filter(None, map(bytes, list))
|
|
self._wbuf_len += sum(map(len, lines))
|
|
self._wbuf.extend(lines)
|
|
if (self._wbufsize <= 1 or
|
|
self._wbuf_len >= self._wbufsize):
|
|
self.flush()
|
|
|
|
def read(self, size=-1):
|
|
# Use max, disallow tiny reads in a loop as they are very inefficient.
|
|
# We never leave read() with any leftover data from a new recv() call
|
|
# in our internal buffer.
|
|
rbufsize = max(self._rbufsize, self.default_bufsize)
|
|
# Our use of BytesIO rather than lists of string objects returned by
|
|
# recv() minimizes memory usage and fragmentation that occurs when
|
|
# rbufsize is large compared to the typical return value of recv().
|
|
buf = self._rbuf
|
|
buf.seek(0, 2) # seek end
|
|
if size < 0:
|
|
# Read until EOF
|
|
self._rbuf = BytesIO() # reset _rbuf. we consume it via buf.
|
|
while True:
|
|
try:
|
|
data = self._sock.recv(rbufsize)
|
|
except error as e:
|
|
if e.args[0] == EINTR:
|
|
continue
|
|
raise
|
|
if not data:
|
|
break
|
|
buf.write(data)
|
|
return buf.getvalue()
|
|
else:
|
|
# Read until size bytes or EOF seen, whichever comes first
|
|
buf_len = buf.tell()
|
|
if buf_len >= size:
|
|
# Already have size bytes in our buffer? Extract and return.
|
|
buf.seek(0)
|
|
rv = buf.read(size)
|
|
self._rbuf = BytesIO()
|
|
self._rbuf.write(buf.read())
|
|
return rv
|
|
|
|
self._rbuf = BytesIO() # reset _rbuf. we consume it via buf.
|
|
while True:
|
|
left = size - buf_len
|
|
# recv() will malloc the amount of memory given as its
|
|
# parameter even though it often returns much less data
|
|
# than that. The returned data string is short lived
|
|
# as we copy it into a BytesIO and free it. This avoids
|
|
# fragmentation issues on many platforms.
|
|
try:
|
|
data = self._sock.recv(left)
|
|
except error as e:
|
|
if e.args[0] == EINTR:
|
|
continue
|
|
raise
|
|
if not data:
|
|
break
|
|
n = len(data)
|
|
if n == size and not buf_len:
|
|
# Shortcut. Avoid buffer data copies when:
|
|
# - We have no data in our buffer.
|
|
# AND
|
|
# - Our call to recv returned exactly the
|
|
# number of bytes we were asked to read.
|
|
return data
|
|
if n == left:
|
|
buf.write(data)
|
|
del data # explicit free
|
|
break
|
|
assert n <= left, "recv(%d) returned %d bytes" % (left, n)
|
|
buf.write(data)
|
|
buf_len += n
|
|
del data # explicit free
|
|
#assert buf_len == buf.tell()
|
|
return buf.getvalue()
|
|
|
|
def readline(self, size=-1):
|
|
buf = self._rbuf
|
|
buf.seek(0, 2) # seek end
|
|
if buf.tell() > 0:
|
|
# check if we already have it in our buffer
|
|
buf.seek(0)
|
|
bline = buf.readline(size)
|
|
if bline.endswith(b'\n') or len(bline) == size:
|
|
self._rbuf = BytesIO()
|
|
self._rbuf.write(buf.read())
|
|
return bline
|
|
del bline
|
|
if size < 0:
|
|
# Read until \n or EOF, whichever comes first
|
|
if self._rbufsize <= 1:
|
|
# Speed up unbuffered case
|
|
buf.seek(0)
|
|
buffers = [buf.read()]
|
|
self._rbuf = BytesIO() # reset _rbuf. we consume it via buf.
|
|
data = None
|
|
recv = self._sock.recv
|
|
while True:
|
|
try:
|
|
while data != b"\n":
|
|
data = recv(1)
|
|
if not data:
|
|
break
|
|
buffers.append(data)
|
|
except error as e:
|
|
# The try..except to catch EINTR was moved outside the
|
|
# recv loop to avoid the per byte overhead.
|
|
if e.args[0] == EINTR:
|
|
continue
|
|
raise
|
|
break
|
|
return "".join(buffers)
|
|
|
|
buf.seek(0, 2) # seek end
|
|
self._rbuf = BytesIO() # reset _rbuf. we consume it via buf.
|
|
while True:
|
|
try:
|
|
data = self._sock.recv(self._rbufsize)
|
|
except error as e:
|
|
if e.args[0] == EINTR:
|
|
continue
|
|
raise
|
|
if not data:
|
|
break
|
|
nl = data.find(b'\n')
|
|
if nl >= 0:
|
|
nl += 1
|
|
buf.write(data[:nl])
|
|
self._rbuf.write(data[nl:])
|
|
del data
|
|
break
|
|
buf.write(data)
|
|
return buf.getvalue()
|
|
else:
|
|
# Read until size bytes or \n or EOF seen, whichever comes first
|
|
buf.seek(0, 2) # seek end
|
|
buf_len = buf.tell()
|
|
if buf_len >= size:
|
|
buf.seek(0)
|
|
rv = buf.read(size)
|
|
self._rbuf = BytesIO()
|
|
self._rbuf.write(buf.read())
|
|
return rv
|
|
self._rbuf = BytesIO() # reset _rbuf. we consume it via buf.
|
|
while True:
|
|
try:
|
|
data = self._sock.recv(self._rbufsize)
|
|
except error as e:
|
|
if e.args[0] == EINTR:
|
|
continue
|
|
raise
|
|
if not data:
|
|
break
|
|
left = size - buf_len
|
|
# did we just receive a newline?
|
|
nl = data.find(b'\n', 0, left)
|
|
if nl >= 0:
|
|
nl += 1
|
|
# save the excess data to _rbuf
|
|
self._rbuf.write(data[nl:])
|
|
if buf_len:
|
|
buf.write(data[:nl])
|
|
break
|
|
else:
|
|
# Shortcut. Avoid data copy through buf when returning
|
|
# a substring of our first recv().
|
|
return data[:nl]
|
|
n = len(data)
|
|
if n == size and not buf_len:
|
|
# Shortcut. Avoid data copy through buf when
|
|
# returning exactly all of our first recv().
|
|
return data
|
|
if n >= left:
|
|
buf.write(data[:left])
|
|
self._rbuf.write(data[left:])
|
|
break
|
|
buf.write(data)
|
|
buf_len += n
|
|
#assert buf_len == buf.tell()
|
|
return buf.getvalue()
|
|
|
|
def readlines(self, sizehint=0):
|
|
total = 0
|
|
list = []
|
|
while True:
|
|
line = self.readline()
|
|
if not line:
|
|
break
|
|
list.append(line)
|
|
total += len(line)
|
|
if sizehint and total >= sizehint:
|
|
break
|
|
return list
|
|
|
|
# Iterator protocols
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def next(self):
|
|
line = self.readline()
|
|
if not line:
|
|
raise StopIteration
|
|
return line
|
|
|