GridFile supports the context manager (with statement) protocol

gridfs/__init__.py

@@ -48,7 +48,8 @@ class GridFS(object):
 
         Only a single opened GridFile instance may exist for a file in gridfs
         at any time. Care must be taken to close GridFile instances when done
-        using them.
+        using them. GridFiles support the context manager protocol (the "with"
+        statement).
 
         :Parameters:
           - `filename`: name of the GridFile to open

gridfs/grid_file.py

@@ -48,7 +48,8 @@ class GridFile(object):
 
         Only a single opened GridFile instance may exist for a file in gridfs
         at any time. Care must be taken to close GridFile instances when done
-        using them.
+        using them. GridFiles support the context manager protocol (the "with"
+        statement).
 
         Raises TypeError if file_spec is not an instance of dict, database is
         not an instance of `pymongo.database.Database`, or collection is not an
@@ -306,3 +307,16 @@ class GridFile(object):
         """
         for line in sequence:
             self.write(line)
+
+    def __enter__(self):
+        """Support for the context manager protocol.
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Support for the context manager protocol.
+
+        Close the file and allow exceptions to propogate.
+        """
+        self.close()
+        return False # propogate exceptions

test/gridfs15.py

@@ -0,0 +1,25 @@
+# Copyright 2009 10gen, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Some tests for the gridfs package that only work under Python >= 1.5.
+"""
+
+from __future__ import with_statement
+
+def test_with_statement(test):
+    with test.fs.open("test", "w") as f:
+        f.write("hello world")
+
+    with test.fs.open("test") as f:
+        test.assertEqual("hello world", f.read())

test/gridfs16.py

@@ -0,0 +1,23 @@
+# Copyright 2009 10gen, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Some tests for the gridfs package that only work under Python >= 1.6.
+"""
+
+def test_with_statement(test):
+    with test.fs.open("test", "w") as f:
+        f.write("hello world")
+
+    with test.fs.open("test") as f:
+        test.assertEqual("hello world", f.read())

test/test_gridfs.py

@@ -204,5 +204,16 @@ class TestGridfs(unittest.TestCase):
         self.assertEqual(f.read(), "hello")
         f.close()
 
+    # NOTE I do recognize how gross this is. There is no good way to test the
+    # with statement because it is a syntax error in older python versions.
+    # One option would be to use eval and skip the test if it is a syntax
+    # error.
+    if sys.version_info[:2] == (2, 5):
+        import gridfs15
+        test_with_statement = gridfs15.test_with_statement
+    elif sys.version_info[:3] >= (2, 6, 0):
+        import gridfs16
+        test_with_statement = gridfs16.test_with_statement
+
 if __name__ == "__main__":
     unittest.main()