Implemented readuntil in StreamResponse (#4734) (#5196)

Co-authored-by: Andrew Svetlov <[email protected]>

Co-authored-by: Anas <[email protected]>
Co-authored-by: Andrew Svetlov <[email protected]>

CHANGES/4054.feature

@@ -0,0 +1 @@
+Implemented readuntil in StreamResponse

aiohttp/streams.py

@@ -310,34 +310,41 @@ async def _wait(self, func_name: str) -> None:
             self._waiter = None
 
     async def readline(self) -> bytes:
+        return await self.readuntil()
+
+    async def readuntil(self, separator: bytes = b"\n") -> bytes:
+        seplen = len(separator)
+        if seplen == 0:
+            raise ValueError("Separator should be at least one-byte string")
+
         if self._exception is not None:
             raise self._exception
 
-        line = []
-        line_size = 0
+        chunk = b""
+        chunk_size = 0
         not_enough = True
 
         while not_enough:
             while self._buffer and not_enough:
                 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.
+                ichar = self._buffer[0].find(separator, offset) + 1
+                # Read from current offset to found separator or to the end.
                 data = self._read_nowait_chunk(ichar - offset if ichar else -1)
-                line.append(data)
-                line_size += len(data)
+                chunk += data
+                chunk_size += len(data)
                 if ichar:
                     not_enough = False
 
-                if line_size > self._high_water:
-                    raise ValueError("Line is too long")
+                if chunk_size > self._high_water:
+                    raise ValueError("Chunk too big")
 
             if self._eof:
                 break
 
             if not_enough:
-                await self._wait("readline")
+                await self._wait("readuntil")
 
-        return b"".join(line)
+        return chunk
 
     async def read(self, n: int = -1) -> bytes:
         if self._exception is not None:
@@ -531,6 +538,8 @@ async def readline(self) -> bytes:
     async def read(self, n: int = -1) -> bytes:
         return b""
 
+    # TODO add async def readuntil
+
     async def readany(self) -> bytes:
         return b""
 

docs/streams.rst

@@ -70,6 +70,19 @@ Reading Methods
 
    :return bytes: the given line
 
+.. comethod:: StreamReader.readuntil(separator="\n")
+
+   Read until separator, where `separator` is a sequence of bytes.
+
+   If EOF is received, and `separator` was not found, the method will
+   return the partial read bytes.
+
+   If the EOF was received and the internal buffer is empty, return an
+   empty bytes object.
+
+   .. versionadded:: 3.8
+
+   :return bytes: the given data
 
 .. comethod:: StreamReader.readchunk()
 

tests/test_streams.py

@@ -394,6 +394,117 @@ async def test_readline_exception(self) -> None:
         with pytest.raises(ValueError):
             await stream.readline()
 
+    async def test_readuntil(self) -> None:
+        loop = asyncio.get_event_loop()
+        # Read one chunk. 'readuntil' will need to wait for the data
+        # to come from 'cb'
+        stream = self._make_one()
+        stream.feed_data(b"chunk1 ")
+        read_task = loop.create_task(stream.readuntil(b"*"))
+
+        def cb():
+            stream.feed_data(b"chunk2 ")
+            stream.feed_data(b"chunk3 ")
+            stream.feed_data(b"* chunk4")
+
+        loop.call_soon(cb)
+
+        line = await read_task
+        assert b"chunk1 chunk2 chunk3 *" == line
+
+        stream.feed_eof()
+        data = await stream.read()
+        assert b" chunk4" == data
+
+    async def test_readuntil_limit_with_existing_data(self) -> None:
+        # Read one chunk. The data is in StreamReader's buffer
+        # before the event loop is run.
+
+        stream = self._make_one(limit=2)
+        stream.feed_data(b"li")
+        stream.feed_data(b"ne1&line2&")
+
+        with pytest.raises(ValueError):
+            await stream.readuntil(b"&")
+        # The buffer should contain the remaining data after exception
+        stream.feed_eof()
+        data = await stream.read()
+        assert b"line2&" == data
+
+    async def test_readuntil_limit(self) -> None:
+        loop = asyncio.get_event_loop()
+        # Read one chunk. StreamReaders are fed with data after
+        # their 'readuntil' methods are called.
+        stream = self._make_one(limit=4)
+
+        def cb():
+            stream.feed_data(b"chunk1")
+            stream.feed_data(b"chunk2$")
+            stream.feed_data(b"chunk3#")
+            stream.feed_eof()
+
+        loop.call_soon(cb)
+
+        with pytest.raises(ValueError):
+            await stream.readuntil(b"$")
+        data = await stream.read()
+        assert b"chunk3#" == data
+
+    async def test_readuntil_nolimit_nowait(self) -> None:
+        # All needed data for the first 'readuntil' call will be
+        # in the buffer.
+        stream = self._make_one()
+        data = b"line1!line2!line3!"
+        stream.feed_data(data[:6])
+        stream.feed_data(data[6:])
+
+        line = await stream.readuntil(b"!")
+        assert b"line1!" == line
+
+        stream.feed_eof()
+        data = await stream.read()
+        assert b"line2!line3!" == data
+
+    async def test_readuntil_eof(self) -> None:
+        stream = self._make_one()
+        stream.feed_data(b"some data")
+        stream.feed_eof()
+
+        line = await stream.readuntil(b"@")
+        assert b"some data" == line
+
+    async def test_readuntil_empty_eof(self) -> None:
+        stream = self._make_one()
+        stream.feed_eof()
+
+        line = await stream.readuntil(b"@")
+        assert b"" == line
+
+    async def test_readuntil_read_byte_count(self) -> None:
+        stream = self._make_one()
+        data = b"line1!line2!line3!"
+        stream.feed_data(data)
+
+        await stream.readuntil(b"!")
+
+        data = await stream.read(7)
+        assert b"line2!l" == data
+
+        stream.feed_eof()
+        data = await stream.read()
+        assert b"ine3!" == data
+
+    async def test_readuntil_exception(self) -> None:
+        stream = self._make_one()
+        stream.feed_data(b"line#")
+
+        data = await stream.readuntil(b"#")
+        assert b"line#" == data
+
+        stream.set_exception(ValueError())
+        with pytest.raises(ValueError):
+            await stream.readuntil(b"#")
+
     async def test_readexactly_zero_or_less(self) -> None:
         # Read exact number of bytes (zero or less).
         stream = self._make_one()