From 602250fdf9515a8127cd567d79afa8134d4cf923 Mon Sep 17 00:00:00 2001 From: Michael Williamson Date: Fri, 4 Jan 2013 23:39:48 +0000 Subject: [PATCH] Turn SFTPFile into a context manager --- paramiko/sftp_file.py | 6 ++++++ tests/test_sftp.py | 11 +++++++++++ 2 files changed, 17 insertions(+) diff --git a/paramiko/sftp_file.py b/paramiko/sftp_file.py index 8c5c7ac..c9fc76f 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 2eadabc..39c8aa8 100755 --- a/tests/test_sftp.py +++ b/tests/test_sftp.py @@ -188,6 +188,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.