[FIX] Support walk() in file storage (#152)

* Fix arg name in mime_type()

* Add new APIs for walk

* Add a comment for invalidating cache

Author: gaya3-zipstack

src/unstract/sdk/__init__.py

@@ -1,4 +1,5 @@
-__version__ = "0.56.0rc1"
+__version__ = "0.56.0rc2"
+
 
 def get_sdk_version():
     """Returns the SDK version."""

src/unstract/sdk/file_storage/impl.py

@@ -383,3 +383,22 @@ def guess_extension(self, path: str) -> str:
             file_type = filetype.guess(sample_contents)
             file_extension = file_type.EXTENSION
         return file_extension
+
+    def walk(self, path: str, max_depth=None, topdown=True):
+        """Walks the dir in the path and returns the list of files/dirs.
+
+        Args:
+            path (str): Root to recurse into
+            maxdepth (int): Maximum recursion depth. None means limitless,
+            but not recommended
+                on link-based file-systems.
+            topdown (bool): Whether to walk the directory tree from the top
+            downwards or from
+                the bottom upwards.
+
+        Returns:
+            Iterator containing the list of files and folders
+        """
+        # Invalidating cache explicitly to avoid any stale listing
+        self.fs.invalidate_cache(path=path)
+        return self.fs.walk(path, maxdepth=max_depth, topdown=topdown)

src/unstract/sdk/file_storage/interface.py

@@ -125,3 +125,7 @@ def yaml_load(
     @abstractmethod
     def guess_extension(self, path: str) -> str:
         pass
+
+    @abstractmethod
+    def walk(self, path: str):
+        pass

tests/test_file_storage.py

@@ -771,3 +771,72 @@ def test_get_storage(storage_type, env_name, expected):
     file_storage = EnvHelper.get_storage(storage_type, env_name)
     assert file_storage.provider == expected
     print(file_storage)
+
+
+@pytest.mark.parametrize(
+    "storage_type, env_name, path",
+    [
+        (
+            StorageType.PERMANENT,
+            "TEST_PERMANENT_STORAGE",
+            "fsspec-test",
+        ),
+        (
+            StorageType.SHARED_TEMPORARY,
+            "TEST_TEMPORARY_STORAGE",
+            "unstract/execution/mock_org/"
+            "13484b52-2127-48c2-b1a3-b517365346c3/"
+            "39fcdcba-90bb-44ce-9446-67253adcb4d7/COPY_TO_FOLDER",
+        ),
+    ],
+)
+def test_dir_walk(storage_type, env_name, path):
+    file_storage = EnvHelper.get_storage(storage_type, env_name)
+    try:
+        root, dirs, files = next(file_storage.walk(path))
+    except StopIteration:
+        return []
+    for dir_name in dirs:
+        print(dir_name)
+    for file_name in files:
+        print(file_name)
+    if storage_type == StorageType.PERMANENT:
+        assert len(files) > 0
+    elif storage_type == StorageType.SHARED_TEMPORARY:
+        assert len(files) == 0
+
+
+def list_print_dir(file_storage, path, iter_num):
+    print(f"PATH: {path}")
+    print(f"\nItertion: {iter_num}")
+    try:
+        root, dirs, files = next(file_storage.walk(path))
+    except StopIteration:
+        return []
+    for dir_name in dirs:
+        print(dir_name)
+    for file_name in files:
+        print(file_name)
+    print(f"Files: {files}")
+
+
+@pytest.mark.parametrize(
+    "storage_type, env_name, path",
+    [
+        (
+            StorageType.SHARED_TEMPORARY,
+            "TEST_TEMPORARY_STORAGE",
+            "unstract/execution/mock_org/"
+            "13484b52-2127-48c2-b1a3-b517365346c3/b"
+            "f7b3d81-d0aa-4e9e-883d-25dd0f3a6466/COPY_TO_FOLDER",
+        ),
+    ],
+)
+def test_dir_ls(storage_type, env_name, path):
+    new_file = os.path.join(path, "tmp.txt")
+    file_storage = EnvHelper.get_storage(storage_type, env_name)
+    if file_storage.exists(new_file):
+        file_storage.rm(new_file)
+    list_print_dir(file_storage, path, "1")
+    file_storage.write(new_file, "w", data="Hello")
+    list_print_dir(file_storage, path, "2")