Fix a couple bugs with Substream

- seek() now supports a second argument.

- peek() and read() no longer look beyond the end of the substream.

- slice() no longer doubles its offset.
This commit is contained in:
Eevee (Lexy Munroe) 2016-12-20 19:20:33 -08:00
parent db1ddc68e7
commit a38cb37e90

View file

@ -29,15 +29,18 @@ class Substream:
def read(self, n=-1): def read(self, n=-1):
self.stream.seek(self.offset + self.pos) self.stream.seek(self.offset + self.pos)
if n < 0: maxread = self.length - self.pos
n = self.length if n < 0 or 0 <= maxread < n:
elif self.length >= 0 and n > self.length: n = maxread
n = self.length
data = self.stream.read(n) data = self.stream.read(n)
self.pos += len(data) self.pos += len(data)
return data return data
def seek(self, offset): def seek(self, offset, whence=0):
if whence == 1:
offset += self.pos
elif whence == 2:
offset += self.length
offset = max(offset, 0) offset = max(offset, 0)
if self.length >= 0: if self.length >= 0:
offset = min(offset, self.length) offset = min(offset, self.length)
@ -60,7 +63,8 @@ class Substream:
def peek(self, n): def peek(self, n):
pos = self.stream.tell() pos = self.stream.tell()
self.stream.seek(self.offset + self.pos) self.stream.seek(self.offset + self.pos)
data = self.stream.read(n) maxread = self.length - self.pos
data = self.stream.read(min(maxread, n))
self.stream.seek(pos) self.stream.seek(pos)
return data return data
@ -71,7 +75,7 @@ class Substream:
def slice(self, offset, length=-1): def slice(self, offset, length=-1):
# TODO limit or warn if length is too long for this slice? # TODO limit or warn if length is too long for this slice?
return Substream(self, self.offset + offset, length) return Substream(self, offset, length)
class _ContainerFile: class _ContainerFile: