From 77869634e20ae5a2646d7455e499866e9cfafbab Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Tue, 9 Oct 2012 16:25:15 +1300 Subject: [PATCH] Limit reads to block length. --- netlib/tcp.py | 8 ++++++-- test/test_tcp.py | 11 +++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/netlib/tcp.py b/netlib/tcp.py index 414c12377..f8f877de5 100644 --- a/netlib/tcp.py +++ b/netlib/tcp.py @@ -106,13 +106,17 @@ class Writer(_FileLike): class Reader(_FileLike): def read(self, length): """ - If length is None, we read until connection closes. + If length is -1, we read until connection closes. """ result = '' start = time.time() while length == -1 or length > 0: + if length == -1 or length > self.BLOCKSIZE: + rlen = self.BLOCKSIZE + else: + rlen = length try: - data = self.o.read(self.BLOCKSIZE if length == -1 else length) + data = self.o.read(rlen) except SSL.ZeroReturnError: break except SSL.WantReadError: diff --git a/test/test_tcp.py b/test/test_tcp.py index c833ce077..5a12da91b 100644 --- a/test/test_tcp.py +++ b/test/test_tcp.py @@ -255,6 +255,17 @@ class TestTCPClient: 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): s = cStringIO.StringIO("foobar\nfoobar") s.flush()