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,7 +61,8 @@ def test_seek_mode_2():
def test_read_n0():
# Arrange
with open(TEST_FILE, "rb") as fh:
for bytesmode in (True, False):
with open(TEST_FILE, "rb" if bytesmode else "r") as fh:
container = ContainerIO.ContainerIO(fh, 22, 100)
# Act
@ -69,12 +70,15 @@ def test_read_n0():
data = container.read()
# 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():
# Arrange
with open(TEST_FILE, "rb") as fh:
for bytesmode in (True, False):
with open(TEST_FILE, "rb" if bytesmode else "r") as fh:
container = ContainerIO.ContainerIO(fh, 22, 100)
# Act
@ -82,12 +86,15 @@ def test_read_n():
data = container.read(3)
# Assert
assert data == b"7\nT"
if bytesmode:
data = data.decode()
assert data == "7\nT"
def test_read_eof():
# Arrange
with open(TEST_FILE, "rb") as fh:
for bytesmode in (True, False):
with open(TEST_FILE, "rb" if bytesmode else "r") as fh:
container = ContainerIO.ContainerIO(fh, 22, 100)
# Act
@ -95,39 +102,46 @@ def test_read_eof():
data = container.read()
# Assert
assert data == b""
if bytesmode:
data = data.decode()
assert data == ""
def test_readline():
# Arrange
with open(TEST_FILE, "rb") as fh:
for bytesmode in (True, False):
with open(TEST_FILE, "rb" if bytesmode else "r") as fh:
container = ContainerIO.ContainerIO(fh, 0, 120)
# Act
data = container.readline()
# Assert
assert data == b"This is line 1\n"
if bytesmode:
data = data.decode()
assert data == "This is line 1\n"
def test_readlines():
# Arrange
for bytesmode in (True, False):
expected = [
b"This is line 1\n",
b"This is line 2\n",
b"This is line 3\n",
b"This is line 4\n",
b"This is line 5\n",
b"This is line 6\n",
b"This is line 7\n",
b"This is line 8\n",
"This is line 1\n",
"This is line 2\n",
"This is line 3\n",
"This is line 4\n",
"This is line 5\n",
"This is line 6\n",
"This is line 7\n",
"This is line 8\n",
]
with open(TEST_FILE, "rb") as fh:
with open(TEST_FILE, "rb" if bytesmode else "r") as fh:
container = ContainerIO.ContainerIO(fh, 0, 120)
# Act
data = container.readlines()
# Assert
if bytesmode:
data = [line.decode() for line in data]
assert data == expected

View File

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