#70 - allow streaming iterables of strings

Author: daggaz

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "json-stream"
-version = "2.4.0"
+version = "2.4.1"
 authors = [
     {name = "Jamie Cockburn", email="jamie_cockburn@hotmail.co.uk"},
 ]

src/json_stream/iterators.py

@@ -4,13 +4,21 @@
 class IterableStream(io.RawIOBase):
     def __init__(self, iterable):
         self.iterator = iter(iterable)
-        self.remainder = None
+        self.remainder = None  # type: bytes | None
+
+    def _normalize_chunk(self, chunk):
+        # Ensure chunk is bytes for writing into a binary buffer
+        if isinstance(chunk, str):
+            return chunk.encode()
+        return chunk
 
     def readinto(self, buffer):
         try:
-            chunk = self.remainder or next(self.iterator)
-            length = min(len(buffer), len(chunk))
-            buffer[:length], self.remainder = chunk[:length], chunk[length:]
+            # Wrap `buffer: WriteableBuffer` in memoryview to ensure len() and slicing
+            mv = memoryview(buffer)
+            chunk = self.remainder or self._normalize_chunk(next(self.iterator))
+            length = min(len(mv), len(chunk))
+            mv[:length], self.remainder = chunk[:length], chunk[length:]
             return length
         except StopIteration:
             return 0    # indicate EOF

src/json_stream/tests/test_iterators.py

@@ -15,3 +15,16 @@ def test_read(self):
         # stream it and check the result
         stream = IterableStream(data)
         self.assertEqual(stream.read(), b"".join(data))
+
+    def test_read_str_chunks(self):
+        # create some chunks of text data
+        data_str = (
+            "a" * io.DEFAULT_BUFFER_SIZE,
+            "b" * (io.DEFAULT_BUFFER_SIZE + 1),
+            "c" * (io.DEFAULT_BUFFER_SIZE - 1),
+        )
+        expected = ("".join(data_str)).encode()
+
+        # stream it and check the result
+        stream = IterableStream(data_str)
+        self.assertEqual(stream.read(), expected)