fs: merge upload_fobj and upload (#6570)

* fs: merge upload_fobj and upload

* Update dvc/fs/base.py

Author: skshetry

dvc/fs/base.py

@@ -217,44 +217,50 @@ def is_dir_hash(cls, hash_):
             return False
         return hash_.endswith(cls.CHECKSUM_DIR_SUFFIX)
 
-    def upload(self, from_info, to_info, name=None, no_progress_bar=False):
-        if not hasattr(self, "_upload"):
-            raise RemoteActionNotImplemented("upload", self.scheme)
+    def upload(
+        self,
+        from_info,
+        to_info,
+        total=None,
+        desc=None,
+        no_progress_bar=False,
+        **pbar_args,
+    ):
+        is_file_obj = hasattr(from_info, "read")
+        method = "upload_fobj" if is_file_obj else "_upload"
+        if not hasattr(self, method):
+            raise RemoteActionNotImplemented(method, self.scheme)
 
         if to_info.scheme != self.scheme:
             raise NotImplementedError
 
-        if from_info.scheme != "local":
-            raise NotImplementedError
-
         logger.debug("Uploading '%s' to '%s'", from_info, to_info)
+        if is_file_obj:
+            with Tqdm.wrapattr(
+                from_info,
+                "read",
+                disable=no_progress_bar,
+                bytes=True,
+                total=total,
+                desc=desc,
+                **pbar_args,
+            ) as wrapped:
+                # `size` is used to provide hints to the WebdavFileSystem
+                # for legacy servers.
+                # pylint: disable=no-member
+                return self.upload_fobj(wrapped, to_info, size=total)
 
-        name = name or from_info.name
+        if from_info.scheme != "local":
+            raise NotImplementedError
 
-        self._upload(  # noqa, pylint: disable=no-member
+        name = desc or from_info.name
+        return self._upload(  # noqa, pylint: disable=no-member
             from_info.fspath,
             to_info,
             name=name,
             no_progress_bar=no_progress_bar,
-        )
-
-    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,
-            total=size,
             **pbar_args,
-        ) as wrapped:
-            self._upload_fobj(  # pylint: disable=no-member
-                wrapped, to_info, size=size
-            )
+        )
 
     def download(
         self,

dvc/fs/fsspec_wrapper.py

@@ -130,7 +130,7 @@ def makedirs(self, path_info, **kwargs):
             self._with_bucket(path_info), exist_ok=kwargs.pop("exist_ok", True)
         )
 
-    def _upload_fobj(self, fobj, to_info, size=None):
+    def upload_fobj(self, fobj, to_info, **kwargs):
         self.makedirs(to_info.parent)
         with self.open(to_info, "wb") as fdest:
             shutil.copyfileobj(fobj, fdest, length=fdest.blocksize)
@@ -141,13 +141,15 @@ def _upload(
         self.makedirs(to_info.parent)
         size = os.path.getsize(from_file)
         with open(from_file, "rb") as fobj:
-            self.upload_fobj(
+            with Tqdm.wrapattr(
                 fobj,
-                to_info,
-                size=size,
+                "read",
+                disable=no_progress_bar,
+                bytes=True,
+                total=size,
                 desc=name,
-                no_progress_bar=no_progress_bar,
-            )
+            ) as wrapped:
+                self.upload_fobj(wrapped, to_info, size=size)
         self.fs.invalidate_cache(self._with_bucket(to_info.parent))
 
     def _download(

dvc/fs/gdrive.py

@@ -298,7 +298,7 @@ def _with_bucket(self, path):
 
         return super()._with_bucket(path)
 
-    def _upload_fobj(self, fobj, to_info, size: int = None):
+    def upload_fobj(self, fobj, to_info, **kwargs):
         rpath = self._with_bucket(to_info)
         self.makedirs(os.path.dirname(rpath))
-        return self.fs.upload_fobj(fobj, rpath, size=size)
+        return self.fs.upload_fobj(fobj, rpath, **kwargs)

dvc/fs/hdfs.py

@@ -244,7 +244,7 @@ def checksum(self, path_info):
             size=self.getsize(path_info),
         )
 
-    def _upload_fobj(self, fobj, to_info, **kwargs):
+    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, self.BLOCK_SIZE)

