Repository: iterative/dvc
Pull request: tree: Introduce the new .ls() API
Resolves #5274 

Introduces a new `Tree.ls()` API for listings

master:
```
Importing 's3://dvc-temp/import-url' -> 'import-url'                  
To track the changes with git, run:
        git add import-url.dvc
dvc import-url s3://dvc-temp/import-url  57,20s user 2,65s system 27% cpu 3:37,47 total
```

this PR:
```
 $ time dvc import-url s3://dvc-temp/import-url
Importing 's3://dvc-temp/import-url' -> 'import-url'                  
To track the changes with git, run:
        git add import-url.dvc
dvc import-url s3://dvc-temp/import-url  37,69s user 2,57s system 65% cpu 1:01,20 total
```

Problem statement:
tree, import-url: Introduce a new .ls() API for faster dir_info generation (including collection of hash_infos for entries)
Currently, we list files and then fetch every file individually to get their e-tag, which is very costly both in terms of time and the number of API requests we do, so instead, we can retrieve both `size` and `e_tag` directly on the listing (object summaries) and reduce them the amount of requests we do from `n/1000 + n` to only `n/1000`. This shows a %60+ speedup on my local system (low internet speed).

Hints from issue discussion:

Commit message: assert recursive

Patch to review:
diff --git a/dvc/tree/azure.py b/dvc/tree/azure.py
index b80ad74c35..c84d5f61c9 100644
--- a/dvc/tree/azure.py
+++ b/dvc/tree/azure.py
@@ -23,6 +23,8 @@ class AzureTree(BaseTree):
         "knack": "knack",
     }
     PARAM_CHECKSUM = "etag"
+    DETAIL_FIELDS = frozenset(("etag", "size"))
+
     COPY_POLL_SECONDS = 5
     LIST_OBJECT_PAGE_SIZE = 5000
 
@@ -151,6 +153,24 @@ def walk_files(self, path_info, **kwargs):
 
             yield path_info.replace(path=fname)
 
+    def ls(self, path_info, recursive=False, detail=False):
+        assert recursive
+
+        container_client = self.blob_service.get_container_client(
+            path_info.bucket
+        )
+        for blob in container_client.list_blobs(
+            name_starts_with=path_info.path
+        ):
+            file_info = path_info.replace(path=blob.name)
+            if detail:
+                yield file_info, {
+                    "size": blob.size,
+                    "etag": blob.etag,
+                }
+            else:
+                yield file_info
+
     def remove(self, path_info):
         if path_info.scheme != self.scheme:
             raise NotImplementedError
diff --git a/dvc/tree/base.py b/dvc/tree/base.py
index d6cb93bd42..08d43ed479 100644
--- a/dvc/tree/base.py
+++ b/dvc/tree/base.py
@@ -2,13 +2,14 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from functools import partial
 from multiprocessing import cpu_count
-from typing import Any, ClassVar, Dict, Optional
+from typing import Any, ClassVar, Dict, FrozenSet, Optional
 from urllib.parse import urlparse
 
 from funcy import cached_property, decorator
 
 from dvc.dir_info import DirInfo
 from dvc.exceptions import DvcException, DvcIgnoreInCollectedDirError
+from dvc.hash_info import HashInfo
 from dvc.ignore import DvcIgnore
 from dvc.path_info import URLInfo
 from dvc.progress import Tqdm
@@ -63,6 +64,7 @@ class BaseTree:
     CHUNK_SIZE = 64 * 1024 * 1024  # 64 MiB
 
     PARAM_CHECKSUM: ClassVar[Optional[str]] = None
+    DETAIL_FIELDS: FrozenSet[str] = frozenset()
 
     state = StateNoop()
 
