From 04b1b8ad3431cdca8841bd5c97cca9980c363289 Mon Sep 17 00:00:00 2001
From: aherlihy <anna.herlihy@10gen.com>
Date: Tue, 11 Aug 2015 12:14:06 -0400
Subject: [PATCH] open_upload_stream returns only GridIn

---
 gridfs/__init__.py         | 15 +++++----------
 test/test_gridfs_bucket.py | 12 ++++++------
 2 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/gridfs/__init__.py b/gridfs/__init__.py
index 550f93bcf..fc22921e7 100644
--- a/gridfs/__init__.py
+++ b/gridfs/__init__.py
@@ -426,8 +426,7 @@ class GridFSBucket(object):
           grid_in.write("data I want to store!")
           grid_in.close()  # uploaded on close
 
-        Returns an instance of :class:`~gridfs.grid_file.GridIn` and the _id
-        of the file to upload.
+        Returns an instance of :class:`~gridfs.grid_file.GridIn`.
 
         Raises :exc:`~gridfs.errors.NoFile` if no such version of
         that file exists.
@@ -449,8 +448,7 @@ class GridFSBucket(object):
         if metadata is not None:
             opts["metadata"] = metadata
 
-        gin = GridIn(self._collection, **opts)
-        return gin, gin._id
+        return GridIn(self._collection, **opts)
 
     def upload_from_stream(self, filename, source, chunk_size_bytes=None,
                            metadata=None):
@@ -484,14 +482,11 @@ class GridFSBucket(object):
             files collection document. If not provided the metadata field will
             be omitted from the files collection document.
         """
-        gin, _id = self.open_upload_stream(filename, chunk_size_bytes,
-                                           metadata)
-        try:
+        with self.open_upload_stream(
+                filename, chunk_size_bytes, metadata) as gin:
             gin.write(source)
-        finally:
-            gin.close()
 
-        return _id
+        return gin._id
 
     def open_download_stream(self, file_id):
         """Opens a Stream from which the application can read the contents of
diff --git a/test/test_gridfs_bucket.py b/test/test_gridfs_bucket.py
index a06adc034..42faa1aac 100644
--- a/test/test_gridfs_bucket.py
+++ b/test/test_gridfs_bucket.py
@@ -53,7 +53,7 @@ class JustWrite(threading.Thread):
 
     def run(self):
         for _ in range(self.num):
-            file, _ = self.gfs.open_upload_stream("test")
+            file = self.gfs.open_upload_stream("test")
             file.write(b"hello")
             file.close()
 
@@ -224,7 +224,7 @@ class TestGridfs(IntegrationTest):
     def test_get_last_version(self):
         one = self.fs.upload_from_stream("test", b"foo")
         time.sleep(0.01)
-        two, _ = self.fs.open_upload_stream("test")
+        two = self.fs.open_upload_stream("test")
         two.write(b"bar")
         two.close()
         time.sleep(0.01)
@@ -356,18 +356,18 @@ class TestGridfs(IntegrationTest):
             "second_name").read())
 
     def test_abort(self):
-        gin, file_id = self.fs.open_upload_stream("test_filename",
-                                                  chunk_size_bytes=5)
+        gin = self.fs.open_upload_stream("test_filename",
+                                         chunk_size_bytes=5)
         gin.write(b"test1")
         gin.write(b"test2")
         gin.write(b"test3")
         self.assertEqual(3, self.db.fs.chunks.count(
-            {"files_id": file_id}))
+            {"files_id": gin._id}))
         gin.abort()
         self.assertTrue(gin.closed)
         self.assertRaises(ValueError, gin.write, b"test4")
         self.assertEqual(0, self.db.fs.chunks.count(
-            {"files_id": file_id}))
+            {"files_id": gin._id}))
 
 class TestGridfsBucketReplicaSet(TestReplicaSetClientBase):