added readline support to json string streams

Author: daggaz

src/json_stream/tokenizer/strings.py

@@ -28,11 +28,19 @@ class JsonStringReader(io.TextIOBase):
     def __init__(self, stream: io.TextIOBase, initial_buffer=''):
         self.stream = stream
         self.buffer = initial_buffer
+        self.readline_buffer = ''
         self.unicode_buffer = ''
         self.state = CHAR
-        self.complete = False
+        self.end_of_string = False
         self.index = 0
 
+    @property
+    def complete(self):
+        return self.end_of_string and not self.readline_buffer
+
+    def readable(self) -> bool:
+        return True
+
     def read(self, size: Union[int, None] = None) -> str:
         result = ''
         length = DEFAULT_BUFFER_SIZE
@@ -43,6 +51,9 @@ def read(self, size: Union[int, None] = None) -> str:
         return result
 
     def _read_chunk(self, size: Union[int, None] = ...) -> str:
+        if self.readline_buffer:
+            result, self.readline_buffer = self.readline_buffer[:size], self.readline_buffer[size:]
+            return result
         chunk = self.buffer or self.stream.read(size)
         if not chunk:
             raise ValueError("Unterminated string at end of file")
@@ -60,7 +71,7 @@ def _read_chunk(self, size: Union[int, None] = ...) -> str:
             if state == CHAR:
                 if c == '"':
                     result += chunk[start:i]
-                    self.complete = True
+                    self.end_of_string = True
                     self.buffer = chunk[i + 1:]
                     break
                 elif c == "\\":
@@ -136,3 +147,23 @@ def _read_chunk(self, size: Union[int, None] = ...) -> str:
         self.state = state
         self.unicode_buffer = unicode_buffer
         return result
+
+    def readline(self, size: int = None) -> str:
+        result = ''
+        read_size = DEFAULT_BUFFER_SIZE
+        while not self.complete:
+            if size:
+                result_length = len(result)
+                if result_length >= size:
+                    result, self.readline_buffer = result[:size], result[size:] + self.readline_buffer
+                    break
+                read_size = size - result_length
+            chunk = self._read_chunk(read_size)
+            i = chunk.find('\n')
+            if i < 0:
+                result += chunk
+            else:
+                chunk, self.readline_buffer = chunk[:i+1], chunk[i+1:]
+                result += chunk
+                break
+        return result

src/json_stream/tokenizer/tests/test_strings.py

@@ -1,6 +1,7 @@
 import re
 from io import StringIO
 from unittest import TestCase
+from unittest.mock import patch
 
 from json_stream.tokenizer.strings import JsonStringReader
 
@@ -76,7 +77,7 @@ def test_unterminated_strings_while_in_escape(self):
         self.assertStringRaises(r'"\u!!!', "Unterminated string at end of file")
 
     def test_with_initial_buffer(self):
