Added FolderIO.walk and FolderIO.get_base_name

jedi/file_io.py

@@ -21,6 +21,9 @@ class AbstractFolderIO(object):
 
 
 class FolderIO(AbstractFolderIO):
+    def get_base_name(self):
+        return os.path.basename(self.path)
+
     def list(self):
         return os.listdir(self.path)
 
@@ -30,6 +33,24 @@ class FolderIO(AbstractFolderIO):
     def get_parent_folder(self):
         return FolderIO(os.path.dirname(self.path))
 
+    def walk(self):
+        for root, dirs, files in os.walk(self.path):
+            root_folder_io = FolderIO(root)
+            original_folder_ios = [FolderIO(os.path.join(root, d)) for d in dirs]
+            modified_folder_ios = list(original_folder_ios)
+            yield (
+                root_folder_io,
+                modified_folder_ios,
+                [FileIO(os.path.join(root, f)) for f in files],
+            )
+            modified_iterator = iter(modified_folder_ios)
+            current = next(modified_iterator, None)
+            for i, folder_io in enumerate(original_folder_ios):
+                if current is folder_io:
+                    current = next(modified_iterator, None)
+                else:
+                    del dirs[i]
+
 
 class FileIOFolderMixin(object):
     def get_parent_folder(self):

test/test_file_io.py

@@ -0,0 +1,19 @@
+from os.path import join
+from jedi.file_io import FolderIO
+from test.helpers import get_example_dir
+
+
+def test_folder_io_walk():
+    root_dir = get_example_dir('namespace_package')
+    iterator = FolderIO(root_dir).walk()
+    root, folder_ios, file_ios = next(iterator)
+    assert {f.path for f in folder_ios} == {join(root_dir, 'ns1'), join(root_dir, 'ns2')}
+    for f in list(folder_ios):
+        if f.path.endswith('ns1'):
+            folder_ios.remove(f)
+
+    root, folder_ios, file_ios = next(iterator)
+    assert folder_ios
+    assert root.path == join(root_dir, 'ns2')
+    folder_ios.clear()
+    assert next(iterator, None) is None