diff --git a/ld_eventsource/async_reader.py b/ld_eventsource/async_reader.py index d71aa9b..2469380 100644 --- a/ld_eventsource/async_reader.py +++ b/ld_eventsource/async_reader.py @@ -12,7 +12,8 @@ class _AsyncBufferedLineReader: @staticmethod async def lines_from(chunks: AsyncIterator[bytes]) -> AsyncIterator[str]: last_char_was_cr = False - partial_line = None + # Join fragments only when a line ends to avoid repeatedly copying the partial line. + pending_fragments: list = [] async for chunk in chunks: if len(chunk) == 0: @@ -25,14 +26,22 @@ async def lines_from(chunks: AsyncIterator[bytes]) -> AsyncIterator[str]: lines.pop(0) if len(lines) == 0: continue - if partial_line is not None: - lines[0] = partial_line + lines[0] - partial_line = None + last_char = chunk[-1] + terminated = last_char in (10, 13) + + if pending_fragments: + pending_fragments.append(lines[0]) + if len(lines) == 1 and not terminated: + continue # this chunk continues the pending line, but does not end it + # This chunk ends the pending line, so join the fragments one time. + lines[0] = b"".join(pending_fragments) + pending_fragments = [] + if last_char == 13: last_char_was_cr = True elif last_char != 10: - partial_line = lines.pop() + pending_fragments = [lines.pop()] for line in lines: yield line.decode() diff --git a/ld_eventsource/reader.py b/ld_eventsource/reader.py index 8b2756e..372006a 100644 --- a/ld_eventsource/reader.py +++ b/ld_eventsource/reader.py @@ -16,7 +16,8 @@ def lines_from(chunks): series of strings, each of which is one line of text. The line does not include the terminator. """ last_char_was_cr = False - partial_line = None + # Join fragments only when a line ends to avoid repeatedly copying the partial line. + pending_fragments = [] for chunk in chunks: if len(chunk) == 0: @@ -35,18 +36,24 @@ def lines_from(chunks): lines.pop(0) if len(lines) == 0: continue # ran out of data, continue to get next chunk - if partial_line is not None: - # On our last time through the loop, we ended up with an unterminated line, so we should - # treat our first parsed line here as a continuation of that. - lines[0] = partial_line + lines[0] - partial_line = None + + last_char = chunk[-1] + terminated = last_char in (10, 13) + + if pending_fragments: + pending_fragments.append(lines[0]) + if len(lines) == 1 and not terminated: + continue # this chunk continues the pending line, but does not end it + # This chunk ends the pending line, so join the fragments one time. + lines[0] = b"".join(pending_fragments) + pending_fragments = [] + # Check whether the buffer really ended in a terminator. If it did not, then the last line in # lines is a partial line and should not be emitted yet. - last_char = chunk[-1] if last_char == 13: last_char_was_cr = True # remember this in case the next chunk starts with \n elif last_char != 10: - partial_line = lines.pop() # remove last element which is the partial line + pending_fragments = [lines.pop()] # remove last element which is the partial line for line in lines: yield line.decode() diff --git a/ld_eventsource/testing/test_async_reader.py b/ld_eventsource/testing/test_async_reader.py index 05c02ca..f221028 100644 --- a/ld_eventsource/testing/test_async_reader.py +++ b/ld_eventsource/testing/test_async_reader.py @@ -66,6 +66,25 @@ async def test_line_reader_empty_chunk(): assert lines == ["hello", "world"] +@pytest.mark.asyncio +@pytest.mark.parametrize("terminator", ["\r", "\n", "\r\n"]) +async def test_line_reader_fragments_with_empty_chunks(terminator): + chunks = [s.replace("*", terminator).encode() + for s in ["first", "", " line", "*", "second*", "unterminated"]] + lines = await lines_from_bytes(*chunks) + assert lines == ["first line", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("terminator", ["\r", "\n", "\r\n"]) +async def test_line_reader_utf8_line_spanning_many_chunks(terminator): + line = "caf\u00e9 " * 1000 + data = (line + terminator + "next line" + terminator).encode() + chunks = [data[i:i + 7] for i in range(0, len(data), 7)] + lines = await lines_from_bytes(*chunks) + assert lines == [line, "next line"] + + @pytest.mark.asyncio async def test_sse_reader_simple_event(): items = await events_from_lines("data: hello", "") diff --git a/ld_eventsource/testing/test_reader.py b/ld_eventsource/testing/test_reader.py index b1bb893..ecaa4ca 100644 --- a/ld_eventsource/testing/test_reader.py +++ b/ld_eventsource/testing/test_reader.py @@ -25,6 +25,10 @@ def terminator(self, request): ["first line*second line*third", " line*fourth line*"], ["first line", "second line", "third line", "fourth line"], ], + [ + ["first", "", " line", "*", "second*", "unterminated"], + ["first line", "second"], + ], ] ) def inputs_outputs(self, terminator, request): @@ -38,7 +42,9 @@ def test_parsing(self, inputs_outputs): def test_mixed_terminators(self): chunks = [ - b"first line\nsecond line\r\nthird line\r", + b"first line\nsecond line\r\nthird ", + b"line\r", + b"", b"\nfourth line\r", b"\r\nlast\r\n", ] @@ -52,6 +58,12 @@ def test_mixed_terminators(self): ] assert list(_BufferedLineReader.lines_from(chunks)) == expected + def test_utf8_line_spanning_many_chunks(self, terminator): + line = "caf\u00e9 " * 1000 + data = (line + terminator + "next line" + terminator).encode() + chunks = [data[i:i + 7] for i in range(0, len(data), 7)] + assert list(_BufferedLineReader.lines_from(chunks)) == [line, "next line"] + class TestSSEReader: def expect_output(self, lines, expected):