fs.put_file: add support for uploading fileobjects

Also changed the callbacks to inherit from our own so that we can
add few more utilities to it.

dvc/fs/_callback.py

@@ -1,9 +1,51 @@
+from functools import wraps
+from typing import IO, TYPE_CHECKING, Optional, TypeVar, cast
+
 import fsspec
 
 from dvc.progress import Tqdm
 
+if TYPE_CHECKING:
+    from typing import Callable
+
+    from typing_extensions import ParamSpec
+
+    _P = ParamSpec("_P")
+    _R = TypeVar("_R")
+
 
 class FsspecCallback(fsspec.Callback):
+    """FsspecCallback usable as a context manager, and a few helper methods."""
+
+    def wrap_attr(self, fobj: IO, method: str = "read") -> IO:
+        from tqdm.utils import CallbackIOWrapper
+
+        wrapped = CallbackIOWrapper(self.relative_update, fobj, method)
+        return cast(IO, wrapped)
+
+    def wrap_fn(self, fn: "Callable[_P, _R]") -> "Callable[_P, _R]":
+        @wraps(fn)
+        def wrapped(*args: "_P.args", **kwargs: "_P.kwargs") -> "_R":
+            res = fn(*args, **kwargs)
+            self.relative_update()
+            return res
+
+        return wrapped
+
+    @classmethod
+    def as_callback(
+        cls, maybe_callback: Optional["FsspecCallback"] = None
+    ) -> "FsspecCallback":
+        if maybe_callback is None:
+            return DEFAULT_CALLBACK
+        return maybe_callback
+
+
+class NoOpCallback(FsspecCallback, fsspec.callbacks.NoOpCallback):
+    pass
+
+
+class TqdmCallback(FsspecCallback):
     def __init__(self, progress_bar):
         self.progress_bar = progress_bar
         super().__init__()
@@ -22,15 +64,6 @@ class FsspecCallback(fsspec.Callback):
         self.progress_bar.update_to(value)
         super().absolute_update(value)
 
