updated archivefs

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

1 files changed, 105 insertions, 12 deletions

diff --git a/fs/contrib/archivefs.py b/fs/contrib/archivefs.py
index b7308bd..863f62c 100644
--- a/fs/contrib/archivefs.py
+++ b/fs/contrib/archivefs.py
@@ -15,11 +15,33 @@ from fs.base import *
 from fs.path import *
 from fs.errors import *
 from fs.filelike import StringIO
+from fs import mountfs
 
 import libarchive
 
 ENCODING = libarchive.ENCODING
 
+
+class SizeUpdater(object):
+    '''A file-like object to allow writing to a file within the archive. When closed
+    this object will update the archive entry's size within the archive.'''
+    def __init__(self, entry, stream):
+        self.entry = entry
+        self.stream = stream
+        self.size = 0
+
+    def __del__(self):
+        self.close()
+
+    def write(self, data):
+        self.size += len(data)
+        self.stream.write(data)
+
+    def close(self):
+        self.stream.close()
+        self.entry.size = self.size
+
+
 class ArchiveFS(FS):
     """A FileSystem that represents an archive supported by libarchive."""
 
@@ -33,12 +55,19 @@ class ArchiveFS(FS):
              }
 
     def __init__(self, f, mode='r', format=None, thread_synchronize=True):
-        """Create a FS that maps on to a zip file.
+        """Create a FS that maps on to an archive file.
 
-        :param path: a (system) path, or a file-like object
+        :param f: a (system) path, or a file-like object
+        :param format: required for 'w' mode. The archive format ('zip, 'tar', etc)
         :param thread_synchronize: set to True (default) to enable thread-safety
         """
         super(ArchiveFS, self).__init__(thread_synchronize=thread_synchronize)
+        if isinstance(f, basestring):
+            self.fileobj = None
+            self.root_path = f
+        else:
+            self.fileobj = f
+            self.root_path = getattr(f, 'name', None)
         self.contents = PathMap()
         self.archive = libarchive.SeekableArchive(f, format=format, mode=mode)
         if mode == 'r':
@@ -51,10 +80,10 @@ class ArchiveFS(FS):
                         self.contents[part] = libarchive.Entry(pathname=part, mode=stat.S_IFDIR, size=0, mtime=item.mtime)
 
     def __str__(self):
-        return "<ArchiveFS>"
+        return "<ArchiveFS: %s>" % self.root_path
 
     def __unicode__(self):
-        return u"<ArchiveFS>"
+        return u"<ArchiveFS: %s>" % self.root_path
 
     def getmeta(self, meta_name, default=NoDefaultMeta):
         if meta_name == 'read_only':
@@ -62,17 +91,21 @@ class ArchiveFS(FS):
         return super(ZipFS, self).getmeta(meta_name, default)
 
     def close(self):
+        if getattr(self, 'archive', None) is None:
+            return
         self.archive.close()
 
     @synchronize
     def open(self, path, mode="r", **kwargs):
         path = normpath(relpath(path))
-        if mode not in ('r', 'w', 'wb'):
+        if 'a' in mode:
             raise Exception('Unsupported mode ' + mode)
         if 'r' in mode:
             return self.archive.readstream(path)
         else:
-            return self.archive.writestream(path)
+            entry = self.archive.entry_class(pathname=path, mode=stat.S_IFREG, size=0, mtime=time.time())
+            self.contents[path] = entry
+            return SizeUpdater(entry, self.archive.writestream(path))
 
     @synchronize
     def getcontents(self, path, mode="rb"):
@@ -86,14 +119,19 @@ class ArchiveFS(FS):
 
     def isdir(self, path):
         info = self.getinfo(path)
-        return stat.S_ISDIR(info.get('mode', 0))
+        # Don't use stat.S_ISDIR, it won't work when mode == S_IFREG | S_IFDIR.
+        return info.get('mode', 0) & stat.S_IFDIR == stat.S_IFDIR
 
     def isfile(self, path):
         info = self.getinfo(path)
-        return stat.S_ISREG(info.get('mode', 0))
+        # Don't use stat.S_ISREG, it won't work when mode == S_IFREG | S_IFDIR.
+        return info.get('mode', 0) & stat.S_IFREG == stat.S_IFREG
 
     def exists(self, path):
         path = normpath(path).lstrip('/')
+        if path == '':
+            # We are being asked about root (the archive itself)
+            return True
         return path in self.contents
 
     def listdir(self, path="/", wildcard=None, full=False, absolute=False, dirs_only=False, files_only=False):
@@ -101,6 +139,7 @@ class ArchiveFS(FS):
 
     def makedir(self, dirname, recursive=False, allow_recreate=False):
         entry = self.archive.entry_class(pathname=dirname, mode=stat.S_IFDIR, size=0, mtime=time.time())
+        self.contents[dirname] = entry
         self.archive.write(entry)
 
     @synchronize
@@ -108,20 +147,74 @@ class ArchiveFS(FS):
         if not self.exists(path):
             raise ResourceNotFoundError(path)
         path = normpath(path).lstrip('/')
-        info = { 'size': 0 }
-        try:
+        if path == '':
+            # We are being asked about root (the archive itself)
+            if self.root_path:
+                st = os.stat(self.root_path)
+            elif hasattr(self.fileobj, 'fileno'):
+                st = os.fstat(self.fileobj.fileno())
+            else:
+                raise Exception('Could not stat archive.')
+            info = dict((k, getattr(st, k)) for k in dir(st) if k.startswith('st_'))
+            for name, longname in (
+                    ('st_ctime', 'created_time'), ('st_atime', 'accessed_time'),
+                    ('st_mtime', 'modified_time'),
+                ):
+                if name in info:
+                    t = info.pop(name)
+                    if t:
+                        info[long_name] = datetime.datetime.fromtimestamp(t)
+            info['size'] = info.pop('st_size')
+            # Masquerade as a directory.
+            info['mode'] |= stat.S_IFDIR
+        else:
+            info = { 'size': 0 }
             entry = self.contents.get(path)
             for attr in dir(entry):
                 if attr.startswith('_'):
                     continue
                 elif attr == 'mtime':
                     info['created_time'] = datetime.datetime.fromtimestamp(entry.mtime)
+                elif attr == 'mode':
+                    info['st_mode'] = entry.mode
                 else:
                     info[attr] = getattr(entry, attr)
-        except KeyError:
-            pass
         return info
 
+
+class ArchiveMountFS(mountfs.MountFS):
+    '''A subclass of MountFS that automatically identifies archives. Once identified
+    archives are mounted in place of the archive file.'''
+    def __init__(self, root, **kwargs):
+        super(ArchiveMountFS, self).__init__(**kwargs)
+        self.root_path = root_path
+        self.mountdir('/', root)
+
+    def ismount(self, path):
+        try:
+            object = self.mount_tree[path]
+        except KeyError:
+            return False
+        return type(object) is mountfs.MountFS.DirMount
+
+    def _delegate(self, path):
+        for ppath in recursepath(path)[1:]:
+            # Don't mount again...
+            if self.ismount(ppath):
+                break
+            if libarchive.is_archive_name(ppath):
+                # It looks like an archive, try mounting it.
+                full_path = pathjoin(self.root_path, relpath(ppath))
+                try:
+                    self.mountdir(ppath, ArchiveFS(full_path, 'r'))
+                except:
+                    pass # Must NOT have been an archive after all
+                # Stop recursing path, we support just one archive per path!
+                # No nested archives yet!
+                break
+        return super(ArchiveMountFS, self)._delegate(path)
+
+
 def main():
     ArchiveFS()