From 042457290a3d571360496601285db42c3b52b0c3 Mon Sep 17 00:00:00 2001
From: Erik Man <erikman@users.noreply.github.com>
Date: Sat, 13 Jul 2024 21:28:25 +0200
Subject: [PATCH] Add header parameter to upload (#176)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Saugat Pachhai (सौगात) <suagatchhetri@outlook.com>
---
 src/webdav4/client.py | 10 +++++++++-
 tests/test_client.py  |  2 +-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/webdav4/client.py b/src/webdav4/client.py
index af48d9a..bb430cd 100644
--- a/src/webdav4/client.py
+++ b/src/webdav4/client.py
@@ -625,6 +625,7 @@ def upload_file(  # noqa: PLR0913
         overwrite: bool = False,
         chunk_size: Optional[int] = None,
         callback: Optional[Callable[[int], Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
     ) -> None:
         """Upload file from local path to a given remote path."""
         with open(from_path, mode="rb") as fobj:
@@ -634,6 +635,7 @@ def upload_file(  # noqa: PLR0913
                 overwrite=overwrite,
                 chunk_size=chunk_size,
                 callback=callback,
+                headers=headers,
             )
 
     def upload_fileobj(  # noqa: PLR0913
@@ -644,8 +646,12 @@ def upload_fileobj(  # noqa: PLR0913
         callback: Optional[Callable[[int], Any]] = None,
         chunk_size: Optional[int] = None,
         size: Optional[int] = None,
+        headers: Optional[Dict[str, str]] = None,
     ) -> None:
         """Upload file from file object to given path."""
+        if headers is None:
+            headers = {}
+
         # we try to avoid chunked transfer as much as possible
         # so we try to use size as a hint if provided.
         # else, we will try to find that out from the file object
@@ -654,7 +660,9 @@ def upload_fileobj(  # noqa: PLR0913
         if size is None:
             size = peek_filelike_length(file_obj)
 
-        headers = {"Content-Length": str(size)} if size is not None else None
+        if size is not None:
+            headers = {"Content-Length": str(size), **headers}
+
         if not overwrite and self.exists(to_path):
             raise ResourceAlreadyExists(to_path)
 
diff --git a/tests/test_client.py b/tests/test_client.py
index 7315182..a2e227d 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -687,7 +687,7 @@ def test_upload_fobj_size_hints(client: Client):
 
         wrapped = ReadWrapper(BytesIO(b"foobar"))  # type: ignore
         client.upload_fileobj(wrapped, "foobar")  # type: ignore
-        assert m.call_args[1]["headers"] is None
+        assert m.call_args[1]["headers"] == {}
 
         wrapped = ReadWrapper(BytesIO(b"foobar"))  # type: ignore
         client.upload_fileobj(wrapped, "foobar", size=6)  # type: ignore