Add readinto to StreamingBody (boto#2746)

This adds a readinto implementation to StreamingBody, following the
contract outlined in the Python io package.

The primary motivation for this is that it allows consumers to wrap the
StreamingBody in an io.BufferedReader. This provides a nice performance
boost for parsing streaming binary data formats, where you want to be
able to read small numbers of bytes in your parsing algorithm, but
benefit from reading large chunks of data from upstream.

Author: aJanuary

botocore/httpchecksum.py

@@ -244,6 +244,17 @@ def read(self, amt=None):
             self._validate_checksum()
         return chunk
 
+    def readinto(self, b):
+        amount_read = super().readinto(b)
+        if amount_read == len(b):
+            view = b
+        else:
+            view = memoryview(b)[:amount_read]
+        self._checksum.update(view)
+        if amount_read == 0 and len(b) > 0:
+            self._validate_checksum()
+        return amount_read
+
     def _validate_checksum(self):
         if self._checksum.digest() != base64.b64decode(self._expected):
             error_msg = (

botocore/response.py

@@ -110,6 +110,22 @@ def read(self, amt=None):
             self._verify_content_length()
         return chunk
 
+    def readinto(self, b):
+        """Read bytes into a pre-allocated, writable bytes-like object b, and return the number of bytes read."""
+        try:
+            amount_read = self._raw_stream.readinto(b)
+        except URLLib3ReadTimeoutError as e:
+            # TODO: the url will be None as urllib3 isn't setting it yet
+            raise ReadTimeoutError(endpoint_url=e.url, error=e)
+        except URLLib3ProtocolError as e:
+            raise ResponseStreamingError(error=e)
+        self._amount_read += amount_read
+        if amount_read == 0 and len(b) > 0:
+            # If the server sends empty contents then we know we need to verify
+            # the content length.
+            self._verify_content_length()
+        return amount_read
+
     def readlines(self):
         return self._raw_stream.readlines()
 

tests/unit/test_httpchecksum.py

@@ -878,3 +878,38 @@ def test_iter_raises_error(self):
         with self.assertRaises(FlexibleChecksumError):
             for chunk in self.wrapper:
                 pass
+
+    def test_readinto_good(self):
+        chunk = bytearray(6)
+        self.assertEqual(6, self.wrapper.readinto(chunk))
+        self.assertEqual(chunk, bytearray(b"hello "))
+        self.assertEqual(5, self.wrapper.readinto(chunk))
+        # Note the trailing space here comes from the fact we've only got 5
+        # bytes left to read from the stream into a 6 byte buffer, so it leaves
+        # the last byte untouched, which is the space character from the
+        # previous read.
+        self.assertEqual(chunk, bytearray(b"world "))
+        # Whole body has been read, next read signals the end of the stream and
+        # validates the checksum of the body contents read
+        self.wrapper.readinto(chunk)
+
+    def test_readinto_bad(self):
+        self._make_wrapper("duorhq==")
+        chunk = bytearray(6)
+        self.assertEqual(6, self.wrapper.readinto(chunk))
+        self.assertEqual(chunk, bytearray(b"hello "))
+        self.assertEqual(5, self.wrapper.readinto(chunk))
+        self.assertEqual(chunk, bytearray(b"world "))
+        # Whole body has been read, next read signals the end of the stream and
+        # validates the checksum of the body contents read
+        with self.assertRaises(FlexibleChecksumError):
+            self.wrapper.readinto(chunk)
+
+    def test_readinto_zero_bytes(self):
+        # Test that readinto returns 0 when 0 bytes are requested
+        chunk = bytearray(0)
+        self.assertEqual(0, self.wrapper.readinto(chunk))
+        self.assertEqual(chunk, bytearray(b""))
+        # Whole body has been read, next read signals the end of the stream and
+        # validates the checksum of the body contents read
+        self.wrapper.readinto(chunk)

tests/unit/test_response.py

@@ -110,6 +110,59 @@ def test_streaming_body_readlines(self):
         chunks = [b'1234567890\n', b'1234567890\n', b'12345']
         self.assertEqual(stream.readlines(), chunks)
 
+    def test_streaming_body_readinto(self):
+        body = BytesIO(b"123456789")
+        stream = response.StreamingBody(body, content_length=9)
+        chunk = bytearray(b"\x00\x00\x00\x00\x00")
+        self.assertEqual(5, stream.readinto(chunk))
+        self.assertEqual(chunk, bytearray(b"\x31\x32\x33\x34\x35"))
+        self.assertEqual(4, stream.readinto(chunk))
+        self.assertEqual(chunk, bytearray(b"\x36\x37\x38\x39\x35"))
+
+    def test_streaming_body_readinto_with_invalid_length(self):
+        body = BytesIO(b"12")
+        stream = response.StreamingBody(body, content_length=9)
+        chunk = bytearray(b"\xde\xad\xbe\xef")
+        self.assertEqual(2, stream.readinto(chunk))
+        self.assertEqual(chunk, bytearray(b"\x31\x32\xbe\xef"))
+        with self.assertRaises(IncompleteReadError):
+            stream.readinto(chunk)
+
+    def test_streaming_body_readinto_with_empty_buffer(self):
+        body = BytesIO(b"12")
+        stream = response.StreamingBody(body, content_length=9)
+        chunk = bytearray(b"")
+        # Does not raise IncompleteReadError
+        self.assertEqual(0, stream.readinto(chunk))
+
+    def test_streaming_body_readinto_catches_urllib3_read_timeout(self):
+        class TimeoutBody:
+            def readinto(*args, **kwargs):
+                raise URLLib3ReadTimeoutError(None, None, None)
+
+            def geturl(*args, **kwargs):
+                return "http://example.com"
+
+        stream = response.StreamingBody(TimeoutBody(), content_length=None)
+        with self.assertRaises(ReadTimeoutError):
+            chunk = bytearray(b"\x00\x00\x00\x00\x00")
+            stream.readinto(chunk)
+
+    def test_streaming_body_readinto_catches_urllib3_protocol_error(self):
+        class ProtocolErrorBody:
+            def readinto(*args, **kwargs):
+                raise URLLib3ProtocolError(None, None, None)
+
+            def geturl(*args, **kwargs):
+                return "http://example.com"
+
+        stream = response.StreamingBody(
+            ProtocolErrorBody(), content_length=None
+        )
+        with self.assertRaises(ResponseStreamingError):
+            chunk = bytearray(b"\x00\x00\x00\x00\x00")
+            stream.readinto(chunk)
+
     def test_streaming_body_tell(self):
         body = BytesIO(b'1234567890')
         stream = response.StreamingBody(body, content_length=10)