tree: implement getsize() (#5261)

* Implement Tree.getsize() API

* Remove azure from tests

* strip support for recursive size calculation

* _getsize() -> getsize()

dvc/tree/azure.py

@@ -160,6 +160,13 @@ class AzureTree(BaseTree):
             path_info.bucket, path_info.path
         ).delete_blob()
 
+    def getsize(self, path_info):
+        blob_client = self.blob_service.get_blob_client(
+            path_info.bucket, path_info.path
+        )
+        properties = blob_client.get_blob_properties()
+        return properties.size
+
     def get_file_hash(self, path_info):
         return HashInfo(self.PARAM_CHECKSUM, self.get_etag(path_info))
 

dvc/tree/base.py

@@ -220,6 +220,9 @@ class BaseTree:
     def is_empty(self, path_info):
         return False
 
+    def getsize(self, path_info):
+        return None
+
     def remove(self, path_info):
         raise RemoteActionNotImplemented("remove", self.scheme)
 

dvc/tree/gdrive.py

@@ -564,6 +564,12 @@ class GDriveTree(BaseTree):
     def get_file_hash(self, path_info):
         raise NotImplementedError
 
+    def getsize(self, path_info):
+        item_id = self._get_item_id(path_info)
+        gdrive_file = self._drive.CreateFile({"id": item_id})
+        gdrive_file.FetchMetadata(fields="fileSize")
+        return gdrive_file.get("fileSize")
+
     def _upload(
         self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
     ):

dvc/tree/gs.py

@@ -138,6 +138,11 @@ class GSTree(BaseTree):
         blob = self.gs.bucket(path_info.bucket).blob(path_info.path)
         return blob.exists()
 
+    def getsize(self, path_info):
+        bucket = self.gs.bucket(path_info.bucket)
+        blob = bucket.get_blob(path_info.path)
+        return blob.size
+
     def _list_paths(self, path_info, max_items=None):
         for blob in self.gs.bucket(path_info.bucket).list_blobs(
             prefix=path_info.path, max_results=max_items

dvc/tree/hdfs.py

@@ -182,6 +182,11 @@ class HDFSTree(BaseTree):
             file_info = hdfs.get_file_info(path_info.path)
             return file_info.type == pyarrow.fs.FileType.Directory
 
+    def getsize(self, path_info):
+        with self.hdfs(path_info) as hdfs:
+            file_info = hdfs.get_file_info(path_info.path)
+            return file_info.size
+
     def hadoop_fs(self, cmd, user=None):
         cmd = "hadoop fs -" + cmd
         if user:

dvc/tree/http.py

@@ -148,6 +148,12 @@ class HTTPTree(BaseTree):  # pylint:disable=abstract-method
             return True
         raise HTTPError(res.status_code, res.reason)
 
+    def getsize(self, path_info):
+        response = self.request("GET", path_info.url, stream=True)
+        if response.status_code != 200:
+            raise HTTPError(response.status_code, response.reason)
+        return self._content_length(response)
+
     def get_file_hash(self, path_info):
         url = path_info.url
 

dvc/tree/s3.py

@@ -201,6 +201,10 @@ class S3Tree(BaseTree):
 
         return path_info.path in self._list_paths(path_info)
 
+    def getsize(self, path_info):
+        with self._get_obj(path_info) as obj:
+            return obj.content_length
+
     def _list_paths(self, path_info, max_items=None):
         with self._get_bucket(path_info.bucket) as bucket:
             obj_summaries = bucket.objects.filter(Prefix=path_info.path)

dvc/tree/webdav.py

@@ -199,7 +199,7 @@ class WebDAVTree(BaseTree):  # pylint:disable=abstract-method
             with Tqdm.wrapattr(
                 fd,
                 "write",
-                total=None if no_progress_bar else self._file_size(from_info),
+                total=None if no_progress_bar else self.getsize(from_info),
                 leave=False,
                 desc=from_info.url if name is None else name,
                 disable=no_progress_bar,
@@ -236,6 +236,6 @@ class WebDAVTree(BaseTree):  # pylint:disable=abstract-method
                 )
 
     # Queries size of file at remote
-    def _file_size(self, path_info):
+    def getsize(self, path_info):
         # Get file size from info dictionary and convert to int (from str)
         return int(self._client.info(path_info.path)["size"])

dvc/tree/webhdfs.py

@@ -122,6 +122,9 @@ class WebHDFSTree(BaseTree):
         status = self.hdfs_client.status(path_info.path, strict=False)
         return status is not None
 
+    def getsize(self, path_info):
+        return self.hdfs_client.status(path_info.path)["length"]
+
     def get_file_hash(self, path_info):
         checksum = self.hdfs_client.checksum(path_info.path)
         hash_info = HashInfo(self.PARAM_CHECKSUM, checksum["bytes"])
@@ -155,7 +158,7 @@ class WebHDFSTree(BaseTree):
     def _download(
         self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs
     ):
-        total = self.hdfs_client.status(from_info.path)["length"]
+        total = self.getsize(from_info)
         with Tqdm(
             desc=name, total=total, disable=no_progress_bar, bytes=True
         ) as pbar:

tests/func/test_tree.py

@@ -2,9 +2,12 @@ import os
 from operator import itemgetter
 from os.path import join
 
+import pytest
+
 from dvc.path_info import PathInfo
 from dvc.repo import Repo
 from dvc.scm import SCM
+from dvc.tree import get_cloud_tree
 from dvc.tree.local import LocalTree
 from dvc.tree.repo import RepoTree
 from dvc.utils.fs import remove
@@ -273,3 +276,24 @@ def test_walk_dont_ignore_subrepos(tmp_dir, scm, dvc):
     kw = {"ignore_subrepos": False}
     assert get_dirs(next(dvc_tree.walk(path, **kw))) == ["subdir"]
     assert get_dirs(next(scm_tree.walk(path, **kw))) == ["subdir"]
+
+
+@pytest.mark.parametrize(
+    "workspace",
+    [
+        pytest.lazy_fixture("local_cloud"),
+        pytest.lazy_fixture("s3"),
+        pytest.lazy_fixture("gs"),
+        pytest.lazy_fixture("gdrive"),
+        pytest.lazy_fixture("hdfs"),
+        pytest.lazy_fixture("http"),
+    ],
+    indirect=True,
+)
+def test_tree_getsize(dvc, workspace):
+    workspace.gen({"data": {"foo": "foo"}, "baz": "baz baz"})
+    tree = get_cloud_tree(dvc, url=workspace.url)
+    path_info = tree.path_info
+
+    assert tree.getsize(path_info / "baz") == 7
+    assert tree.getsize(path_info / "data" / "foo") == 3