dvc/fs/local.py

@@ -108,7 +108,7 @@ def copy(self, from_info, to_info):
             self.remove(tmp_info)
             raise
 
-    def _upload_fobj(self, fobj, to_info, **kwargs):
+    def upload_fobj(self, fobj, to_info, **kwargs):
         self.makedirs(to_info.parent)
         tmp_info = to_info.parent / tmp_fname("")
         try:

dvc/fs/ssh.py

@@ -120,9 +120,9 @@ def fs(self):
     # Ensure that if an interrupt happens during the transfer, we don't
     # pollute the cache.
 
-    def _upload_fobj(self, fobj, to_info, *args, **kwargs):
+    def upload_fobj(self, fobj, to_info, **kwargs):
         with as_atomic(self, to_info) as tmp_file:
-            super()._upload_fobj(fobj, tmp_file, *args, **kwargs)
+            super().upload_fobj(fobj, tmp_file, **kwargs)
 
     def _upload(self, from_file, to_info, *args, **kwargs):
         with as_atomic(self, to_info) as tmp_file:

dvc/fs/utils.py

@@ -33,7 +33,8 @@ def transfer(
             return from_fs.download_file(from_info, to_info)
 
         with from_fs.open(from_info, mode="rb") as fobj:
-            return to_fs.upload_fobj(fobj, to_info)
+            size = from_fs.getsize(from_info)
+            return to_fs.upload(fobj, to_info, total=size)
     except OSError as exc:
         # If the target file already exists, we are going to simply
         # ignore the exception (#4992).

dvc/fs/webdav.py

@@ -1,5 +1,4 @@
 import logging
-import os
 import threading
 from functools import lru_cache
 
@@ -71,15 +70,13 @@ def fs(self):
 
         return WebdavFileSystem(**self.fs_args)
 
-    def _upload_fobj(self, fobj, to_info, size: int = None):
+    def upload_fobj(self, fobj, to_info, **kwargs):
         rpath = self.translate_path_info(to_info)
-        self.makedirs(os.path.dirname(rpath))
+        size = kwargs.get("size")
         # 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, size=size
-        )
+        return self.fs.upload_fileobj(fobj, rpath, overwrite=True, size=size)
 
     @lru_cache(512)
     def translate_path_info(self, path):

dvc/fs/webhdfs.py

@@ -141,7 +141,7 @@ def move(self, from_info, to_info):
         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, **kwargs):
+    def upload_fobj(self, fobj, to_info, **kwargs):
         with self.hdfs_client.write(to_info.path) as fdest:
             shutil.copyfileobj(fobj, fdest)
 

dvc/objects/db/reference.py

@@ -70,7 +70,7 @@ def _add_file(
         ref_fobj = io.BytesIO(ref_file.to_bytes())
         ref_fobj.seek(0)
         try:
-            self.fs.upload_fobj(ref_fobj, to_info)
+            self.fs.upload(ref_fobj, to_info)
         except OSError as exc:
             if isinstance(exc, FileExistsError) or (
                 os.name == "nt"

dvc/objects/stage.py

@@ -34,9 +34,8 @@ def _upload_file(path_info, fs, odb, upload_odb):
     tmp_info = upload_odb.path_info / tmp_fname()
     with fs.open(path_info, mode="rb", chunk_size=fs.CHUNK_SIZE) as stream:
         stream = HashedStreamReader(stream)
-        upload_odb.fs.upload_fobj(
-            stream, tmp_info, desc=path_info.name, size=fs.getsize(path_info)
-        )
+        size = fs.getsize(path_info)
+        upload_odb.fs.upload(stream, tmp_info, desc=path_info.name, total=size)
 
     odb.add(tmp_info, upload_odb.fs, stream.hash_info)
     return path_info, odb.get(stream.hash_info)

tests/func/test_fs.py

@@ -405,7 +405,7 @@ def test_fs_makedirs_on_upload_and_copy(dvc, cloud):
     fs = cls(**config)
 
     with io.BytesIO(b"foo") as stream:
-        fs.upload_fobj(stream, cloud / "dir" / "foo")
+        fs.upload(stream, cloud / "dir" / "foo")
 
     assert fs.isdir(cloud / "dir")
     assert fs.exists(cloud / "dir" / "foo")