From 0f038b1f899a6621213dba18ed98a430ae23f495 Mon Sep 17 00:00:00 2001 From: Eyal Posener Date: Tue, 22 Dec 2015 19:09:44 +0200 Subject: [PATCH] Add number of bytes to stream.read_nowait --- aiohttp/streams.py | 4 ++-- tests/test_streams.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/aiohttp/streams.py b/aiohttp/streams.py index 2086444cef7..a732c856860 100644 --- a/aiohttp/streams.py +++ b/aiohttp/streams.py @@ -297,7 +297,7 @@ def readexactly(self, n): return b''.join(blocks) - def read_nowait(self): + def read_nowait(self, n=None): if self._exception is not None: raise self._exception @@ -305,7 +305,7 @@ def read_nowait(self): raise RuntimeError( 'Called while some coroutine is waiting for incoming data.') - return self._read_nowait() + return self._read_nowait(n) def _read_nowait(self, n=None): if not self._buffer: diff --git a/tests/test_streams.py b/tests/test_streams.py index 360ccbcaed9..b760e15b81c 100644 --- a/tests/test_streams.py +++ b/tests/test_streams.py @@ -452,6 +452,20 @@ def test_read_nowait(self): data = self.loop.run_until_complete(stream.read()) self.assertEqual(b'', data) + def test_read_nowait_n(self): + stream = self._make_one() + stream.feed_data(b'line1\nline2\n') + + self.assertEqual( + stream.read_nowait(4), b'line') + self.assertEqual( + stream.read_nowait(), b'1\nline2\n') + self.assertIs( + stream.read_nowait(), streams.EOF_MARKER) + stream.feed_eof() + data = self.loop.run_until_complete(stream.read()) + self.assertEqual(b'', data) + def test_read_nowait_exception(self): stream = self._make_one() stream.feed_data(b'line\n')