From c840cecc91c7012e3bebd5b810f10ac21a37ebd4 Mon Sep 17 00:00:00 2001 From: Batuhan Taskaya Date: Wed, 13 Jan 2021 19:35:57 +0300 Subject: [PATCH] tree: implement getsize() (#5261) * Implement Tree.getsize() API * Remove azure from tests * strip support for recursive size calculation * _getsize() -> getsize() --- dvc/tree/azure.py | 7 +++++++ dvc/tree/base.py | 3 +++ dvc/tree/gdrive.py | 6 ++++++ dvc/tree/gs.py | 5 +++++ dvc/tree/hdfs.py | 5 +++++ dvc/tree/http.py | 6 ++++++ dvc/tree/s3.py | 4 ++++ dvc/tree/webdav.py | 4 ++-- dvc/tree/webhdfs.py | 5 ++++- tests/func/test_tree.py | 24 ++++++++++++++++++++++++ 10 files changed, 66 insertions(+), 3 deletions(-) diff --git a/dvc/tree/azure.py b/dvc/tree/azure.py index 3a14aac33..57307cecc 100644 --- a/dvc/tree/azure.py +++ b/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)) diff --git a/dvc/tree/base.py b/dvc/tree/base.py index ff66a2ba8..705a4a509 100644 --- a/dvc/tree/base.py +++ b/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) diff --git a/dvc/tree/gdrive.py b/dvc/tree/gdrive.py index e8d2b0b3f..00fbcfbd8 100644 --- a/dvc/tree/gdrive.py +++ b/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 ): diff --git a/dvc/tree/gs.py b/dvc/tree/gs.py index 04aa6fc05..7141d7a33 100644 --- a/dvc/tree/gs.py +++ b/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 diff --git a/dvc/tree/hdfs.py b/dvc/tree/hdfs.py index f564f4b9b..2c2db68c3 100644 --- a/dvc/tree/hdfs.py +++ b/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: diff --git a/dvc/tree/http.py b/dvc/tree/http.py index c003a7233..7ffd78504 100644 --- a/dvc/tree/http.py +++ b/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 diff --git a/dvc/tree/s3.py b/dvc/tree/s3.py index ca4efdd4d..7b3eb0c10 100644 --- a/dvc/tree/s3.py +++ b/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) diff --git a/dvc/tree/webdav.py b/dvc/tree/webdav.py index 1daf3fe67..88d557bbb 100644 --- a/dvc/tree/webdav.py +++ b/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"]) diff --git a/dvc/tree/webhdfs.py b/dvc/tree/webhdfs.py index f6e77ef90..a65568f09 100644 --- a/dvc/tree/webhdfs.py +++ b/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: diff --git a/tests/func/test_tree.py b/tests/func/test_tree.py index 7bb5d015f..2d5db2672 100644 --- a/tests/func/test_tree.py +++ b/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