@@ -202,6 +204,9 @@ def walk_files(self, path_info, **kwargs):
         """
         raise NotImplementedError
 
+    def ls(self, path_info, recursive=False, detail=False):
+        raise RemoteActionNotImplemented("ls", self.scheme)
+
     def is_empty(self, path_info):
         return False
 
@@ -287,7 +292,6 @@ def hash_to_path_info(self, hash_):
         return self.path_info / hash_[0:2] / hash_[2:]
 
     def _calculate_hashes(self, file_infos):
-        file_infos = list(file_infos)
         with Tqdm(
             total=len(file_infos),
             unit="md5",
@@ -298,24 +302,38 @@ def _calculate_hashes(self, file_infos):
                 hash_infos = executor.map(worker, file_infos)
                 return dict(zip(file_infos, hash_infos))
 
-    def _collect_dir(self, path_info, **kwargs):
-
-        file_infos = set()
-
-        for fname in self.walk_files(path_info, **kwargs):
-            if DvcIgnore.DVCIGNORE_FILE == fname.name:
-                raise DvcIgnoreInCollectedDirError(fname.parent)
+    def _iter_hashes(self, path_info, **kwargs):
+        if self.PARAM_CHECKSUM in self.DETAIL_FIELDS:
+            for file_info, details in self.ls(
+                path_info, recursive=True, detail=True
+            ):
+                hash_info = HashInfo(
+                    self.PARAM_CHECKSUM,
+                    details[self.PARAM_CHECKSUM],
+                    size=details.get("size"),
+                )
+                yield file_info, hash_info
 
-            file_infos.add(fname)
+            return None
 
-        hash_infos = {fi: self.state.get(fi) for fi in file_infos}
-        not_in_state = {fi for fi, hi in hash_infos.items() if hi is None}
+        file_infos = []
+        for file_info in self.walk_files(path_info, **kwargs):
+            hash_info = self.state.get(  # pylint: disable=assignment-from-none
+                file_info
+            )
+            if not hash_info:
+                file_infos.append(file_info)
+                continue
+            yield file_info, hash_info
 
-        new_hash_infos = self._calculate_hashes(not_in_state)
-        hash_infos.update(new_hash_infos)
+        yield from self._calculate_hashes(file_infos).items()
 
+    def _collect_dir(self, path_info, **kwargs):
         dir_info = DirInfo()
-        for fi, hi in hash_infos.items():
+        for fi, hi in self._iter_hashes(path_info, **kwargs):
+            if DvcIgnore.DVCIGNORE_FILE == fi.name:
+                raise DvcIgnoreInCollectedDirError(fi.parent)
+
             # NOTE: this is lossy transformation:
             #   "hey\there" -> "hey/there"
             #   "hey/there" -> "hey/there"
diff --git a/dvc/tree/s3.py b/dvc/tree/s3.py
index eb62786fd4..b0574d2aee 100644
--- a/dvc/tree/s3.py
+++ b/dvc/tree/s3.py
@@ -23,6 +23,7 @@ class S3Tree(BaseTree):
     PATH_CLS = CloudURLInfo
     REQUIRES = {"boto3": "boto3"}
     PARAM_CHECKSUM = "etag"
+    DETAIL_FIELDS = frozenset(("etag", "size"))
 
     def __init__(self, repo, config):
         super().__init__(repo, config)
@@ -227,6 +228,20 @@ def walk_files(self, path_info, **kwargs):
 
             yield path_info.replace(path=fname)
 
+    def ls(self, path_info, recursive=False, detail=False):
+        assert recursive
+
+        with self._get_bucket(path_info.bucket) as bucket:
+            for obj_summary in bucket.objects.filter(Prefix=path_info.path):
+                file_info = path_info.replace(path=obj_summary.key)
+                if detail:
+                    yield file_info, {
+                        "size": obj_summary.size,
+                        "etag": obj_summary.e_tag.strip('"'),
+                    }
+                else:
+                    yield file_info
+
     def remove(self, path_info):
         if path_info.scheme != "s3":
             raise NotImplementedError
diff --git a/tests/func/test_tree.py b/tests/func/test_tree.py
index 3df94ee9ef..c88de70114 100644
--- a/tests/func/test_tree.py
+++ b/tests/func/test_tree.py
@@ -325,3 +325,31 @@ def test_tree_upload_fobj(dvc, tmp_dir, cloud):
     assert tree.exists(to_info)
     with tree.open(to_info, "rb") as stream:
         assert stream.read() == b"foo"
+
+
+@pytest.mark.parametrize(
+    "cloud", [pytest.lazy_fixture("s3"), pytest.lazy_fixture("azure")],
+)
+def test_tree_ls(dvc, cloud):
+    cloud.gen({"data": {"foo": "foo", "bar": {"baz": "baz"}, "quux": "quux"}})
+    tree = get_cloud_tree(dvc, **cloud.config)
+    path_info = tree.path_info
+
+    assert {
+        file_info.name
+        for file_info in tree.ls(path_info / "data", recursive=True)
+    } == {"foo", "baz", "quux"}
+
+
+@pytest.mark.parametrize(
+    "cloud", [pytest.lazy_fixture("s3"), pytest.lazy_fixture("azure")],
+)
+def test_tree_ls_with_etag(dvc, cloud):
+    cloud.gen({"data": {"foo": "foo", "bar": {"baz": "baz"}, "quux": "quux"}})
+    tree = get_cloud_tree(dvc, **cloud.config)
+    path_info = tree.path_info
+
+    for file_info, details in tree.ls(
+        path_info / "data", recursive=True, detail=True
+    ):
+        assert tree.get_file_hash(file_info).value == details["etag"]


Code review comments:
diff_hunk: @@ -227,6 +228,20 @@ def walk_files(self, path_info, **kwargs):
 
             yield path_info.replace(path=fname)
 
+    def ls(self, path_info, recursive=False, detail=False):
+        assert recursive
+
+        with self._get_bucket(path_info.bucket) as bucket:
+            for obj_summary in bucket.objects.filter(Prefix=path_info.path):
+                file_info = path_info.replace(path=obj_summary.key)
+                if detail:
+                    yield file_info, {
+                        "size": obj_summary.size,
+                        "etag": obj_summary.e_tag.strip('"'),
+                    }
path: dvc/tree/s3.py
text: @user1:
This should return just a dict with a string as a name. Creating lots of `file_info`s will be really costly. Also doesn't comply with fsspec. I might be missing a good reason though.

@author:
> Also doesn't comply with fsspec

I am aware of this, though I thought it would comply with the general sense of our trees (exchanging file_info's everywhere [walk_files, upload, download, etc.]). 

@user1:
I see, you are right. But let's try to comply with fsspec if we can. This looks like a good place to do it, as you'll need to call `.replace` only once instead of in every `ls` implementation.