-    @staticmethod
-    def wrap_fn(cb, fn):
-        def wrapped(*args, **kwargs):
-            res = fn(*args, **kwargs)
-            cb.relative_update()
-            return res
-
-        return wrapped
-
 
 def tdqm_or_callback_wrapped(
     fobj, method, total, callback=None, **pbar_kwargs
@@ -46,4 +79,4 @@ def tdqm_or_callback_wrapped(
     return Tqdm.wrapattr(fobj, method, total=total, bytes=True, **pbar_kwargs)
 
 
-DEFAULT_CALLBACK = fsspec.callbacks.NoOpCallback()
+DEFAULT_CALLBACK = NoOpCallback()

dvc/fs/base.py

@@ -15,18 +15,19 @@ from typing import (
     List,
     Optional,
     Union,
+    cast,
     overload,
 )
 
 from funcy import cached_property
-from tqdm.utils import CallbackIOWrapper
 
 from dvc.exceptions import DvcException
-from dvc.fs._callback import DEFAULT_CALLBACK, FsspecCallback
 from dvc.ui import ui
 from dvc.utils import tmp_fname
 from dvc.utils.fs import makedirs, move
 
+from ._callback import DEFAULT_CALLBACK, FsspecCallback
+
 if TYPE_CHECKING:
     from fsspec.spec import AbstractFileSystem
     from typing_extensions import Literal
@@ -271,19 +272,29 @@ class FileSystem:
 
     def put_file(
         self,
-        from_file: AnyFSPath,
+        from_file: Union[AnyFSPath, IO],
         to_info: AnyFSPath,
-        callback: Any = DEFAULT_CALLBACK,
+        callback: FsspecCallback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> None:
-        self.fs.put_file(from_file, to_info, callback=callback, **kwargs)
+        size = kwargs.get("size")
+        if size:
+            callback.set_size(size)
+        if hasattr(from_file, "read"):
+            stream = callback.wrap_attr(cast("IO", from_file))
+            self.upload_fobj(stream, to_info, size=size)
+        else:
+            assert isinstance(from_file, str)
+            self.fs.put_file(
+                os.fspath(from_file), to_info, callback=callback, **kwargs
+            )
         self.fs.invalidate_cache(self.path.parent(to_info))
 
     def get_file(
         self,
         from_info: AnyFSPath,
         to_info: AnyFSPath,
-        callback: Any = DEFAULT_CALLBACK,
+        callback: FsspecCallback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> None:
         self.fs.get_file(from_info, to_info, callback=callback, **kwargs)
@@ -320,15 +331,11 @@ class FileSystem:
         no_progress_bar: bool = False,
         **pbar_args: Any,
     ):
-        is_file_obj = hasattr(from_info, "read")
-        method = "upload_fobj" if is_file_obj else "put_file"
-        if not hasattr(self, method):
-            raise RemoteActionNotImplemented(method, self.scheme)
-
-        if not is_file_obj:
+        if not hasattr(from_info, "read"):
             from .local import localfs
 
             desc = desc or localfs.path.name(from_info)
+            logger.debug("Uploading '%s' to '%s'", from_info, to_info)
 
         stack = contextlib.ExitStack()
         if not callback:
@@ -341,24 +348,12 @@ class FileSystem:
             )
             stack.enter_context(pbar)
             callback = pbar.as_callback()
-            if total:
-                callback.set_size(total)
 
         with stack:
-            if is_file_obj:
-                wrapped = CallbackIOWrapper(
-                    callback.relative_update, from_info
-                )
-                # `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)
-
-            assert isinstance(from_info, str)
-            logger.debug("Uploading '%s' to '%s'", from_info, to_info)
-            # pylint: disable=no-member
+            if total:
+                callback.set_size(total)
             return self.put_file(
-                os.fspath(from_info), to_info, callback=callback
+                from_info, to_info, callback=callback, size=total
             )
 
     def download(
@@ -406,7 +401,7 @@ class FileSystem:
         self,
         from_info: AnyFSPath,
         to_info: AnyFSPath,
-        callback=DEFAULT_CALLBACK,
+        callback: FsspecCallback = DEFAULT_CALLBACK,
         jobs: int = None,
         **kwargs,
     ):
@@ -423,7 +418,7 @@ class FileSystem:
         )
         callback.set_size(len(from_infos))
 
-        download_files = FsspecCallback.wrap_fn(callback, self._download_file)
+        download_files = callback.wrap_fn(self._download_file)
         max_workers = jobs or self.jobs
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
             futures = [

dvc/fs/http.py

@@ -1,12 +1,12 @@
 import threading
-from typing import Any
+from typing import IO, Union
 
 from funcy import cached_property, memoize, wrap_with
 
 from dvc import prompt
 from dvc.scheme import Schemes
 
-from ._callback import DEFAULT_CALLBACK
+from ._callback import DEFAULT_CALLBACK, FsspecCallback
 from .base import AnyFSPath, FileSystem
 
 
@@ -141,9 +141,9 @@ class HTTPFileSystem(FileSystem):
 
     def put_file(
         self,
-        from_file: AnyFSPath,
+        from_file: Union[AnyFSPath, IO],
         to_info: AnyFSPath,
-        callback: Any = DEFAULT_CALLBACK,
+        callback: FsspecCallback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> None:
         kwargs.setdefault("method", self.upload_method)

dvc/progress.py

@@ -136,9 +136,9 @@ class Tqdm(tqdm):
         return wrapped
 
     def as_callback(self):
-        from dvc.fs._callback import FsspecCallback
+        from dvc.fs._callback import TqdmCallback
 
-        return FsspecCallback(self)
+        return TqdmCallback(self)
 
     def close(self):
         self.postfix["info"] = ""

dvc/repo/experiments/executor/ssh.py

@@ -257,7 +257,7 @@ class SSHExecutor(BaseExecutor):
             with TemporaryFile(mode="w+", encoding="utf-8") as fobj:
                 json.dump(info.asdict(), fobj)
                 fobj.seek(0)
-                fs.upload_fobj(fobj, infofile)
+                fs.put_file(fobj, infofile)
             cmd = ["source ~/.profile"]
             script_path = cls._setup_script_path(info.dvc_dir)
             if fs.exists(posixpath.join(info.root_dir, script_path)):

tests/func/test_fs.py

@@ -3,10 +3,10 @@ import os
 from operator import itemgetter
 from os.path import join
 
-import fsspec
 import pytest
 
 from dvc.fs import get_cloud_fs
+from dvc.fs._callback import FsspecCallback
 from dvc.fs.local import LocalFileSystem
 from dvc.repo import Repo
 
@@ -201,7 +201,7 @@ def test_walk_dont_ignore_subrepos(tmp_dir, scm, dvc):
     assert set(get_dirs(next(scm_fs.walk(path)))) == {".dvc", "subdir"}
 
 
-def test_fs_getsize(dvc, cloud):
+def test_fs_size(dvc, cloud):
     cloud.gen({"data": {"foo": "foo"}, "baz": "baz baz"})
     cls, config, path = get_cloud_fs(dvc, **cloud.config)
     fs = cls(**config)
@@ -248,7 +248,7 @@ def test_upload_callback(tmp_dir, dvc, cloud):
     fs = cls(**config)
     expected_size = os.path.getsize(tmp_dir / "foo")
 
-    callback = fsspec.Callback()
+    callback = FsspecCallback()
     fs.upload(
         (tmp_dir / "foo").fs_path,
         (cloud / "foo").fs_path,
@@ -267,7 +267,7 @@ def test_download_callback(tmp_dir, dvc, cloud, local_cloud):
     fs.upload((tmp_dir / "to_upload").fs_path, (cloud / "foo").fs_path)
     expected_size = fs.size((cloud / "foo").fs_path)
 
-    callback = fsspec.Callback()
+    callback = FsspecCallback()
     fs.download_file(
         (cloud / "foo").fs_path,
         (tmp_dir / "foo").fs_path,
@@ -284,7 +284,7 @@ def test_download_dir_callback(tmp_dir, dvc, cloud):
     fs = cls(**config)
     cloud.gen({"dir": {"foo": "foo", "bar": "bar"}})
 
-    callback = fsspec.Callback()
+    callback = FsspecCallback()
     fs.download(
         (cloud / "dir").fs_path, (tmp_dir / "dir").fs_path, callback=callback
     )
@@ -303,7 +303,7 @@ def test_download_callbacks_on_dvc_git_fs(tmp_dir, dvc, scm, fs_type):
 
     fs = dvc.dvcfs if fs_type == "dvc" else GitFileSystem(scm=scm, rev="HEAD")
 
-    callback = fsspec.Callback()
+    callback = FsspecCallback()
     fs.download_file(
         "file",
         (tmp_dir / "file2").fs_path,
@@ -315,7 +315,7 @@ def test_download_callbacks_on_dvc_git_fs(tmp_dir, dvc, scm, fs_type):
     assert callback.size == size
     assert callback.value == size
 
-    callback = fsspec.Callback()
+    callback = FsspecCallback()
     fs.download(
         "dir",
         (tmp_dir / "dir2").fs_path,
@@ -333,7 +333,7 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm):
 
     fs = dvc.repo_fs
 
-    callback = fsspec.Callback()
+    callback = FsspecCallback()
     fs.download(
         "dir",
         (tmp_dir / "dir2").fs_path,
@@ -344,7 +344,7 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm):
     assert callback.size == 2
     assert callback.value == 2
 
-    callback = fsspec.Callback()
+    callback = FsspecCallback()
     fs.download(
         os.path.join("dir", "foo"),
         (tmp_dir / "foo").fs_path,
@@ -356,7 +356,7 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm):
     assert callback.size == size
     assert callback.value == size
 
-    callback = fsspec.Callback()
+    callback = FsspecCallback()
     fs.download(
         os.path.join("dir", "bar"),
         (tmp_dir / "bar").fs_path,