Return strings or bytes from ContainerIO according to the file object mode

This commit is contained in:
Andrew Murray 2020-01-07 21:06:02 +11:00
parent e817ed0d3e
commit f958e2f8ed
2 changed files with 61 additions and 47 deletions

View File

@ -61,73 +61,87 @@ def test_seek_mode_2():
def test_read_n0(): def test_read_n0():
# Arrange # Arrange
with open(TEST_FILE, "rb") as fh: for bytesmode in (True, False):
container = ContainerIO.ContainerIO(fh, 22, 100) with open(TEST_FILE, "rb" if bytesmode else "r") as fh:
container = ContainerIO.ContainerIO(fh, 22, 100)
# Act # Act
container.seek(81) container.seek(81)
data = container.read() data = container.read()
# Assert # Assert
assert data == b"7\nThis is line 8\n" if bytesmode:
data = data.decode()
assert data == "7\nThis is line 8\n"
def test_read_n(): def test_read_n():
# Arrange # Arrange
with open(TEST_FILE, "rb") as fh: for bytesmode in (True, False):
container = ContainerIO.ContainerIO(fh, 22, 100) with open(TEST_FILE, "rb" if bytesmode else "r") as fh:
container = ContainerIO.ContainerIO(fh, 22, 100)
# Act # Act
container.seek(81) container.seek(81)
data = container.read(3) data = container.read(3)
# Assert # Assert
assert data == b"7\nT" if bytesmode:
data = data.decode()
assert data == "7\nT"
def test_read_eof(): def test_read_eof():
# Arrange # Arrange
with open(TEST_FILE, "rb") as fh: for bytesmode in (True, False):
container = ContainerIO.ContainerIO(fh, 22, 100) with open(TEST_FILE, "rb" if bytesmode else "r") as fh:
container = ContainerIO.ContainerIO(fh, 22, 100)
# Act # Act
container.seek(100) container.seek(100)
data = container.read() data = container.read()
# Assert # Assert
assert data == b"" if bytesmode:
data = data.decode()
assert data == ""
def test_readline(): def test_readline():
# Arrange # Arrange
with open(TEST_FILE, "rb") as fh: for bytesmode in (True, False):
container = ContainerIO.ContainerIO(fh, 0, 120) with open(TEST_FILE, "rb" if bytesmode else "r") as fh:
container = ContainerIO.ContainerIO(fh, 0, 120)
# Act # Act
data = container.readline() data = container.readline()
# Assert # Assert
assert data == b"This is line 1\n" if bytesmode:
data = data.decode()
assert data == "This is line 1\n"
def test_readlines(): def test_readlines():
# Arrange # Arrange
expected = [ for bytesmode in (True, False):
b"This is line 1\n", expected = [
b"This is line 2\n", "This is line 1\n",
b"This is line 3\n", "This is line 2\n",
b"This is line 4\n", "This is line 3\n",
b"This is line 5\n", "This is line 4\n",
b"This is line 6\n", "This is line 5\n",
b"This is line 7\n", "This is line 6\n",
b"This is line 8\n", "This is line 7\n",
] "This is line 8\n",
with open(TEST_FILE, "rb") as fh: ]
container = ContainerIO.ContainerIO(fh, 0, 120) with open(TEST_FILE, "rb" if bytesmode else "r") as fh:
container = ContainerIO.ContainerIO(fh, 0, 120)
# Act # Act
data = container.readlines() data = container.readlines()
# Assert # Assert
if bytesmode:
assert data == expected data = [line.decode() for line in data]
assert data == expected

View File

@ -82,7 +82,7 @@ class ContainerIO:
else: else:
n = self.length - self.pos n = self.length - self.pos
if not n: # EOF if not n: # EOF
return b"" return b"" if "b" in self.fh.mode else ""
self.pos = self.pos + n self.pos = self.pos + n
return self.fh.read(n) return self.fh.read(n)
@ -92,13 +92,13 @@ class ContainerIO:
:returns: An 8-bit string. :returns: An 8-bit string.
""" """
s = b"" s = b"" if "b" in self.fh.mode else ""
while True: while True:
c = self.read(1) c = self.read(1)
if not c: if not c:
break break
s = s + c s = s + c
if c == b"\n": if c == (b"\n" if "b" in self.fh.mode else "\n"):
break break
return s return s