Repository: boto/botocore
Pull request #1195 (2017-05-15 16:52:47)

Add line iterator to StreamingBody
This way, we can readlines from a streaming body (without having to load
all bytes into memory!)

This will allow us to use StreamingBody in places where a Python
file-like object is expected.
(E.g. csv.reader!)


Patch:
diff --git a/botocore/response.py b/botocore/response.py
index ff024472f..1846c13ba 100644
--- a/botocore/response.py
+++ b/botocore/response.py
@@ -80,6 +80,51 @@ class StreamingBody(object):
             self._verify_content_length()
         return chunk
 
+    def __iter__(self):
+        """Return an iterator to yield lines from the raw stream.
+        """
+        # Reading 1024 bytes at a time seems a sane choice [citation needed].
+        # If the user needs to specify something different they can always call
+        # streaming_body._iter_lines(<other-chunk-size>)
+        default_chunk_size = 1024
+        return self._iter_lines(default_chunk_size)
+
+    def _iter_lines(self, chunk_size):
+        """Return an iterator to yield lines from the raw stream.
+
+        This is achieved by reading chunk of bytes (of size chunk_size) at a
+        time from the raw stream, and then yielding lines from there.
+        """
+        pending = None
+        for chunk in self._iter_chunks(chunk_size):
+            if pending is not None:
+                chunk = pending + chunk
+
+            lines = chunk.splitlines()
+
+            if lines and lines[-1] and chunk and lines[-1][-1] == chunk[-1]:
+                # We might be in the 'middle' of a line. Hence we keep the
+                # last line as pending.
+                pending = lines.pop()
+            else:
+                pending = None
+
+            for line in lines:
+                yield line
+
+        if pending is not None:
+            yield pending
+
+    def _iter_chunks(self, chunk_size):
+        """Return an iterator to yield chunks of chunk_size bytes from the raw
+        stream.
+        """
+        while True:
+            current_chunk = self.read(chunk_size)
+            if current_chunk == b"":
+                break
+            yield current_chunk
+
     def _verify_content_length(self):
         # See: https://github.com/kennethreitz/requests/issues/1855
         # Basically, our http library doesn't do this for us, so we have


Test patch:
diff --git a/tests/unit/test_response.py b/tests/unit/test_response.py
index 7b2eb2597..69c7ad666 100644
--- a/tests/unit/test_response.py
+++ b/tests/unit/test_response.py
@@ -68,6 +68,41 @@ class TestStreamWrapper(unittest.TestCase):
         stream.close()
         self.assertTrue(body.closed)
 
+    def test_streaming_line_iterator(self):
+        body = six.BytesIO(b'1234567890\n1234567890\n12345')
+        stream = response.StreamingBody(body, content_length=27)
+        self._assert_lines(
+            iter(stream),
+            [b'1234567890', b'1234567890', b'12345'],
+        )
+
+    def test_streaming_line_iterator_ends_newline(self):
+        body = six.BytesIO(b'1234567890\n1234567890\n12345\n')
+        stream = response.StreamingBody(body, content_length=28)
+        self._assert_lines(
+            iter(stream),
+            [b'1234567890', b'1234567890', b'12345'],
+        )
+
+    def test_streaming_line_iter_chunk_sizes(self):
+        for chunk_size in range(1, 30):
+            body = six.BytesIO(b'1234567890\n1234567890\n12345')
+            stream = response.StreamingBody(body, content_length=27)
+            self._assert_lines(
+                stream._iter_lines(chunk_size),
+                [b'1234567890', b'1234567890', b'12345'],
+            )
+
+    def _assert_lines(self, line_iterator, expected_lines):
+        for expected_line in expected_lines:
+            self.assertEqual(
+                next(line_iterator),
+                expected_line,
+            )
+        # We should have exhausted the iterator.
+        with self.assertRaises(StopIteration):
+            next(line_iterator)
+
 
 class TestGetResponse(BaseResponseTest):
     maxDiff = None