Skip to content

Commit

Permalink
Fixed StreamReader._read_nowait (#1297)
Browse files Browse the repository at this point in the history
* Fixed StreamReader._read_nowait

* Added an entry to CHANGES
  • Loading branch information
dalazx authored and asvetlov committed Oct 11, 2016
1 parent 3ea78fb commit 869e133
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 6 deletions.
3 changes: 2 additions & 1 deletion CHANGES.rst
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@ CHANGES
1.0.5 (XXXX-XX-XX)
------------------

-
- Fix StreamReader._read_nowait to return all available
data up to the requested amount

-

Expand Down
20 changes: 15 additions & 5 deletions aiohttp/streams.py
Original file line number Diff line number Diff line change
Expand Up @@ -242,7 +242,7 @@ def readline(self):
offset = self._buffer_offset
ichar = self._buffer[0].find(b'\n', offset) + 1
# Read from current offset to found b'\n' or to the end.
data = self._read_nowait(ichar - offset if ichar else -1)
data = self._read_nowait_chunk(ichar - offset if ichar else -1)
line.append(data)
line_size += len(data)
if ichar:
Expand Down Expand Up @@ -340,10 +340,7 @@ def read_nowait(self, n=-1):

return self._read_nowait(n)

def _read_nowait(self, n):
if not self._buffer:
return EOF_MARKER

def _read_nowait_chunk(self, n):
first_buffer = self._buffer[0]
offset = self._buffer_offset
if n != -1 and len(first_buffer) - offset > n:
Expand All @@ -361,6 +358,19 @@ def _read_nowait(self, n):
self._buffer_size -= len(data)
return data

def _read_nowait(self, n):
chunks = []

while self._buffer:
chunk = self._read_nowait_chunk(n)
chunks.append(chunk)
if n != -1:
n -= len(chunk)
if n == 0:
break

return b''.join(chunks) if chunks else EOF_MARKER


class EmptyStreamReader(AsyncStreamReaderMixin):

Expand Down
22 changes: 22 additions & 0 deletions tests/test_streams.py
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,28 @@ def test_read_line_breaks(self):
data = self.loop.run_until_complete(stream.read(5))
self.assertEqual(b'line2', data)

def test_read_all(self):
# Read all avaliable buffered bytes
stream = self._make_one()
stream.feed_data(b'line1')
stream.feed_data(b'line2')
stream.feed_eof()

data = self.loop.run_until_complete(stream.read())
self.assertEqual(b'line1line2', data)

def test_read_up_to(self):
# Read available buffered bytes up to requested amount
stream = self._make_one()
stream.feed_data(b'line1')
stream.feed_data(b'line2')

data = self.loop.run_until_complete(stream.read(8))
self.assertEqual(b'line1lin', data)

data = self.loop.run_until_complete(stream.read(8))
self.assertEqual(b'e2', data)

def test_read_eof(self):
# Read bytes, stop at eof.
stream = self._make_one()
Expand Down

0 comments on commit 869e133

Please sign in to comment.