From 87b49fd9f9829caab351a81560d2b90d791feae2 Mon Sep 17 00:00:00 2001 From: Batuhan Taskaya Date: Wed, 7 Jul 2021 22:53:43 +0300 Subject: [PATCH] staging: optimize the number of info() calls (#6281) * staging: optimize the number of info() calls * fs: make info() give details about the object --- dvc/fs/dvc.py | 2 +- dvc/fs/gdrive.py | 11 +++++++---- dvc/fs/git.py | 6 +++++- dvc/fs/hdfs.py | 16 ++++++++++++++-- dvc/fs/http.py | 2 +- dvc/fs/local.py | 2 +- dvc/fs/ssh/connection.py | 2 +- dvc/fs/webhdfs.py | 2 +- dvc/objects/stage.py | 12 +++--------- tests/unit/remote/test_http.py | 6 +++++- 10 files changed, 39 insertions(+), 22 deletions(-) diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py index 64a864f33..5755ab7ca 100644 --- a/dvc/fs/dvc.py +++ b/dvc/fs/dvc.py @@ -233,7 +233,7 @@ class DvcFileSystem(BaseFileSystem): # pylint:disable=abstract-method def info(self, path_info): meta = self.metadata(path_info) - ret = {"type": "dir" if meta.isdir else "file"} + ret = {"type": "directory" if meta.isdir else "file"} if meta.is_output and len(meta.outs) == 1 and meta.outs[0].hash_info: hash_info = meta.outs[0].hash_info ret["size"] = hash_info.size diff --git a/dvc/fs/gdrive.py b/dvc/fs/gdrive.py index 11ca9cf22..4bbf583f7 100644 --- a/dvc/fs/gdrive.py +++ b/dvc/fs/gdrive.py @@ -1,3 +1,4 @@ +import errno import io import logging import os @@ -11,7 +12,7 @@ from urllib.parse import urlparse from funcy import cached_property, retry, wrap_prop, wrap_with from funcy.py3 import cat -from dvc.exceptions import DvcException, FileMissingError +from dvc.exceptions import DvcException from dvc.path_info import CloudURLInfo from dvc.progress import Tqdm from dvc.scheme import Schemes @@ -521,12 +522,14 @@ class GDriveFileSystem(BaseFileSystem): # pylint:disable=abstract-method return min(item_ids) assert not create - raise FileMissingError(path_info, hint) + raise FileNotFoundError( + errno.ENOENT, os.strerror(errno.ENOENT), hint or path_info + ) def exists(self, path_info) -> bool: try: self._get_item_id(path_info) - except FileMissingError: + except FileNotFoundError: return False else: return True @@ -616,7 +619,7 @@ class GDriveFileSystem(BaseFileSystem): # pylint:disable=abstract-method item_id = self._get_item_id(path_info) gdrive_file = self._drive.CreateFile({"id": item_id}) gdrive_file.FetchMetadata(fields="fileSize") - return {"size": gdrive_file.get("fileSize")} + return {"size": gdrive_file.get("fileSize"), "type": "file"} def _upload_fobj(self, fobj, to_info, **kwargs): dirname = to_info.parent diff --git a/dvc/fs/git.py b/dvc/fs/git.py index c02028136..05acdbb90 100644 --- a/dvc/fs/git.py +++ b/dvc/fs/git.py @@ -1,5 +1,6 @@ import errno import os +import stat from dvc.utils import is_exec, relpath @@ -106,7 +107,10 @@ class GitFileSystem(BaseFileSystem): # pylint:disable=abstract-method raise FileNotFoundError( errno.ENOENT, os.strerror(errno.ENOENT), path_info ) - return {"size": st.st_size} + return { + "size": st.st_size, + "type": "directory" if stat.S_ISDIR(st.st_mode) else "file", + } def stat(self, path): key = self._get_key(path) diff --git a/dvc/fs/hdfs.py b/dvc/fs/hdfs.py index e409b0d56..b2893d255 100644 --- a/dvc/fs/hdfs.py +++ b/dvc/fs/hdfs.py @@ -1,3 +1,4 @@ +import errno import io import logging import os @@ -220,8 +221,19 @@ class HDFSFileSystem(BaseFileSystem): def info(self, path_info): with self.hdfs(path_info) as hdfs: - finfo = hdfs.get_file_info(path_info.path) - return {"size": finfo.size} + from pyarrow.fs import FileType + + file_info = hdfs.get_file_info(path_info.path) + if file_info.type is FileType.Directory: + kind = "directory" + elif file_info.type is FileType.File: + kind = "file" + else: + raise FileNotFoundError( + errno.ENOENT, os.strerror(errno.ENOENT), path_info + ) + + return {"size": file_info.size, "type": kind} def checksum(self, path_info): return HashInfo( diff --git a/dvc/fs/http.py b/dvc/fs/http.py index c265db206..1805bbf16 100644 --- a/dvc/fs/http.py +++ b/dvc/fs/http.py @@ -146,7 +146,7 @@ class HTTPFileSystem(BaseFileSystem): # pylint:disable=abstract-method resp = self._head(path_info.url) etag = resp.headers.get("ETag") or resp.headers.get("Content-MD5") size = self._content_length(resp) - return {"etag": etag, "size": size} + return {"etag": etag, "size": size, "type": "file"} def _upload_fobj(self, fobj, to_info, **kwargs): def chunks(fobj): diff --git a/dvc/fs/local.py b/dvc/fs/local.py index ed0bc958f..996c53ff9 100644 --- a/dvc/fs/local.py +++ b/dvc/fs/local.py @@ -157,7 +157,7 @@ class LocalFileSystem(BaseFileSystem): st = os.stat(path_info) return { "size": st.st_size, - "type": "dir" if stat.S_ISDIR(st.st_mode) else "file", + "type": "directory" if stat.S_ISDIR(st.st_mode) else "file", } def _upload( diff --git a/dvc/fs/ssh/connection.py b/dvc/fs/ssh/connection.py index 9baad43b1..465172b27 100644 --- a/dvc/fs/ssh/connection.py +++ b/dvc/fs/ssh/connection.py @@ -82,7 +82,7 @@ class SSHConnection: st = self.sftp.stat(path) return { "size": st.st_size, - "type": "dir" if stat.S_ISDIR(st.st_mode) else "file", + "type": "directory" if stat.S_ISDIR(st.st_mode) else "file", } def getsize(self, path): diff --git a/dvc/fs/webhdfs.py b/dvc/fs/webhdfs.py index 0d54216e5..f2cda6e8f 100644 --- a/dvc/fs/webhdfs.py +++ b/dvc/fs/webhdfs.py @@ -123,7 +123,7 @@ class WebHDFSFileSystem(BaseFileSystem): # pylint:disable=abstract-method def info(self, path_info): st = self.hdfs_client.status(path_info.path) - return {"size": st["length"]} + return {"size": st["length"], "type": "file"} def checksum(self, path_info): return HashInfo( diff --git a/dvc/objects/stage.py b/dvc/objects/stage.py index d5920576b..63bcafa05 100644 --- a/dvc/objects/stage.py +++ b/dvc/objects/stage.py @@ -1,5 +1,3 @@ -import errno -import os from concurrent.futures import ThreadPoolExecutor from functools import partial @@ -174,11 +172,7 @@ def stage(odb, path_info, fs, name, upload=False, **kwargs): isinstance(path_info, str) or path_info.scheme == fs.scheme ) - if not fs.exists(path_info): - raise FileNotFoundError( - errno.ENOENT, os.strerror(errno.ENOENT), path_info - ) - + details = fs.info(path_info) state = odb.state # pylint: disable=assignment-from-none hash_info = state.get(path_info, fs) @@ -210,12 +204,12 @@ def stage(odb, path_info, fs, name, upload=False, **kwargs): obj.hash_info.size = hash_info.size return obj - if fs.isdir(path_info): + if details["type"] == "directory": obj = _get_tree_obj(path_info, fs, name, odb, state, upload, **kwargs) else: _, obj = _get_file_obj(path_info, fs, name, odb, state, upload) - if obj.hash_info and fs.exists(path_info): + if obj.hash_info: state.save(path_info, fs, obj.hash_info) return obj diff --git a/tests/unit/remote/test_http.py b/tests/unit/remote/test_http.py index 779ba7126..cfefff2a8 100644 --- a/tests/unit/remote/test_http.py +++ b/tests/unit/remote/test_http.py @@ -175,5 +175,9 @@ def test_content_length(mocker, headers, expected_size): url = URLInfo("https://example.org/file.txt") - assert fs.info(url) == {"etag": None, "size": expected_size} + assert fs.info(url) == { + "etag": None, + "size": expected_size, + "type": "file", + } assert fs._content_length(res) == expected_size