refactor upload_fileobj in fsspec and add tests

Author: skshetry

src/webdav4/client.py

@@ -4,7 +4,7 @@
 import threading
 from contextlib import contextmanager, suppress
 from http import HTTPStatus
-from io import TextIOWrapper, UnsupportedOperation
+from io import TextIOWrapper
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -662,8 +662,7 @@ def upload_fileobj(
         # if we are not successfull in that, we gracefully fallback
         # to the chunked encoding.
         if size is None:
-            with suppress(TypeError, AttributeError, UnsupportedOperation):
-                size = peek_filelike_length(file_obj)
+            size = peek_filelike_length(file_obj)
 
         headers = {"Content-Length": str(size)} if size is not None else None
         if not overwrite and self.exists(to_path):

src/webdav4/fsspec.py

@@ -22,6 +22,7 @@
     cast,
 )
 
+from fsspec import Callback
 from fsspec.spec import AbstractBufferedFile, AbstractFileSystem
 
 from .client import (
@@ -32,6 +33,7 @@
     ResourceConflict,
     ResourceNotFound,
 )
+from .fs_utils import peek_filelike_length
 from .stream import read_into
 
 if TYPE_CHECKING:
@@ -41,8 +43,7 @@
     from os import PathLike
     from typing import AnyStr
 
-    from fsspec import Callback
-
+    from .callback import CallbackFn
     from .types import AuthTypes, URLTypes
 
 
@@ -319,12 +320,41 @@ def sign(self, path: str, expiration: int = 100, **kwargs: Any) -> None:
 
     def pipe_file(self, path: str, value: bytes, **kwargs: Any) -> None:
         """Upload the contents to given file in the remote webdav server."""
-        path = self._strip_protocol(path)
         buff = io.BytesIO(value)
         kwargs.setdefault("overwrite", True)
-        # maybe it's not a bad idea to make a `self.open` for `mode="rb"`
-        # on top of `io.BytesIO`?
-        self.client.upload_fileobj(buff, path, **kwargs)
+        return self.upload_fileobj(buff, path, **kwargs)
+
+    def upload_fileobj(
+        self,
+        fobj: BinaryIO,
+        rpath: str,
+        callback: "Callback" = None,
+        overwrite: bool = True,
+        size: int = None,
+        **kwargs: Any,
+    ) -> None:
+        """Upload contents from the fileobj to the remote path."""
+        rpath = self._strip_protocol(rpath)
+        self.mkdirs(os.path.dirname(rpath), exist_ok=True)
+
+        if size is None:
+            size = peek_filelike_length(fobj)
+
+        callback = cast("Callback", Callback.as_callback(callback))
+        if size is not None:  # pragma: no cover
+            callback.set_size(size)
+        progress_callback = cast("CallbackFn", callback.relative_update)
+
+        return self.client.upload_fileobj(
+            fobj,
+            rpath,
+            overwrite=overwrite,
+            callback=progress_callback,
+            size=size,
+            **kwargs,
+        )
+
+    put_fileobj = upload_fileobj
 
     def put_file(
         self,
@@ -334,16 +364,19 @@ def put_file(
         **kwargs: Any,
     ) -> None:
         """Copy file to remote webdav server."""
-        rpath = self._strip_protocol(rpath)
         if os.path.isdir(lpath):
-            self.makedirs(rpath, exist_ok=True)
-        else:
-            if callback is not None:
-                callback.set_size(os.path.getsize(lpath))
-                kwargs.setdefault("callback", callback.relative_update)
-            self.mkdirs(os.path.dirname(rpath), exist_ok=True)
+            rpath = self._strip_protocol(rpath)
+            return self.makedirs(rpath, exist_ok=True)
+
+        with open(lpath, mode="rb") as fobj:
             kwargs.setdefault("overwrite", True)
-            self.client.upload_file(lpath, rpath, **kwargs)
+            kwargs.setdefault("size", None)
+            return self.upload_fileobj(
+                fobj,
+                rpath,
+                callback=callback,
+                **kwargs,
+            )
 
 
 class WebdavFile(AbstractBufferedFile):

tests/test_fs_utils.py

@@ -0,0 +1,15 @@
+"""Test fs utils."""
+from io import BytesIO
+
+from webdav4.fs_utils import peek_filelike_length
+
+from .test_callback import ReadWrapper
+
+
+def test_peek_filelike_length():
+    """Test peek_filelike length for the fileobj."""
+    fobj = BytesIO(b"Hello, World!")
+
+    assert peek_filelike_length(fobj) == 13
+    assert peek_filelike_length(ReadWrapper(fobj)) is None  # type: ignore
+    assert peek_filelike_length(object()) is None

tests/test_fsspec.py

@@ -9,6 +9,7 @@
 from webdav4.fsspec import WebdavFileSystem
 from webdav4.urls import URL, join_url
 
+from .test_callback import ReadWrapper
 from .utils import TmpDir, approx_datetime
 
 
@@ -828,15 +829,44 @@ def test_touch_not_truncate(storage_dir: TmpDir, fs: WebdavFileSystem):
         fs.touch("foo", truncate=False)
 
 
+@pytest.mark.parametrize("with_size", [True, False])
+@pytest.mark.parametrize("wrap_fobj", [True, False])
+def test_upload_fileobj(
+    storage_dir: TmpDir,
+    fs: WebdavFileSystem,
+    with_size: bool,
+    wrap_fobj: bool,
+):
+    """Test upload_fileobj.
+
+    If with_size, size hints are provided to the API.
+    If wrap_fobj, the file object is wrapped with the ReadWrapper
+       that makes the peek_fileobj_length util to not be able to
+       figure out the size of the given fileobj. For the most part,
+       it should still work in this case.
+    """
+    if wrap_fobj and not with_size:
+        pytest.skip("the test server does not work without content-length :(")
+
+    foo = storage_dir / "foo"
+    length = foo.write_bytes(b"foo")
+
+    with foo.open(mode="rb") as fobj:
+        if wrap_fobj:
+            fobj = ReadWrapper(fobj)  # type: ignore
+        size = length if with_size else None
+        fs.upload_fileobj(fobj, "data/foobar", size=size)
+
+    assert (storage_dir / "data").cat() == {"foobar": "foo"}
+
+
 def test_callbacks(storage_dir: TmpDir, fs: WebdavFileSystem):
     """Test fsspec callbacks."""
     src_file = storage_dir / "source"
     dest_file = "data/get_put_file/dest"
 
     data = b"test" * 4
-
-    with open(src_file, "wb") as stream:
-        stream.write(data)
+    size = src_file.write_bytes(data)
 
     class EventLogger(fsspec.Callback):
         """Log callback values."""
@@ -855,9 +885,9 @@ def relative_update(self, value: int) -> None:
 
     event_logger = EventLogger()
     fs.put_file(src_file, dest_file, chunk_size=4, callback=event_logger)
-    assert fs.exists(dest_file)
+    assert (storage_dir / dest_file).cat() == data.decode()
 
-    assert event_logger.events[0] == ("set_size", len(data))
+    assert event_logger.events[0] == ("set_size", size)
     assert event_logger.events[1:] == [
         ("relative_update", 4),
         ("relative_update", 4),