Implement walk/walkfiles/walkdirs in WrapFS.

This allows FS subclasses to provide more efficient implementations, and have
them transparently used through various wrapper classes.


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

3 files changed, 61 insertions, 2 deletions

diff --git a/fs/__init__.py b/fs/__init__.py
index 570c1ed..ab49c7b 100644
--- a/fs/__init__.py
+++ b/fs/__init__.py
@@ -15,7 +15,7 @@ implementations of this interface such as:
 
 """
 
-__version__ = "0.4.0a8"
+__version__ = "0.4.0a9"
 __author__ = "Will McGugan (will@willmcgugan.com)"
 
 #  'base' imports * from 'path' and 'errors', so their
diff --git a/fs/base.py b/fs/base.py
index 847ec33..b5249dc 100644
--- a/fs/base.py
+++ b/fs/base.py
@@ -887,9 +887,10 @@ class FS(object):
         :param search: same as the walk method
         :param ignore_errors: ignore any errors reading the directory
         """
-        for p, files in self.walk(path, wildcard=wildcard, search=search, ignore_errors=ignore_errors):
+        for p, files in self.walk(path, dir_wildcard=wildcard, search=search, ignore_errors=ignore_errors):
             yield p
 
+
     def getsize(self, path):
         """Returns the size (in bytes) of a resource.
 
diff --git a/fs/wrapfs/__init__.py b/fs/wrapfs/__init__.py
index 46d1c48..2fafb08 100644
--- a/fs/wrapfs/__init__.py
+++ b/fs/wrapfs/__init__.py
@@ -278,6 +278,64 @@ class WrapFS(FS):
             yield (nm,info)
 
     @rewrite_errors
+    def walk(self,path="/",wildcard=None,dir_wildcard=None,search="breadth",ignore_errors=False):
+        if dir_wildcard is not None:
+            #  If there is a dir_wildcard, fall back to the default impl
+            #  that uses listdir().  Otherwise we run the risk of enumerating
+            #  lots of directories that will just be thrown away.
+            for item in super(WrapFS,self).walk(path,wildcard,dir_wildcard,search,ignore_errors):
+                yield item
+        #  Otherwise, the wrapped FS may provide a more efficient impl
+        #  which we can use directly.
+        else:
+            if wildcard is not None and not callable(wildcard):
+                wildcard_re = re.compile(fnmatch.translate(wildcard))
+                wildcard = lambda fn:bool (wildcard_re.match(fn))         
+            for (dirpath,filepaths) in self.wrapped_fs.walk(self._encode(path),search=search,ignore_errors=ignore_errors):
+                filepaths = [basename(self._decode(pathjoin(dirpath,p)))
+                                 for p in filepaths]
+                dirpath = abspath(self._decode(dirpath))
+                if wildcard is not None:
+                    filepaths = [p for p in filepaths if wildcard(p)]
+                yield (dirpath,filepaths)
+
+    @rewrite_errors
+    def walkfiles(self,path="/",wildcard=None,dir_wildcard=None,search="breadth",ignore_errors=False):
+        if dir_wildcard is not None:
+            #  If there is a dir_wildcard, fall back to the default impl
+            #  that uses listdir().  Otherwise we run the risk of enumerating
+            #  lots of directories that will just be thrown away.
+            for item in super(WrapFS,self).walkfiles(path,wildcard,dir_wildcard,search,ignore_errors):
+                yield item
+        #  Otherwise, the wrapped FS may provide a more efficient impl
+        #  which we can use directly.
+        else:
+            if wildcard is not None and not callable(wildcard):
+                wildcard_re = re.compile(fnmatch.translate(wildcard))
+                wildcard = lambda fn:bool (wildcard_re.match(fn))         
+            for filepath in self.wrapped_fs.walkfiles(self._encode(path),search=search,ignore_errors=ignore_errors):
+                filepath = abspath(self._decode(filepath))
+                if wildcard is not None:
+                    if not wildcard(basename(filepath)):
+                        continue
+                yield filepath
+
+    @rewrite_errors
+    def walkdirs(self,path="/",wildcard=None,search="breadth",ignore_errors=False):
+        if wildcard is not None:
+            #  If there is a wildcard, fall back to the default impl
+            #  that uses listdir().  Otherwise we run the risk of enumerating
+            #  lots of directories that will just be thrown away.
+            for item in super(WrapFS,self).walkdirs(path,wildcard,search,ignore_errors):
+                yield item
+        #  Otherwise, the wrapped FS may provide a more efficient impl
+        #  which we can use directly.
+        else:
+            for dirpath in self.wrapped_fs.walkdirs(self._encode(path),search=search,ignore_errors=ignore_errors):
+                yield abspath(self._decode(dirpath))
+
+
+    @rewrite_errors
     def makedir(self, path, *args, **kwds):
         return self.wrapped_fs.makedir(self._encode(path),*args,**kwds)