Limit reads to block length.

This commit is contained in:
Aldo Cortesi 2012-10-09 16:25:15 +13:00
parent 15679e010d
commit 77869634e2
2 changed files with 17 additions and 2 deletions

View File

@ -106,13 +106,17 @@ class Writer(_FileLike):
class Reader(_FileLike): class Reader(_FileLike):
def read(self, length): def read(self, length):
""" """
If length is None, we read until connection closes. If length is -1, we read until connection closes.
""" """
result = '' result = ''
start = time.time() start = time.time()
while length == -1 or length > 0: while length == -1 or length > 0:
if length == -1 or length > self.BLOCKSIZE:
rlen = self.BLOCKSIZE
else:
rlen = length
try: try:
data = self.o.read(self.BLOCKSIZE if length == -1 else length) data = self.o.read(rlen)
except SSL.ZeroReturnError: except SSL.ZeroReturnError:
break break
except SSL.WantReadError: except SSL.WantReadError:

View File

@ -255,6 +255,17 @@ class TestTCPClient:
class TestFileLike: class TestFileLike:
def test_blocksize(self):
s = cStringIO.StringIO("1234567890abcdefghijklmnopqrstuvwxyz")
s = tcp.Reader(s)
s.BLOCKSIZE = 2
assert s.read(1) == "1"
assert s.read(2) == "23"
assert s.read(3) == "456"
assert s.read(4) == "7890"
d = s.read(-1)
assert d.startswith("abc") and d.endswith("xyz")
def test_wrap(self): def test_wrap(self):
s = cStringIO.StringIO("foobar\nfoobar") s = cStringIO.StringIO("foobar\nfoobar")
s.flush() s.flush()