tree: Introduce the new .ls() API (#5275)

* s3: faster directory hashing

* Introduce Tree.info() API

* tests for s3

* Implement support for azure

* New .ls() API

* remove _iter_hashes from other trees

* assert recursive

* yield strings instead of file_infos

* Add GSTree.ls()

* Add DETAIL_FIELDS to the GSTree

* WebDAVTree.ls()

* proper webdav root directory initalization

* disable arguments-differ

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
 
@@ -149,6 +151,27 @@ class AzureTree(BaseTree):
 
             yield path_info.replace(path=fname)
 
+    def ls(
+        self, path_info, detail=False, recursive=False
+    ):  # pylint: disable=arguments-differ
+        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
+        ):
+            if detail:
+                yield {
+                    "type": "file",
+                    "name": blob.name,
+                    "size": blob.size,
+                    "etag": blob.etag,
+                }
+            else:
+                yield blob.name
+
     def remove(self, path_info):
         if path_info.scheme != self.scheme:
             raise NotImplementedError

dvc/tree/base.py

@@ -2,13 +2,14 @@ import logging
 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 @@ class BaseTree:
         """
         raise NotImplementedError
 
+    def ls(self, path_info, detail=False, **kwargs):
+        raise RemoteActionNotImplemented("ls", self.scheme)
+
     def is_empty(self, path_info):
         return False
 
@@ -287,7 +292,6 @@ class BaseTree:
         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,37 @@ class BaseTree:
                 hash_infos = executor.map(worker, file_infos)
                 return dict(zip(file_infos, hash_infos))
 
+    def _iter_hashes(self, path_info, **kwargs):
+        if self.PARAM_CHECKSUM in self.DETAIL_FIELDS:
+            for details in self.ls(path_info, recursive=True, detail=True):
+                file_info = path_info.replace(path=details["name"])
+                hash_info = HashInfo(
+                    self.PARAM_CHECKSUM,
+                    details[self.PARAM_CHECKSUM],
+                    size=details.get("size"),
+                )
+                yield file_info, hash_info
+
+            return 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
+
+        yield from self._calculate_hashes(file_infos).items()
+
     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)
-
-            file_infos.add(fname)
-
-        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}
-
-        new_hash_infos = self._calculate_hashes(not_in_state)
-        hash_infos.update(new_hash_infos)
-
         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"

dvc/tree/gs.py

@@ -56,6 +56,7 @@ class GSTree(BaseTree):
     PATH_CLS = CloudURLInfo
     REQUIRES = {"google-cloud-storage": "google.cloud.storage"}
     PARAM_CHECKSUM = "md5"
+    DETAIL_FIELDS = frozenset(("md5", "size"))
 
     def __init__(self, repo, config):
         super().__init__(repo, config)
@@ -143,6 +144,24 @@ class GSTree(BaseTree):
                 continue
             yield path_info.replace(fname)
 
+    def ls(
+        self, path_info, detail=False, recursive=False
+    ):  # pylint: disable=arguments-differ
+        assert recursive
+
+        for blob in self.gs.bucket(path_info.bucket).list_blobs(
+            prefix=path_info.path
+        ):
+            if detail:
+                yield {
+                    "type": "file",
+                    "name": blob.name,
+                    "md5": blob.md5_hash.encode().hex(),
+                    "size": blob.size,
+                }
+            else:
+                yield blob.name
+
     def remove(self, path_info):
         if path_info.scheme != "gs":
             raise NotImplementedError

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,23 @@ class S3Tree(BaseTree):
 
             yield path_info.replace(path=fname)
 
+    def ls(
+        self, path_info, detail=False, recursive=False
+    ):  # pylint: disable=arguments-differ
+        assert recursive
+
+        with self._get_bucket(path_info.bucket) as bucket:
+            for obj_summary in bucket.objects.filter(Prefix=path_info.path):
+                if detail:
+                    yield {
+                        "type": "file",
+                        "name": obj_summary.key,
+                        "size": obj_summary.size,
+                        "etag": obj_summary.e_tag.strip('"'),
+                    }
+                else:
+                    yield obj_summary.key
+
     def remove(self, path_info):
         if path_info.scheme != "s3":
             raise NotImplementedError

dvc/tree/webdav.py

@@ -44,6 +44,7 @@ class WebDAVTree(BaseTree):  # pylint:disable=abstract-method
     CHUNK_SIZE = 2 ** 16
 
     PARAM_CHECKSUM = "etag"
+    DETAIL_FIELDS = frozenset(("etag", "size"))
 
     # Constructor
     def __init__(self, repo, config):
@@ -189,6 +190,36 @@ class WebDAVTree(BaseTree):  # pylint:disable=abstract-method
                     # Yield path info to non directory
                     yield info
 
+    def ls(
+        self, path_info, detail=False, recursive=False
+    ):  # pylint: disable=arguments-differ
+        dirs = deque([path_info.path])
+
+        while dirs:
+            for entry in self._client.list(dirs.pop(), get_info=True):
+                path = entry["path"]
+                if entry["isdir"]:
+                    dirs.append(path)
+                    continue
+
+                if detail:
+                    yield {
+                        "type": "file",
+                        "name": path,
+                        "size": entry["size"],
+                        "etag": entry["etag"],
+                    }
+                else:
+                    yield path
+
+            if not recursive:
+                for entry in dirs:
+                    if detail:
+                        yield {"type": "directory", "name": entry}
+                    else:
+                        yield entry
+                return None
+
     # Removes file/directory
     def remove(self, path_info):
         # Use webdav client clean (DELETE) method to remove file/directory

tests/func/test_tree.py

@@ -279,3 +279,44 @@ 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"),
+        pytest.lazy_fixture("gs"),
+        pytest.lazy_fixture("webdav"),
+    ],
+)
+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 {
+        os.path.basename(file_key)
+        for file_key in tree.ls(path_info / "data", recursive=True)
+    } == {"foo", "baz", "quux"}
+
+
+@pytest.mark.parametrize(
+    "cloud",
+    [
+        pytest.lazy_fixture("s3"),
+        pytest.lazy_fixture("azure"),
+        pytest.lazy_fixture("gs"),
+        pytest.lazy_fixture("webdav"),
+    ],
+)
+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 details in tree.ls(path_info / "data", recursive=True, detail=True):
+        assert (
+            tree.get_file_hash(path_info.replace(path=details["name"])).value
+            == details[tree.PARAM_CHECKSUM]
+        )

tests/remotes/azure.py

@@ -49,7 +49,6 @@ class Azure(Base, CloudURLInfo):
     def mkdir(self, mode=0o777, parents=False, exist_ok=False):
         assert mode == 0o777
         assert parents
-        assert not exist_ok
 
     def write_bytes(self, contents):
         self.blob_client.upload_blob(contents, overwrite=True)

tests/remotes/base.py

@@ -36,11 +36,11 @@ class Base(URLInfo):
 
             if isinstance(contents, dict):
                 if not contents:
-                    path.mkdir(parents=True)
+                    path.mkdir(parents=True, exist_ok=True)
                 else:
                     self._gen(contents, prefix=path)
             else:
-                path.parent.mkdir(parents=True)
+                path.parent.mkdir(parents=True, exist_ok=True)
                 if isinstance(contents, bytes):
                     path.write_bytes(contents)
                 else:

tests/remotes/gs.py

@@ -79,7 +79,6 @@ class GCP(Base, CloudURLInfo):
     def mkdir(self, mode=0o777, parents=False, exist_ok=False):
         assert mode == 0o777
         assert parents
-        assert not exist_ok
 
     def write_bytes(self, contents):
         assert isinstance(contents, bytes)

tests/remotes/hdfs.py

@@ -48,7 +48,6 @@ class HDFS(Base, URLInfo):  # pylint: disable=abstract-method
     def mkdir(self, mode=0o777, parents=False, exist_ok=False):
         assert mode == 0o777
         assert parents
-        assert not exist_ok
 
         with self._hdfs() as _hdfs:
             # NOTE: fs.create_dir creates parents by default

tests/remotes/http.py

@@ -20,7 +20,6 @@ class HTTP(Base, HTTPURLInfo):
     def mkdir(self, mode=0o777, parents=False, exist_ok=False):
         assert mode == 0o777
         assert parents
-        assert not exist_ok
 
     def write_bytes(self, contents):
         assert isinstance(contents, bytes)

tests/remotes/s3.py

@@ -74,7 +74,6 @@ class S3(Base, CloudURLInfo):
     def mkdir(self, mode=0o777, parents=False, exist_ok=False):
         assert mode == 0o777
         assert parents
-        assert not exist_ok
 
     def write_bytes(self, contents):
         self._s3.put_object(

tests/remotes/ssh.py

@@ -81,7 +81,6 @@ class SSHMocked(Base, URLInfo):
     def mkdir(self, mode=0o777, parents=False, exist_ok=False):
         assert mode == 0o777
         assert parents
-        assert not exist_ok
 
         with self._ssh() as _ssh:
             _ssh.makedirs(self.path)

tests/remotes/webdav.py

@@ -1,6 +1,8 @@
 # pylint:disable=abstract-method
 import os
+import uuid
 from pathlib import Path
+from tempfile import TemporaryDirectory
 from wsgiref.simple_server import make_server
 
 import pytest
@@ -8,28 +10,32 @@ from funcy import first
 from wsgidav.wsgidav_app import WsgiDAVApp
 
 from dvc.path_info import WebDAVURLInfo
-from dvc.utils.fs import makedirs
 from tests.utils.httpd import run_server_on_thread
 
 from .base import Base
 
 AUTH = {"user1": {"password": "password1"}}
+_WEBDAV_ROOT = TemporaryDirectory()
+_WEBDAV_DIR = Path(_WEBDAV_ROOT.name)
 
 
 class Webdav(Base, WebDAVURLInfo):
     @staticmethod
-    def get_url(port):  # pylint: disable=arguments-differ
-        return f"webdav://localhost:{port}"
+    def get_url(port, root_id):  # pylint: disable=arguments-differ
+        return f"webdav://localhost:{port}/{root_id}/"
 
     def mkdir(self, mode=0o777, parents=False, exist_ok=False):
-        assert parents
-        makedirs(self.path, exist_ok=exist_ok)
+        self.dir_path.mkdir(parents=parents, exist_ok=True)
 
     def write_bytes(self, contents):
-        Path(self.path).write_bytes(contents)
+        self.dir_path.write_bytes(contents)
 
     def write_text(self, contents, encoding=None, errors=None):
-        Path(self.path).write_text(contents, encoding=encoding, errors=errors)
+        self.dir_path.write_text(contents, encoding=encoding, errors=errors)
+
+    @property
+    def dir_path(self):
+        return _WEBDAV_DIR / self.path[1:]
 
 
 @pytest.fixture
@@ -37,8 +43,10 @@ def webdav_server(test_config, tmp_path_factory):
     test_config.requires("webdav")
 
     host, port = "localhost", 0
-    directory = os.fspath(tmp_path_factory.mktemp("http"))
-    dirmap = {"/": directory}
+    root_id = str(uuid.uuid4())
+    root_dir = _WEBDAV_DIR / root_id
+    root_dir.mkdir()
+    dirmap = {f"/{root_id}": os.fspath(_WEBDAV_DIR / root_id)}
 
     app = WsgiDAVApp(
         {
@@ -49,13 +57,14 @@ def webdav_server(test_config, tmp_path_factory):
         }
     )
     server = make_server(host, port, app)
+    server.root_id = root_id
     with run_server_on_thread(server) as httpd:
         yield httpd
 
 
 @pytest.fixture
 def webdav(webdav_server):
-    url = Webdav.get_url(webdav_server.server_port)
+    url = Webdav.get_url(webdav_server.server_port, webdav_server.root_id)
     ret = Webdav(url)
     user, secrets = first(AUTH.items())
     ret.config = {"url": url, "user": user, **secrets}

tests/remotes/webhdfs.py

@@ -35,7 +35,6 @@ class WebHDFS(Base, URLInfo):  # pylint: disable=abstract-method
     def mkdir(self, mode=0o777, parents=False, exist_ok=False):
         assert mode == 0o777
         assert parents
-        assert not exist_ok
 
         with self._webhdfs() as _hdfs:
             # NOTE: hdfs.makekdirs always creates parents