Added closing context manager to files returned by open

git-svn-id: http://pyfilesystem.googlecode.com/svn/trunk@796 67cdc799-7952-0410-af00-57a81ceafa0f

1 files changed, 26 insertions, 24 deletions

diff --git a/fs/zipfs.py b/fs/zipfs.py
index 35d5aaf..968b0e2 100644
--- a/fs/zipfs.py
+++ b/fs/zipfs.py
@@ -8,6 +8,7 @@ A FS object that represents the contents of a Zip file
 
 import datetime
 import os.path
+from contextlib import closing
 
 from fs.base import *
 from fs.path import *
@@ -32,7 +33,6 @@ class ZipNotFoundError(CreateFailedError):
 
 
 class _TempWriteFile(object):
-
     """Proxies a file object and calls a callback when the file is closed."""
 
     def __init__(self, fs, filename, close_callback):
@@ -50,13 +50,18 @@ class _TempWriteFile(object):
     def close(self):
         self._file.close()
         self.close_callback(self.filename)
-        
+
     def flush(self):
         self._file.flush()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.close()
 
-class _ExceptionProxy(object):
 
+class _ExceptionProxy(object):
     """A placeholder for an object that may no longer be used."""
 
     def __getattr__(self, name):
@@ -70,9 +75,8 @@ class _ExceptionProxy(object):
 
 
 class ZipFS(FS):
-
     """A FileSystem that represents a zip file."""
-    
+
     _meta = { 'thread_safe' : True,
               'virtual' : False,
               'read_only' : False,
@@ -107,27 +111,27 @@ class ZipFS(FS):
             raise ValueError("mode must be 'r', 'w' or 'a'")
 
         self.zip_mode = mode
-        self.encoding = encoding                
-        
+        self.encoding = encoding
+
         if isinstance(zip_file, basestring):
             zip_file = os.path.expanduser(os.path.expandvars(zip_file))
             zip_file = os.path.normpath(os.path.abspath(zip_file))
             self._zip_file_string = True
         else:
             self._zip_file_string = False
-        
+
         try:
             self.zf = ZipFile(zip_file, mode, compression_type, allow_zip_64)
-        except BadZipfile, bzf:            
+        except BadZipfile, bzf:
             raise ZipOpenError("Not a zip file or corrupt (%s)" % str(zip_file),
                                details=bzf)
-        except IOError, ioe:            
+        except IOError, ioe:
             if str(ioe).startswith('[Errno 22] Invalid argument'):
                 raise ZipOpenError("Not a zip file or corrupt (%s)" % str(zip_file),
                                    details=ioe)
             raise ZipNotFoundError("Zip file not found (%s)" % str(zip_file),
                                   details=ioe)
-                
+
         self.zip_path = str(zip_file)
         self.temp_fs = None
         if mode in 'wa':
@@ -136,8 +140,8 @@ class ZipFS(FS):
         self._path_fs = MemoryFS()
         if mode in 'ra':
             self._parse_resource_list()
-            
-        self.read_only = mode == 'r'        
+
+        self.read_only = mode == 'r'
 
     def __str__(self):
         return "<ZipFS: %s>" % self.zip_path
@@ -149,7 +153,7 @@ class ZipFS(FS):
         if PY3:
             return path
         return path.decode(self.encoding)
-    
+
     def _encode_path(self, path):
         if PY3:
             return path
@@ -172,11 +176,10 @@ class ZipFS(FS):
             f = self._path_fs.open(path, 'w')
             f.close()
 
-    def getmeta(self, meta_name, default=NoDefaultMeta):        
+    def getmeta(self, meta_name, default=NoDefaultMeta):
         if meta_name == 'read_only':
             return self.read_only
         return super(ZipFS, self).getmeta(meta_name, default)
-        
 
     def close(self):
         """Finalizes the zip file so that it can be read.
@@ -187,8 +190,8 @@ class ZipFS(FS):
             self.zf = _ExceptionProxy()
 
     @synchronize
-    def open(self, path, mode="r", **kwargs):        
-        path = normpath(relpath(path))        
+    def open(self, path, mode="r", **kwargs):
+        path = normpath(relpath(path))
 
         if 'r' in mode:
             if self.zip_mode not in 'ra':
@@ -202,10 +205,10 @@ class ZipFS(FS):
                     contents = self.zf.read(self._encode_path(path))
             except KeyError:
                 raise ResourceNotFoundError(path)
-            return StringIO(contents)
+            return closing(StringIO(contents))
 
-        if 'w' in mode:            
-            if self.zip_mode not in 'wa':                
+        if 'w' in mode:
+            if self.zip_mode not in 'wa':
                 raise OperationFailedError("open file",
                                            path=path,
                                            msg="2 Zip file must be opened for writing ('w') or appending ('a')")
@@ -215,7 +218,6 @@ class ZipFS(FS):
 
             self._add_resource(path)
             f = _TempWriteFile(self.temp_fs, path, self._on_write_close)
-
             return f
 
         raise ValueError("Mode must contain be 'r' or 'w'")
@@ -238,8 +240,8 @@ class ZipFS(FS):
         sys_path = self.temp_fs.getsyspath(filename)
         self.zf.write(sys_path, self._encode_path(filename))
 
-    def desc(self, path):        
-        return "%s in zip file %s" % (path, self.zip_path)        
+    def desc(self, path):
+        return "%s in zip file %s" % (path, self.zip_path)
 
     def isdir(self, path):
         return self._path_fs.isdir(path)