Merge pull request #2677 from dlm6693/streaming-body-context-manager

allow StreamingBody class to be used as a context manager
boto · May 18, 2022 · 49bbfb3 · 49bbfb3
2 parents 02abf54 + b9c5409
commit 49bbfb3
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/botocore/response.py b/botocore/response.py
@@ -124,6 +124,12 @@ def __next__(self):
             return current_chunk
         raise StopIteration()
 
+    def __enter__(self):
+        return self._raw_stream
+
+    def __exit__(self, type, value, traceback):
+        self._raw_stream.close()
+
     next = __next__
 
     def iter_lines(self, chunk_size=_DEFAULT_CHUNK_SIZE, keepends=False):

diff --git a/tests/unit/test_response.py b/tests/unit/test_response.py
@@ -230,6 +230,13 @@ def test_streaming_line_empty_body(self):
         )
         self.assert_lines(stream.iter_lines(), [])
 
+    def test_streaming_body_as_context_manager(self):
+        body = BytesIO(b'1234567890')
+        with response.StreamingBody(body, content_length=10) as stream:
+            self.assertEqual(stream.read(), b'1234567890')
+            self.assertFalse(body.closed)
+        self.assertTrue(body.closed)
+
 
 class FakeRawResponse(BytesIO):
     def stream(self, amt=1024, decode_content=None):