Use 'with' for opening most file and SFTPFIle objects

1 files changed, 129 insertions, 149 deletions

diff --git a/tests/test_sftp_big.py b/tests/test_sftp_big.py
index dc9cba93..6870c6b4 100644
--- a/tests/test_sftp_big.py
+++ b/tests/test_sftp_big.py
@@ -66,9 +66,8 @@ class BigSFTPTest (unittest.TestCase):
         numfiles = 100
         try:
             for i in range(numfiles):
-                f = sftp.open('%s/file%d.txt' % (FOLDER, i), 'w', 1)
-                f.write('this is file #%d.\n' % i)
-                f.close()
+                with sftp.open('%s/file%d.txt' % (FOLDER, i), 'w', 1) as f:
+                    f.write('this is file #%d.\n' % i)
                 sftp.chmod('%s/file%d.txt' % (FOLDER, i), o660)
 
             # now make sure every file is there, by creating a list of filenmes
@@ -76,9 +75,8 @@ class BigSFTPTest (unittest.TestCase):
             numlist = list(range(numfiles))
             while len(numlist) > 0:
                 r = numlist[random.randint(0, len(numlist) - 1)]
-                f = sftp.open('%s/file%d.txt' % (FOLDER, r))
-                self.assertEqual(f.readline(), 'this is file #%d.\n' % r)
-                f.close()
+                with sftp.open('%s/file%d.txt' % (FOLDER, r)) as f:
+                    self.assertEqual(f.readline(), 'this is file #%d.\n' % r)
                 numlist.remove(r)
         finally:
             for i in range(numfiles):
@@ -95,12 +93,11 @@ class BigSFTPTest (unittest.TestCase):
         kblob = (1024 * 'x')
         start = time.time()
         try:
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'w')
-            for n in range(1024):
-                f.write(kblob)
-                if n % 128 == 0:
-                    sys.stderr.write('.')
-            f.close()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'w') as f:
+                for n in range(1024):
+                    f.write(kblob)
+                    if n % 128 == 0:
+                        sys.stderr.write('.')
             sys.stderr.write(' ')
 
             self.assertEqual(sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
@@ -108,11 +105,10 @@ class BigSFTPTest (unittest.TestCase):
             sys.stderr.write('%ds ' % round(end - start))
             
             start = time.time()
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'r')
-            for n in range(1024):
-                data = f.read(1024)
-                self.assertEqual(data, kblob)
-            f.close()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'r') as f:
+                for n in range(1024):
+                    data = f.read(1024)
+                    self.assertEqual(data, kblob)
 
             end = time.time()
             sys.stderr.write('%ds ' % round(end - start))
@@ -127,13 +123,12 @@ class BigSFTPTest (unittest.TestCase):
         kblob = bytes().join([struct.pack('>H', n) for n in range(512)])
         start = time.time()
         try:
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'wb')
-            f.set_pipelined(True)
-            for n in range(1024):
-                f.write(kblob)
-                if n % 128 == 0:
-                    sys.stderr.write('.')
-            f.close()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
+                f.set_pipelined(True)
+                for n in range(1024):
+                    f.write(kblob)
+                    if n % 128 == 0:
+                        sys.stderr.write('.')
             sys.stderr.write(' ')
 
             self.assertEqual(sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
@@ -141,22 +136,21 @@ class BigSFTPTest (unittest.TestCase):
             sys.stderr.write('%ds ' % round(end - start))
             
             start = time.time()
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'rb')
-            f.prefetch()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
+                f.prefetch()
 
-            # read on odd boundaries to make sure the bytes aren't getting scrambled
-            n = 0
-            k2blob = kblob + kblob
-            chunk = 629
-            size = 1024 * 1024
-            while n < size:
-                if n + chunk > size:
-                    chunk = size - n
-                data = f.read(chunk)
-                offset = n % 1024
-                self.assertEqual(data, k2blob[offset:offset + chunk])
-                n += chunk
-            f.close()
+                # read on odd boundaries to make sure the bytes aren't getting scrambled
+                n = 0
+                k2blob = kblob + kblob
+                chunk = 629
+                size = 1024 * 1024
+                while n < size:
+                    if n + chunk > size:
+                        chunk = size - n
+                    data = f.read(chunk)
+                    offset = n % 1024
+                    self.assertEqual(data, k2blob[offset:offset + chunk])
+                    n += chunk
 
             end = time.time()
             sys.stderr.write('%ds ' % round(end - start))
