Merge pull request #127 from mwilliamson/sftp-file-context-manager

Turn SFTPFile into a context manager

2 files changed, 19 insertions, 0 deletions

diff --git a/paramiko/sftp_file.py b/paramiko/sftp_file.py
index 8c5c7aca..c9fc76f5 100644
--- a/paramiko/sftp_file.py
+++ b/paramiko/sftp_file.py
@@ -474,3 +474,9 @@ class SFTPFile (BufferedFile):
             x = self._saved_exception
             self._saved_exception = None
             raise x
+            
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
diff --git a/tests/test_sftp.py b/tests/test_sftp.py
index 2eadabcd..f95da69c 100755
--- a/tests/test_sftp.py
+++ b/tests/test_sftp.py
@@ -23,6 +23,8 @@ a real actual sftp server is contacted, and a new folder is created there to
 do test file operations in (so no existing files will be harmed).
 """
 
+from __future__ import with_statement
+
 from binascii import hexlify
 import logging
 import os
@@ -188,6 +190,17 @@ class SFTPTest (unittest.TestCase):
         finally:
             sftp.remove(FOLDER + '/duck.txt')
 
+    def test_3_sftp_file_can_be_used_as_context_manager(self):
+        """
+        verify that an opened file can be used as a context manager
+        """
+        try:
+            with sftp.open(FOLDER + '/duck.txt', 'w') as f:
+                f.write(ARTICLE)
+            self.assertEqual(sftp.stat(FOLDER + '/duck.txt').st_size, 1483)
+        finally:
+            sftp.remove(FOLDER + '/duck.txt')
+
     def test_4_append(self):
         """
         verify that a file can be opened for append, and tell() still works.