168 lines
4.7 KiB
Python
168 lines
4.7 KiB
Python
WHENCE_ABSOLUTE = 0
|
|
WHENCE_RELATIVE = 1
|
|
WHENCE_RELATIVE_END = 2
|
|
|
|
READ_UNTIL_END = -1
|
|
|
|
|
|
class BaseStreamFilelike(object):
|
|
def __init__(self, fileobj):
|
|
self._fileobj = fileobj
|
|
self._cursor_position = 0
|
|
|
|
def close(self):
|
|
self._fileobj.close()
|
|
|
|
def read(self, size=READ_UNTIL_END):
|
|
buf = self._fileobj.read(size)
|
|
if buf is None:
|
|
return None
|
|
self._cursor_position += len(buf)
|
|
return buf
|
|
|
|
def tell(self):
|
|
return self._cursor_position
|
|
|
|
def seek(self, index, whence=WHENCE_ABSOLUTE):
|
|
num_bytes_to_ff = 0
|
|
if whence == WHENCE_ABSOLUTE:
|
|
if index < self._cursor_position:
|
|
raise IOError('Cannot seek backwards')
|
|
num_bytes_to_ff = index - self._cursor_position
|
|
|
|
elif whence == WHENCE_RELATIVE:
|
|
if index < 0:
|
|
raise IOError('Cannnot seek backwards')
|
|
num_bytes_to_ff = index
|
|
|
|
elif whence == WHENCE_RELATIVE_END:
|
|
raise IOError('Stream does not have a known end point')
|
|
|
|
bytes_forward = num_bytes_to_ff
|
|
while num_bytes_to_ff > 0:
|
|
buf = self._fileobj.read(num_bytes_to_ff)
|
|
if not buf:
|
|
raise IOError('Seek past end of file')
|
|
num_bytes_to_ff -= len(buf)
|
|
|
|
self._cursor_position += bytes_forward
|
|
return bytes_forward
|
|
|
|
|
|
class SocketReader(BaseStreamFilelike):
|
|
def __init__(self, fileobj):
|
|
super(SocketReader, self).__init__(fileobj)
|
|
self.handlers = []
|
|
|
|
def add_handler(self, handler):
|
|
self.handlers.append(handler)
|
|
|
|
def read(self, size=READ_UNTIL_END):
|
|
buf = super(SocketReader, self).read(size)
|
|
for handler in self.handlers:
|
|
handler(buf)
|
|
return buf
|
|
|
|
|
|
def wrap_with_handler(in_fp, handler):
|
|
wrapper = SocketReader(in_fp)
|
|
wrapper.add_handler(handler)
|
|
return wrapper
|
|
|
|
|
|
class FilelikeStreamConcat(object):
|
|
""" A file-like object which concats all the file-like objects in the specified generator into
|
|
a single stream.
|
|
"""
|
|
def __init__(self, file_generator):
|
|
self._file_generator = file_generator
|
|
self._current_file = file_generator.next()
|
|
self._current_position = 0
|
|
self._closed = False
|
|
|
|
def tell(self):
|
|
return self._current_position
|
|
|
|
def close(self):
|
|
self._closed = True
|
|
|
|
def read(self, size=READ_UNTIL_END):
|
|
buf = ''
|
|
current_size = size
|
|
|
|
while size == READ_UNTIL_END or len(buf) < size:
|
|
current_buf = self._current_file.read(current_size)
|
|
if current_buf:
|
|
buf += current_buf
|
|
self._current_position += len(current_buf)
|
|
if size != READ_UNTIL_END:
|
|
current_size -= len(current_buf)
|
|
|
|
else:
|
|
# That file was out of data, prime a new one
|
|
self._current_file.close()
|
|
try:
|
|
self._current_file = self._file_generator.next()
|
|
except StopIteration:
|
|
return buf
|
|
|
|
return buf
|
|
|
|
|
|
class StreamSlice(BaseStreamFilelike):
|
|
""" A file-like object which returns a file-like object that represents a slice of the data in
|
|
the specified file obj. All methods will act as if the slice is its own file.
|
|
"""
|
|
|
|
def __init__(self, fileobj, start_offset=0, end_offset_exclusive=READ_UNTIL_END):
|
|
super(StreamSlice, self).__init__(fileobj)
|
|
self._end_offset_exclusive = end_offset_exclusive
|
|
self._start_offset = start_offset
|
|
|
|
if start_offset > 0:
|
|
self.seek(start_offset)
|
|
|
|
def read(self, size=READ_UNTIL_END):
|
|
if self._end_offset_exclusive == READ_UNTIL_END:
|
|
# We weren't asked to limit the end of the stream
|
|
return super(StreamSlice, self).read(size)
|
|
|
|
# Compute the max bytes to read until the end or until we reach the user requested max
|
|
max_bytes_to_read = self._end_offset_exclusive - super(StreamSlice, self).tell()
|
|
if size != READ_UNTIL_END:
|
|
max_bytes_to_read = min(max_bytes_to_read, size)
|
|
|
|
return super(StreamSlice, self).read(max_bytes_to_read)
|
|
|
|
def _file_min(self, first, second):
|
|
if first == READ_UNTIL_END:
|
|
return second
|
|
|
|
if second == READ_UNTIL_END:
|
|
return first
|
|
|
|
return min(first, second)
|
|
|
|
def tell(self):
|
|
return super(StreamSlice, self).tell() - self._start_offset
|
|
|
|
def seek(self, index, whence=WHENCE_ABSOLUTE):
|
|
index = self._file_min(self._end_offset_exclusive, index)
|
|
super(StreamSlice, self).seek(index, whence)
|
|
|
|
|
|
class LimitingStream(StreamSlice):
|
|
""" A file-like object which mimics the specified file stream being limited to the given number
|
|
of bytes. All calls after that limit (if specified) will act as if the file has no additional
|
|
data.
|
|
"""
|
|
def __init__(self, fileobj, read_limit=READ_UNTIL_END, seekable=True):
|
|
super(LimitingStream, self).__init__(fileobj, 0, read_limit)
|
|
self._seekable = seekable
|
|
|
|
def seek(self, index, whence=WHENCE_ABSOLUTE):
|
|
if not self._seekable:
|
|
raise AttributeError
|
|
|
|
super(LimitingStream, self).seek(index, whence)
|