[HfFileSystem] Optimize maxdepth: do less /tree calls in glob() (#3389)

* optimize maxdepth

* raise an error on invalid maxdepth

* mypy

* test

* fix test

Author: lhoestq

src/huggingface_hub/hf_file_system.py

@@ -386,6 +386,7 @@ def _ls_tree(
         refresh: bool = False,
         revision: Optional[str] = None,
         expand_info: bool = False,
+        maxdepth: Optional[int] = None,
     ):
         resolved_path = self.resolve_path(path, revision=revision)
         path = resolved_path.unresolve()
@@ -405,19 +406,25 @@ def _ls_tree(
             if recursive:
                 # Use BFS to traverse the cache and build the "recursive "output
                 # (The Hub uses a so-called "tree first" strategy for the tree endpoint but we sort the output to follow the spec so the result is (eventually) the same)
+                depth = 2
                 dirs_to_visit = deque(
-                    [path_info for path_info in cached_path_infos if path_info["type"] == "directory"]
+                    [(depth, path_info) for path_info in cached_path_infos if path_info["type"] == "directory"]
                 )
                 while dirs_to_visit:
-                    dir_info = dirs_to_visit.popleft()
-                    if dir_info["name"] not in self.dircache:
-                        dirs_not_in_dircache.append(dir_info["name"])
-                    else:
-                        cached_path_infos = self.dircache[dir_info["name"]]
-                        out.extend(cached_path_infos)
-                        dirs_to_visit.extend(
-                            [path_info for path_info in cached_path_infos if path_info["type"] == "directory"]
-                        )
+                    depth, dir_info = dirs_to_visit.popleft()
+                    if maxdepth is None or depth <= maxdepth:
+                        if dir_info["name"] not in self.dircache:
+                            dirs_not_in_dircache.append(dir_info["name"])
+                        else:
+                            cached_path_infos = self.dircache[dir_info["name"]]
+                            out.extend(cached_path_infos)
+                            dirs_to_visit.extend(
+                                [
+                                    (depth + 1, path_info)
+                                    for path_info in cached_path_infos
+                                    if path_info["type"] == "directory"
+                                ]
+                            )
 
             dirs_not_expanded = []
             if expand_info:
@@ -436,6 +443,9 @@ def _ls_tree(
                     or common_prefix in chain(dirs_not_in_dircache, dirs_not_expanded)
                     else self._parent(common_prefix)
                 )
+                if maxdepth is not None:
+                    common_path_depth = common_path[len(path) :].count("/")
+                    maxdepth -= common_path_depth
                 out = [o for o in out if not o["name"].startswith(common_path + "/")]
                 for cached_path in self.dircache:
                     if cached_path.startswith(common_path + "/"):
@@ -448,6 +458,7 @@ def _ls_tree(
                         refresh=True,
                         revision=revision,
                         expand_info=expand_info,
+                        maxdepth=maxdepth,
                     )
                 )
         else:
@@ -460,9 +471,10 @@ def _ls_tree(
                 repo_type=resolved_path.repo_type,
             )
             for path_info in tree:
+                cache_path = root_path + "/" + path_info.path
                 if isinstance(path_info, RepoFile):
                     cache_path_info = {
-                        "name": root_path + "/" + path_info.path,
+                        "name": cache_path,
                         "size": path_info.size,
                         "type": "file",
                         "blob_id": path_info.blob_id,
@@ -472,15 +484,17 @@ def _ls_tree(
                     }
                 else:
                     cache_path_info = {
-                        "name": root_path + "/" + path_info.path,
+                        "name": cache_path,
                         "size": 0,
                         "type": "directory",
                         "tree_id": path_info.tree_id,
                         "last_commit": path_info.last_commit,
                     }
                 parent_path = self._parent(cache_path_info["name"])
                 self.dircache.setdefault(parent_path, []).append(cache_path_info)
-                out.append(cache_path_info)
+                depth = cache_path[len(path) :].count("/")
+                if maxdepth is None or depth <= maxdepth:
+                    out.append(cache_path_info)
         return out
 
     def walk(self, path: str, *args, **kwargs) -> Iterator[Tuple[str, List[str], List[str]]]:
@@ -547,19 +561,22 @@ def find(
         Returns:
             `Union[List[str], Dict[str, Dict[str, Any]]]`: List of paths or dict of file information.
         """
-        if maxdepth:
-            return super().find(
-                path, maxdepth=maxdepth, withdirs=withdirs, detail=detail, refresh=refresh, revision=revision, **kwargs
-            )
+        if maxdepth is not None and maxdepth < 1:
+            raise ValueError("maxdepth must be at least 1")
         resolved_path = self.resolve_path(path, revision=revision)
         path = resolved_path.unresolve()
         try:
-            out = self._ls_tree(path, recursive=True, refresh=refresh, revision=resolved_path.revision, **kwargs)
+            out = self._ls_tree(
+                path, recursive=True, refresh=refresh, revision=resolved_path.revision, maxdepth=maxdepth, **kwargs
+            )
         except EntryNotFoundError:
             # Path could be a file
-            if self.info(path, revision=revision, **kwargs)["type"] == "file":
-                out = {path: {}}
-            else:
+            try:
+                if self.info(path, revision=revision, **kwargs)["type"] == "file":
+                    out = {path: {}}
+                else:
+                    out = {}
+            except FileNotFoundError:
                 out = {}
         else:
             if not withdirs:

tests/test_hf_file_system.py

@@ -60,7 +60,8 @@ def setUp(self):
             repo_type="dataset",
         )
 
-        self.text_file = self.hf_path + "/data/text_data.txt"
+        self.text_file_path_in_repo = "data/text_data.txt"
+        self.text_file = self.hf_path + "/" + self.text_file_path_in_repo
 
     def tearDown(self):
         self.api.delete_repo(self.repo_id, repo_type="dataset")
@@ -421,6 +422,16 @@ def test_find_data_file_no_revision(self):
         files = self.hffs.find(self.text_file, detail=False)
         self.assertEqual(files, [self.text_file])
 
+    def test_find_maxdepth(self):
+        text_file_depth = self.text_file_path_in_repo.count("/") + 1
+        files = self.hffs.find(self.hf_path, detail=False, maxdepth=text_file_depth - 1)
+        self.assertNotIn(self.text_file, files)
+        files = self.hffs.find(self.hf_path, detail=False, maxdepth=text_file_depth)
+        self.assertIn(self.text_file, files)
+        # we do it again once the cache is updated
+        files = self.hffs.find(self.hf_path, detail=False, maxdepth=text_file_depth - 1)
+        self.assertNotIn(self.text_file, files)
+
     def test_read_bytes(self):
         data = self.hffs.read_bytes(self.text_file)
         self.assertEqual(data, b"dummy text data")