diff --git a/dvc/fs/base.py b/dvc/fs/base.py index e80c1e436..d04617836 100644 --- a/dvc/fs/base.py +++ b/dvc/fs/base.py @@ -235,14 +235,23 @@ class BaseFileSystem: no_progress_bar=no_progress_bar, ) - def upload_fobj(self, fobj, to_info, no_progress_bar=False, **pbar_args): + def upload_fobj( + self, fobj, to_info, no_progress_bar=False, size=None, **pbar_args + ): if not hasattr(self, "_upload_fobj"): raise RemoteActionNotImplemented("upload_fobj", self.scheme) with Tqdm.wrapattr( - fobj, "read", disable=no_progress_bar, bytes=True, **pbar_args + fobj, + "read", + disable=no_progress_bar, + bytes=True, + total=size, + **pbar_args, ) as wrapped: - self._upload_fobj(wrapped, to_info) # pylint: disable=no-member + self._upload_fobj( # pylint: disable=no-member + wrapped, to_info, size=size + ) def download( self, diff --git a/dvc/fs/fsspec_wrapper.py b/dvc/fs/fsspec_wrapper.py index 8b303f9f5..2c49ce859 100644 --- a/dvc/fs/fsspec_wrapper.py +++ b/dvc/fs/fsspec_wrapper.py @@ -118,20 +118,20 @@ class FSSpecWrapper(BaseFileSystem): info["name"] = self._strip_bucket(info["name"]) return info - def _upload_fobj(self, fobj, to_info): + def _upload_fobj(self, fobj, to_info, size=None): with self.open(to_info, "wb") as fdest: shutil.copyfileobj(fobj, fdest, length=fdest.blocksize) def _upload( self, from_file, to_info, name=None, no_progress_bar=False, **kwargs ): - total = os.path.getsize(from_file) + size = os.path.getsize(from_file) with open(from_file, "rb") as fobj: self.upload_fobj( fobj, self._with_bucket(to_info), + size=size, desc=name, - total=total, no_progress_bar=no_progress_bar, ) self.fs.invalidate_cache(self._with_bucket(to_info.parent)) diff --git a/dvc/fs/gdrive.py b/dvc/fs/gdrive.py index 379d7ef6f..11ca9cf22 100644 --- a/dvc/fs/gdrive.py +++ b/dvc/fs/gdrive.py @@ -618,22 +618,27 @@ class GDriveFileSystem(BaseFileSystem): # pylint:disable=abstract-method gdrive_file.FetchMetadata(fields="fileSize") return {"size": gdrive_file.get("fileSize")} - def _upload_fobj(self, fobj, to_info): + def _upload_fobj(self, fobj, to_info, **kwargs): dirname = to_info.parent assert dirname parent_id = self._get_item_id(dirname, create=True) self._gdrive_upload_fobj(fobj, parent_id, to_info.name) def _upload( - self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs + self, + from_file, + to_info, + name=None, + no_progress_bar=False, + **_kwargs, ): with open(from_file, "rb") as fobj: self.upload_fobj( fobj, to_info, + size=os.path.getsize(from_file), no_progress_bar=no_progress_bar, desc=name or to_info.name, - total=os.path.getsize(from_file), ) def _download(self, from_info, to_file, name=None, no_progress_bar=False): diff --git a/dvc/fs/hdfs.py b/dvc/fs/hdfs.py index 91e4fc203..e409b0d56 100644 --- a/dvc/fs/hdfs.py +++ b/dvc/fs/hdfs.py @@ -230,7 +230,7 @@ class HDFSFileSystem(BaseFileSystem): size=self.getsize(path_info), ) - def _upload_fobj(self, fobj, to_info): + def _upload_fobj(self, fobj, to_info, **kwargs): with self.hdfs(to_info) as hdfs: with hdfs.open_output_stream(to_info.path) as fdest: shutil.copyfileobj(fobj, fdest) diff --git a/dvc/fs/http.py b/dvc/fs/http.py index 022890db9..a05c4112f 100644 --- a/dvc/fs/http.py +++ b/dvc/fs/http.py @@ -147,7 +147,7 @@ class HTTPFileSystem(BaseFileSystem): # pylint:disable=abstract-method size = self._content_length(resp) return {"etag": etag, "size": size} - def _upload_fobj(self, fobj, to_info): + def _upload_fobj(self, fobj, to_info, **kwargs): def chunks(fobj): while True: chunk = fobj.read(self.CHUNK_SIZE) @@ -184,9 +184,9 @@ class HTTPFileSystem(BaseFileSystem): # pylint:disable=abstract-method self.upload_fobj( fobj, to_info, + size=None if no_progress_bar else os.path.getsize(from_file), no_progress_bar=no_progress_bar, desc=name or to_info.url, - total=None if no_progress_bar else os.path.getsize(from_file), ) @staticmethod diff --git a/dvc/fs/local.py b/dvc/fs/local.py index 9cb406c40..d614646e1 100644 --- a/dvc/fs/local.py +++ b/dvc/fs/local.py @@ -103,7 +103,7 @@ class LocalFileSystem(BaseFileSystem): self.remove(tmp_info) raise - def _upload_fobj(self, fobj, to_info): + def _upload_fobj(self, fobj, to_info, **kwargs): self.makedirs(to_info.parent) tmp_info = to_info.parent / tmp_fname("") try: diff --git a/dvc/fs/oss.py b/dvc/fs/oss.py index 19ed1795f..bf97b9b95 100644 --- a/dvc/fs/oss.py +++ b/dvc/fs/oss.py @@ -104,7 +104,7 @@ class OSSFileSystem(BaseFileSystem): # pylint:disable=abstract-method logger.debug(f"Removing oss://{path_info}") self._get_bucket(path_info.bucket).delete_object(path_info.path) - def _upload_fobj(self, fobj, to_info): + def _upload_fobj(self, fobj, to_info, **kwargs): self._get_bucket(to_info.bucket).put_object(to_info.path, fobj) def _upload( diff --git a/dvc/fs/ssh/__init__.py b/dvc/fs/ssh/__init__.py index 93a9e4021..192e3ef3f 100644 --- a/dvc/fs/ssh/__init__.py +++ b/dvc/fs/ssh/__init__.py @@ -237,7 +237,7 @@ class SSHFileSystem(BaseFileSystem): # pylint:disable=abstract-method with self.ssh(path_info) as ssh: return ssh.info(path_info.path) - def _upload_fobj(self, fobj, to_info): + def _upload_fobj(self, fobj, to_info, **kwargs): self.makedirs(to_info.parent) with self.open(to_info, mode="wb") as fdest: shutil.copyfileobj(fobj, fdest) diff --git a/dvc/fs/webdav.py b/dvc/fs/webdav.py index 49b36b5e6..d7e5e11a7 100644 --- a/dvc/fs/webdav.py +++ b/dvc/fs/webdav.py @@ -71,13 +71,15 @@ class WebDAVFileSystem(FSSpecWrapper): # pylint:disable=abstract-method return WebdavFileSystem(**self.fs_args) - def _upload_fobj(self, fobj, to_info): + def _upload_fobj(self, fobj, to_info, size: int = None): rpath = self.translate_path_info(to_info) self.makedirs(os.path.dirname(rpath)) # using upload_fileobj to directly upload fileobj rather than buffering # and using overwrite=True to avoid check for an extra exists call, # as caller should ensure that the file does not exist beforehand. - return self.fs.client.upload_fileobj(fobj, rpath, overwrite=True) + return self.fs.client.upload_fileobj( + fobj, rpath, overwrite=True, size=size + ) def makedirs(self, path_info): path = self.translate_path_info(path_info) diff --git a/dvc/fs/webhdfs.py b/dvc/fs/webhdfs.py index 2d14fd737..0d54216e5 100644 --- a/dvc/fs/webhdfs.py +++ b/dvc/fs/webhdfs.py @@ -141,7 +141,7 @@ class WebHDFSFileSystem(BaseFileSystem): # pylint:disable=abstract-method self.hdfs_client.makedirs(to_info.parent.path) self.hdfs_client.rename(from_info.path, to_info.path) - def _upload_fobj(self, fobj, to_info): + def _upload_fobj(self, fobj, to_info, **kwargs): with self.hdfs_client.write(to_info.path) as fdest: shutil.copyfileobj(fobj, fdest) diff --git a/dvc/objects/stage.py b/dvc/objects/stage.py index e66d426e8..ee22210cf 100644 --- a/dvc/objects/stage.py +++ b/dvc/objects/stage.py @@ -19,7 +19,7 @@ def _upload_file(path_info, fs, odb): with fs.open(path_info, mode="rb", chunk_size=fs.CHUNK_SIZE) as stream: stream = HashedStreamReader(stream) odb.fs.upload_fobj( - stream, tmp_info, desc=path_info.name, total=fs.getsize(path_info) + stream, tmp_info, desc=path_info.name, size=fs.getsize(path_info) ) obj = HashFile(tmp_info, odb.fs, stream.hash_info) diff --git a/setup.py b/setup.py index b0423aa4f..f19374f54 100644 --- a/setup.py +++ b/setup.py @@ -106,7 +106,7 @@ ssh = ["paramiko[invoke]>=2.7.0"] # Remove the env marker if/when pyarrow is available for Python3.9 hdfs = ["pyarrow>=2.0.0"] webhdfs = ["hdfs==2.5.8"] -webdav = ["webdav4>=0.7.0"] +webdav = ["webdav4>=0.8.1"] # gssapi should not be included in all_remotes, because it doesn't have wheels # for linux and mac, so it will fail to compile if user doesn't have all the # requirements, including kerberos itself. Once all the wheels are available,