diff --git a/jedi/file_io.py b/jedi/file_io.py index 012bff41..60895fff 100644 --- a/jedi/file_io.py +++ b/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): diff --git a/test/test_file_io.py b/test/test_file_io.py new file mode 100644 index 00000000..a24b7edf --- /dev/null +++ b/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