@@ -167,13 +161,12 @@ class BigSFTPTest (unittest.TestCase):
         sftp = get_sftp()
         kblob = bytes().join([struct.pack('>H', n) for n in range(512)])
         try:
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'wb')
-            f.set_pipelined(True)
-            for n in range(1024):
-                f.write(kblob)
-                if n % 128 == 0:
-                    sys.stderr.write('.')
-            f.close()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
+                f.set_pipelined(True)
+                for n in range(1024):
+                    f.write(kblob)
+                    if n % 128 == 0:
+                        sys.stderr.write('.')
             sys.stderr.write(' ')
             
             self.assertEqual(sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
@@ -182,20 +175,19 @@ class BigSFTPTest (unittest.TestCase):
             k2blob = kblob + kblob
             chunk = 793
             for i in range(10):
-                f = sftp.open('%s/hongry.txt' % FOLDER, 'rb')
-                f.prefetch()
-                base_offset = (512 * 1024) + 17 * random.randint(1000, 2000)
-                offsets = [base_offset + j * chunk for j in range(100)]
-                # randomly seek around and read them out
-                for j in range(100):
-                    offset = offsets[random.randint(0, len(offsets) - 1)]
-                    offsets.remove(offset)
-                    f.seek(offset)
-                    data = f.read(chunk)
-                    n_offset = offset % 1024
-                    self.assertEqual(data, k2blob[n_offset:n_offset + chunk])
-                    offset += chunk
-                f.close()
+                with sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
+                    f.prefetch()
+                    base_offset = (512 * 1024) + 17 * random.randint(1000, 2000)
+                    offsets = [base_offset + j * chunk for j in range(100)]
+                    # randomly seek around and read them out
+                    for j in range(100):
+                        offset = offsets[random.randint(0, len(offsets) - 1)]
+                        offsets.remove(offset)
+                        f.seek(offset)
+                        data = f.read(chunk)
+                        n_offset = offset % 1024
+                        self.assertEqual(data, k2blob[n_offset:n_offset + chunk])
+                        offset += chunk
             end = time.time()
             sys.stderr.write('%ds ' % round(end - start))
         finally:
@@ -205,13 +197,12 @@ class BigSFTPTest (unittest.TestCase):
         sftp = get_sftp()
         kblob = bytes().join([struct.pack('>H', n) for n in range(512)])
         try:
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'wb')
-            f.set_pipelined(True)
-            for n in range(1024):
-                f.write(kblob)
-                if n % 128 == 0:
-                    sys.stderr.write('.')
-            f.close()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
+                f.set_pipelined(True)
+                for n in range(1024):
+                    f.write(kblob)
+                    if n % 128 == 0:
+                        sys.stderr.write('.')
             sys.stderr.write(' ')
 
             self.assertEqual(sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
@@ -220,21 +211,20 @@ class BigSFTPTest (unittest.TestCase):
             k2blob = kblob + kblob
             chunk = 793
             for i in range(10):
-                f = sftp.open('%s/hongry.txt' % FOLDER, 'rb')
-                base_offset = (512 * 1024) + 17 * random.randint(1000, 2000)
-                # make a bunch of offsets and put them in random order
-                offsets = [base_offset + j * chunk for j in range(100)]
-                readv_list = []
-                for j in range(100):
-                    o = offsets[random.randint(0, len(offsets) - 1)]
-                    offsets.remove(o)
-                    readv_list.append((o, chunk))
-                ret = f.readv(readv_list)
-                for i in range(len(readv_list)):
-                    offset = readv_list[i][0]
-                    n_offset = offset % 1024
-                    self.assertEqual(next(ret), k2blob[n_offset:n_offset + chunk])
-                f.close()
+                with sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
+                    base_offset = (512 * 1024) + 17 * random.randint(1000, 2000)
+                    # make a bunch of offsets and put them in random order
+                    offsets = [base_offset + j * chunk for j in range(100)]
+                    readv_list = []
+                    for j in range(100):
+                        o = offsets[random.randint(0, len(offsets) - 1)]
+                        offsets.remove(o)
+                        readv_list.append((o, chunk))
+                    ret = f.readv(readv_list)
+                    for i in range(len(readv_list)):
+                        offset = readv_list[i][0]
+                        n_offset = offset % 1024
+                        self.assertEqual(next(ret), k2blob[n_offset:n_offset + chunk])
             end = time.time()
             sys.stderr.write('%ds ' % round(end - start))
         finally:
@@ -248,29 +238,26 @@ class BigSFTPTest (unittest.TestCase):
         sftp = get_sftp()
         kblob = (1024 * 'x')
         try:
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'w')
-            f.set_pipelined(True)
-            for n in range(1024):
-                f.write(kblob)
-                if n % 128 == 0:
-                    sys.stderr.write('.')
-            f.close()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'w') as f:
+                f.set_pipelined(True)
+                for n in range(1024):
+                    f.write(kblob)
+                    if n % 128 == 0:
+                        sys.stderr.write('.')
             sys.stderr.write(' ')
 
             self.assertEqual(sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
 
             for i in range(10):
-                f = sftp.open('%s/hongry.txt' % FOLDER, 'r')
+                with sftp.open('%s/hongry.txt' % FOLDER, 'r') as f:
+                    f.prefetch()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'r') as f:
                 f.prefetch()
-                f.close()
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'r')
-            f.prefetch()
-            for n in range(1024):
-                data = f.read(1024)
-                self.assertEqual(data, kblob)
-                if n % 128 == 0:
-                    sys.stderr.write('.')
-            f.close()
+                for n in range(1024):
+                    data = f.read(1024)
+                    self.assertEqual(data, kblob)
+                    if n % 128 == 0:
+                        sys.stderr.write('.')
             sys.stderr.write(' ')
         finally:
             sftp.remove('%s/hongry.txt' % FOLDER)
