diff --git a/gridfs/__init__.py b/gridfs/__init__.py index 2dc41be89..ecff4c84f 100644 --- a/gridfs/__init__.py +++ b/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 diff --git a/gridfs/grid_file.py b/gridfs/grid_file.py index 63281f4cd..c68deb749 100644 --- a/gridfs/grid_file.py +++ b/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 diff --git a/test/gridfs15.py b/test/gridfs15.py new file mode 100644 index 000000000..be19bb009 --- /dev/null +++ b/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()) diff --git a/test/gridfs16.py b/test/gridfs16.py new file mode 100644 index 000000000..253291c31 --- /dev/null +++ b/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()) diff --git a/test/test_gridfs.py b/test/test_gridfs.py index c0f1e482e..bf83fbeff 100644 --- a/test/test_gridfs.py +++ b/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()