PYTHON-1031 - GridFsBucket.download_to_stream now uses GridOutIterator

Author: aherlihy

gridfs/__init__.py

@@ -525,7 +525,7 @@ def download_to_stream(self, file_id, destination):
  # Get _id of file to read
  file_id = fs.upload_from_stream("test_file", "data I want to store!")
  # Get file to write to
- file = open('myfile','rw')
+ file = open('myfile','rwb')
  fs.download_to_stream(file_id, file)
  contents = file.read()
 
@@ -536,7 +536,8 @@ def download_to_stream(self, file_id, destination):
  -`destination`: a file-like object implementing :meth:`write`.
  """
  gout = self.open_download_stream(file_id)
- destination.write(gout)
+ for chunk in gout:
+ destination.write(chunk)
 
  def delete(self, file_id):
  """Given an file_id, delete this stored file's files collection document
@@ -663,7 +664,7 @@ def download_to_stream_by_name(self, filename, destination, revision=-1):
  my_db = MongoClient().test
  fs = GridFSBucket(my_db)
  # Get file to write to
- file = open('myfile','w')
+ file = open('myfile','wb')
  fs.download_to_stream_by_name("test_file", file)
 
  Raises :exc:`~gridfs.errors.NoFile` if no such version of
@@ -687,8 +688,8 @@ def download_to_stream_by_name(self, filename, destination, revision=-1):
  -1 = the most recent revision
  """
  gout = self.open_download_stream_by_name(filename, revision)
-
- destination.write(gout)
+ for chunk in gout:
+  destination.write(chunk)
 
  def rename(self, file_id, new_filename):
  """Renames the stored file with the specified file_id.

test/test_gridfs_bucket.py

@@ -369,6 +369,56 @@ def test_abort(self):
  self.assertEqual(0, self.db.fs.chunks.count(
  {"files_id": gin._id}))
 
+ def test_download_to_stream(self):
+ file1 = StringIO(b"hello world")
+ # Test with one chunk.
+ oid = self.fs.upload_from_stream("one_chunk", file1)
+ self.assertEqual(1, self.db.fs.chunks.count())
+ file2 = StringIO()
+ self.fs.download_to_stream(oid, file2)
+ file1.seek(0)
+ file2.seek(0)
+ self.assertEqual(file1.read(), file2.read())
+
+ # Test with many chunks.
+ self.db.drop_collection("fs.files")
+ self.db.drop_collection("fs.chunks")
+ file1.seek(0)
+ oid = self.fs.upload_from_stream("many_chunks",
+ file1,
+ chunk_size_bytes=1)
+ self.assertEqual(11, self.db.fs.chunks.count())
+ file2 = StringIO()
+ self.fs.download_to_stream(oid, file2)
+ file1.seek(0)
+ file2.seek(0)
+ self.assertEqual(file1.read(), file2.read())
+
+ def test_download_to_stream_by_name(self):
+ file1 = StringIO(b"hello world")
+ # Test with one chunk.
+ oid = self.fs.upload_from_stream("one_chunk", file1)
+ self.assertEqual(1, self.db.fs.chunks.count())
+ file2 = StringIO()
+ self.fs.download_to_stream_by_name("one_chunk", file2)
+ file1.seek(0)
+ file2.seek(0)
+ self.assertEqual(file1.read(), file2.read())
+
+ # Test with many chunks.
+ self.db.drop_collection("fs.files")
+ self.db.drop_collection("fs.chunks")
+ file1.seek(0)
+ self.fs.upload_from_stream("many_chunks", file1, chunk_size_bytes=1)
+ self.assertEqual(11, self.db.fs.chunks.count())
+
+ file2 = StringIO()
+ self.fs.download_to_stream_by_name("many_chunks", file2)
+ file1.seek(0)
+ file2.seek(0)
+ self.assertEqual(file1.read(), file2.read())
+
+
 class TestGridfsBucketReplicaSet(TestReplicaSetClientBase):
 
  def test_gridfs_replica_set(self):