-        self.assertStringEquals("there will be more string", buffer='"there will be ', stream='more string"')
+        self.assertStringEquals("there will be more string", buffer='"there will be ', stream='more string"')  # x   x x
 
     def test_remainder(self):
         reader, f = self.assertStringEquals(
@@ -132,6 +133,166 @@ def test_read_over_split_escape(self):
             buffer, stream = json[:i], json[i:]
             self.assertStringEquals("abcdeÄedcba", buffer=buffer, stream=stream)
 
+    def test_readable(self):
+        reader = JsonStringReader(StringIO())
+        self.assertTrue(reader.readable())
+
+    def test_readline(self):
+        stream = StringIO(r'some\nlines\nof\ntext"')
+        reader = JsonStringReader(stream)
+        self.assertReadline(
+            reader, stream,
+            result='some\n',
+            remaining_readline_buffer='lines\nof\ntext',
+            complete=False,
+        )
+        self.assertReadline(
+            reader, stream,
+            result='lines\n',
+            remaining_readline_buffer='of\ntext',
+            complete=False,
+        )
+        self.assertReadline(
+            reader, stream,
+            result='of\n',
+            remaining_readline_buffer='text',
+            complete=False,
+        )
+        self.assertReadline(
+            reader, stream,
+            result='text',
+        )
+
+    @patch('json_stream.tokenizer.strings.DEFAULT_BUFFER_SIZE', 10)
+    def test_readline_needs_multiple_reads(self):
+        stream = StringIO(r'aaaaaaaaaabbbbb\ncccdddddddd"')
+        reader = JsonStringReader(stream)
+        self.assertReadline(
+            reader, stream,
+            result='aaaaaaaaaabbbbb\n',
+            remaining_readline_buffer='ccc',
+            remaining_stream='dddddddd"',
+            complete=False,
+        )
+        self.assertReadline(reader, stream, 'cccdddddddd')
+
+    def test_readline_eof_without_newline(self):
+        stream = StringIO(r'aaaaaaaaaabbbbbcccdddddddd"')
+        reader = JsonStringReader(stream)
+        self.assertReadline(
+            reader, stream,
+            result='aaaaaaaaaabbbbbcccdddddddd',
+        )
+        self.assertReadline(reader, stream, '')
+
+    @patch('json_stream.tokenizer.strings.DEFAULT_BUFFER_SIZE', 10)
+    def test_readline_then_read(self):
+        stream = StringIO(r'aaaaaaaaaabbbbbbbb\ndddddddd"')
+        reader = JsonStringReader(stream)
+        self.assertReadline(
+            reader, stream,
+            result='aaaaaaaaaabbbbbbbb\n',
+            remaining_stream='dddddddd"',
+            complete=False,
+        )
+        self.assertRead(reader, stream, result='dddddddd')
+
+    @patch('json_stream.tokenizer.strings.DEFAULT_BUFFER_SIZE', 10)
+    def test_readline_then_read_with_data_in_buffer(self):
+        stream = StringIO(r'aaaaaaaaaabbbbb\ncccdddddddd"')
+        reader = JsonStringReader(stream)
+        self.assertReadline(
+            reader, stream,
+            result='aaaaaaaaaabbbbb\n',
+            remaining_readline_buffer='ccc',
+            remaining_stream='dddddddd"',
+            complete=False,
+        )
+        self.assertRead(reader, stream, result='cccdddddddd')
+
+    def test_read_then_readline(self):
+        stream = StringIO(r'aaaaaaaaaabbbbb\ncccdddddddd"')
+        reader = JsonStringReader(stream)
+        self.assertRead(
+            reader, stream,
+            result='aaaaaaaaaa',
+            remaining_stream=r'bbbbb\ncccdddddddd"',
+            amount=10,
+            complete=False,
+        )
+        self.assertReadline(
+            reader, stream,
+            result='bbbbb\n',
+            remaining_readline_buffer='cccdddddddd',
+            complete=False,
+        )
+        self.assertReadline(
+            reader, stream,
+            result='cccdddddddd',
+        )
+
+    def test_readline_with_size_shorter_than_line(self):
+        stream = StringIO(r'aaaaaaaaaabbbbb\ncccdddddddd"')
+        reader = JsonStringReader(stream)
+        self.assertReadline(
+            reader, stream,
+            result='aaaaaaaaaa',
+            remaining_stream=r'bbbbb\ncccdddddddd"',
+            amount=10,
+            complete=False,
+        )
+        self.assertReadline(
+            reader, stream,
+            result='bbbbb\n',
+            remaining_readline_buffer='cccdddddddd',
+            complete=False,
+        )
+        self.assertReadline(
+            reader, stream,
+            result='cccdddddddd',
+        )
+
+    def test_readline_with_size_longer_than_line(self):
+        stream = StringIO(r'aaaaaaaaaabbbbb\ncccdddddddd"')
+        reader = JsonStringReader(stream)
+        self.assertReadline(
+            reader, stream,
+            result='aaaaaaaaaabbbbb\n',
+            remaining_readline_buffer='ccc',
+            remaining_stream='dddddddd"',
+            amount=20,
+            complete=False,
+        )
+        self.assertReadline(reader, stream, 'cccdddddddd')
+
+    def test_readline_trailing_newline(self):
+        stream = StringIO(r'a\n"')
+        reader = JsonStringReader(stream)
+        self.assertReadline(
+            reader, stream,
+            result='a\n',
+        )
+
+    def test_readline_no_trailing_newline(self):
+        stream = StringIO(r'a\nb"')
+        reader = JsonStringReader(stream)
+        self.assertReadline(
+            reader, stream,
+            result='a\n',
+            remaining_readline_buffer='b',
+            complete=False
+        )
+        self.assertReadline(reader, stream, 'b')
+
+    def test_readlines(self):
+        stream = StringIO(r'some\nlines\nof\ntext"')
+        reader = JsonStringReader(stream)
+        self.assertListEqual(["some\n", "lines\n", "of\n", "text"], reader.readlines())
+        self.assertEqual('', reader.readline_buffer)
+        self.assertEqual('', reader.buffer)
+        self.assertEqual('', stream.read())
+        self.assertTrue(reader.complete)
+
     def assertStringEquals(self, result, stream, buffer='', remaining_buffer='', remaining_stream='', amount=None,
                            complete=True):
         if buffer:
@@ -145,6 +306,17 @@ def assertStringEquals(self, result, stream, buffer='', remaining_buffer='', rem
 
     def assertRead(self, reader, stream, result, remaining_buffer='', remaining_stream='', amount=None, complete=True):
         self.assertEqual(result, reader.read(amount))
+        self.assertEqual(reader.readline_buffer, '')
+        self.assertEqual(remaining_buffer, reader.buffer)
+        pos = stream.tell()
+        self.assertEqual(remaining_stream, stream.read())
+        stream.seek(pos)
+        self.assertEqual(complete, reader.complete)
+
+    def assertReadline(self, reader, stream, result, remaining_readline_buffer='', remaining_buffer='',
+                       remaining_stream='', amount=None, complete=True):
+        self.assertEqual(result, reader.readline(amount))
+        self.assertEqual(remaining_readline_buffer, reader.readline_buffer)
         self.assertEqual(remaining_buffer, reader.buffer)
         pos = stream.tell()
         self.assertEqual(remaining_stream, stream.read())