@@ -282,33 +269,31 @@ class BigSFTPTest (unittest.TestCase):
         sftp = get_sftp()
         kblob = bytes().join([struct.pack('>H', n) for n in range(512)])
         try:
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'wb')
-            f.set_pipelined(True)
-            for n in range(1024):
-                f.write(kblob)
-                if n % 128 == 0:
-                    sys.stderr.write('.')
-            f.close()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
+                f.set_pipelined(True)
+                for n in range(1024):
+                    f.write(kblob)
+                    if n % 128 == 0:
+                        sys.stderr.write('.')
             sys.stderr.write(' ')
             
             self.assertEqual(sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
 
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'rb')
-            f.prefetch()
-            data = f.read(1024)
-            self.assertEqual(data, kblob)
-            
-            chunk_size = 793
-            base_offset = 512 * 1024
-            k2blob = kblob + kblob
-            chunks = [(base_offset + (chunk_size * i), chunk_size) for i in range(20)]
-            for data in f.readv(chunks):
-                offset = base_offset % 1024
-                self.assertEqual(chunk_size, len(data))
-                self.assertEqual(k2blob[offset:offset + chunk_size], data)
-                base_offset += chunk_size
-
-            f.close()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
+                f.prefetch()
+                data = f.read(1024)
+                self.assertEqual(data, kblob)
+
+                chunk_size = 793
+                base_offset = 512 * 1024
+                k2blob = kblob + kblob
+                chunks = [(base_offset + (chunk_size * i), chunk_size) for i in range(20)]
+                for data in f.readv(chunks):
+                    offset = base_offset % 1024
+                    self.assertEqual(chunk_size, len(data))
+                    self.assertEqual(k2blob[offset:offset + chunk_size], data)
+                    base_offset += chunk_size
+
             sys.stderr.write(' ')
         finally:
             sftp.remove('%s/hongry.txt' % FOLDER)
@@ -321,24 +306,22 @@ class BigSFTPTest (unittest.TestCase):
         sftp = get_sftp()
         kblob = bytes().join([struct.pack('>H', n) for n in range(512)])
         try:
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'wb')
-            f.set_pipelined(True)
-            for n in range(1024):
-                f.write(kblob)
-                if n % 128 == 0:
-                    sys.stderr.write('.')
-            f.close()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'wb') as f:
+                f.set_pipelined(True)
+                for n in range(1024):
+                    f.write(kblob)
+                    if n % 128 == 0:
+                        sys.stderr.write('.')
             sys.stderr.write(' ')
 
             self.assertEqual(sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
             
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'rb')
-            data = list(f.readv([(23 * 1024, 128 * 1024)]))
-            self.assertEqual(1, len(data))
-            data = data[0]
-            self.assertEqual(128 * 1024, len(data))
+            with sftp.open('%s/hongry.txt' % FOLDER, 'rb') as f:
+                data = list(f.readv([(23 * 1024, 128 * 1024)]))
+                self.assertEqual(1, len(data))
+                data = data[0]
+                self.assertEqual(128 * 1024, len(data))
             
-            f.close()
             sys.stderr.write(' ')
         finally:
             sftp.remove('%s/hongry.txt' % FOLDER)
@@ -350,9 +333,8 @@ class BigSFTPTest (unittest.TestCase):
         sftp = get_sftp()
         mblob = (1024 * 1024 * 'x')
         try:
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'w', 128 * 1024)
-            f.write(mblob)
-            f.close()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'w', 128 * 1024) as f:
+                f.write(mblob)
 
             self.assertEqual(sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
         finally:
@@ -367,21 +349,19 @@ class BigSFTPTest (unittest.TestCase):
         t.packetizer.REKEY_BYTES = 512 * 1024
         k32blob = (32 * 1024 * 'x')
         try:
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'w', 128 * 1024)
-            for i in range(32):
-                f.write(k32blob)
-            f.close()
-            
+            with sftp.open('%s/hongry.txt' % FOLDER, 'w', 128 * 1024) as f:
+                for i in range(32):
+                    f.write(k32blob)
+
             self.assertEqual(sftp.stat('%s/hongry.txt' % FOLDER).st_size, 1024 * 1024)
             self.assertNotEqual(t.H, t.session_id)
             
             # try to read it too.
-            f = sftp.open('%s/hongry.txt' % FOLDER, 'r', 128 * 1024)
-            f.prefetch()
-            total = 0
-            while total < 1024 * 1024:
-                total += len(f.read(32 * 1024))
-            f.close()
+            with sftp.open('%s/hongry.txt' % FOLDER, 'r', 128 * 1024) as f:
+                f.prefetch()
+                total = 0
+                while total < 1024 * 1024:
+                    total += len(f.read(32 * 1024))
         finally:
             sftp.remove('%s/hongry.txt' % FOLDER)
             t.packetizer.REKEY_BYTES = pow(2, 30)