feat: add update existing file function

Author: silentworks

storage3/_async/file_api.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass, field
 from io import BufferedReader, FileIO
 from pathlib import Path
-from typing import Any, Optional, Union, cast
+from typing import Any, Literal, Optional, Union, cast
 
 from httpx import HTTPError, Response
 
@@ -344,8 +344,9 @@ async def download(self, path: str, options: DownloadOptions = {}) -> bytes:
         )
         return response.content
 
-    async def upload(
+    async def _upload_or_update(
         self,
+        method: Literal["POST", "PUT"],
         path: str,
         file: Union[BufferedReader, bytes, FileIO, str, Path],
         file_options: Optional[FileOptions] = None,
@@ -367,9 +368,6 @@ async def upload(
             file_options = {}
         cache_control = file_options.get("cache-control")
         _data = {}
-        if cache_control:
-            file_options["cache-control"] = f"max-age={cache_control}"
-            _data = {"cacheControl": cache_control}
 
         headers = {
             **self._client.headers,
@@ -378,6 +376,10 @@ async def upload(
         }
         filename = path.rsplit("/", maxsplit=1)[-1]
 
+        if cache_control:
+            headers["cache-control"] = f"max-age={cache_control}"
+            _data = {"cacheControl": cache_control}
+
         if (
             isinstance(file, BufferedReader)
             or isinstance(file, bytes)
@@ -398,9 +400,38 @@ async def upload(
         _path = self._get_final_path(path)
 
         return await self._request(
-            "POST", f"/object/{_path}", files=files, headers=headers, data=_data
+            method, f"/object/{_path}", files=files, headers=headers, data=_data
         )
 
+    async def upload(
+        self,
+        path: str,
+        file: Union[BufferedReader, bytes, FileIO, str, Path],
+        file_options: Optional[FileOptions] = None,
+    ) -> Response:
+        """
+        Uploads a file to an existing bucket.
+
+        Parameters
+        ----------
+        path
+            The relative file path including the bucket ID. Should be of the format `bucket/folder/subfolder/filename.png`.
+            The bucket must already exist before attempting to upload.
+        file
+            The File object to be stored in the bucket. or a async generator of chunks
+        file_options
+            HTTP headers.
+        """
+        return await self._upload_or_update("POST", path, file, file_options)
+
+    async def update(
+        self,
+        path: str,
+        file: Union[BufferedReader, bytes, FileIO, str, Path],
+        file_options: Optional[FileOptions] = None,
+    ) -> Response:
+        return await self._upload_or_update("PUT", path, file, file_options)
+
     def _get_final_path(self, path: str) -> str:
         return f"{self.id}/{path}"
 

tests/_async/test_client.py

@@ -156,6 +156,67 @@ def file(tmp_path: Path, uuid_factory: Callable[[], str]) -> FileForTesting:
     )
 
 
+@pytest.fixture
+def two_files(tmp_path: Path, uuid_factory: Callable[[], str]) -> list[FileForTesting]:
+    """Creates multiple test files (different content, same bucket/folder path, different file names)"""
+    file_name_1 = "test_image_1.svg"
+    file_name_2 = "test_image_2.svg"
+    file_content = (
+        b'<svg width="109" height="113" viewBox="0 0 109 113" fill="none" xmlns="http://www.w3.org/2000/svg"> '
+        b'<path d="M63.7076 110.284C60.8481 113.885 55.0502 111.912 54.9813 107.314L53.9738 40.0627L99.1935 '
+        b'40.0627C107.384 40.0627 111.952 49.5228 106.859 55.9374L63.7076 110.284Z" fill="url(#paint0_linear)"/> '
+        b'<path d="M63.7076 110.284C60.8481 113.885 55.0502 111.912 54.9813 107.314L53.9738 40.0627L99.1935 '
+        b'40.0627C107.384 40.0627 111.952 49.5228 106.859 55.9374L63.7076 110.284Z" fill="url(#paint1_linear)" '
+        b'fill-opacity="0.2"/> <path d="M45.317 2.07103C48.1765 -1.53037 53.9745 0.442937 54.0434 5.041L54.4849 '
+        b'72.2922H9.83113C1.64038 72.2922 -2.92775 62.8321 2.1655 56.4175L45.317 2.07103Z" fill="#3ECF8E"/> <defs>'
+        b'<linearGradient id="paint0_linear" x1="53.9738" y1="54.974" x2="94.1635" y2="71.8295"'
+        b'gradientUnits="userSpaceOnUse"> <stop stop-color="#249361"/> <stop offset="1" stop-color="#3ECF8E"/> '
+        b'</linearGradient> <linearGradient id="paint1_linear" x1="36.1558" y1="30.578" x2="54.4844" y2="65.0806" '
+        b'gradientUnits="userSpaceOnUse"> <stop/> <stop offset="1" stop-opacity="0"/> </linearGradient> </defs> </svg>'
+    )
+    file_content_2 = (
+        b'<svg width="119" height="123" viewBox="0 0 119 123" fill="none" xmlns="http://www.w3.org/2000/svg"> '
+        b'<path d="M63.7076 110.284C60.8481 113.885 55.0502 111.912 54.9813 107.314L53.9738 40.0627L99.1935 '
+        b'40.0627C107.384 40.0627 111.952 49.5228 106.859 55.9374L63.7076 110.284Z" fill="url(#paint0_linear)"/> '
+        b'<path d="M63.7076 110.284C60.8481 113.885 55.0502 111.912 54.9813 107.314L53.9738 40.0627L99.1935 '
+        b'40.0627C107.384 40.0627 111.952 49.5228 106.859 55.9374L63.7076 110.284Z" fill="url(#paint1_linear)" '
+        b'fill-opacity="0.2"/> <path d="M45.317 2.07103C48.1765 -1.53037 53.9745 0.442937 54.0434 5.041L54.4849 '
+        b'72.2922H9.83113C1.64038 72.2922 -2.92775 62.8321 2.1655 56.4175L45.317 2.07103Z" fill="#3FDF8E"/> <defs>'
+        b'<linearGradient id="paint0_linear" x1="53.9738" y1="54.974" x2="94.1635" y2="71.8295"'
+        b'gradientUnits="userSpaceOnUse"> <stop stop-color="#249361"/> <stop offset="1" stop-color="#3FDF8E"/> '
+        b'</linearGradient> <linearGradient id="paint1_linear" x1="36.1558" y1="30.578" x2="54.4844" y2="65.0806" '
+        b'gradientUnits="userSpaceOnUse"> <stop/> <stop offset="1" stop-opacity="0"/> </linearGradient> </defs> </svg>'
+    )
+    bucket_folder = uuid_factory()
+    bucket_path_1 = f"{bucket_folder}/{file_name_1}"
+    bucket_path_2 = f"{bucket_folder}/{file_name_2}"
+    file_path_1 = tmp_path / file_name_1
+    file_path_2 = tmp_path / file_name_2
+    with open(file_path_1, "wb") as f:
+        f.write(file_content)
+    with open(file_path_2, "wb") as f:
+        f.write(file_content_2)
+
+    return [
+        FileForTesting(
+            name=file_name_1,
+            local_path=str(file_path_1),
+            bucket_folder=bucket_folder,
+            bucket_path=bucket_path_1,
+            mime_type="image/svg+xml",
+            file_content=file_content,
+        ),
+        FileForTesting(
+            name=file_name_2,
+            local_path=str(file_path_2),
+            bucket_folder=bucket_folder,
+            bucket_path=bucket_path_2,
+            mime_type="image/svg+xml",
+            file_content=file_content_2,
+        ),
+    ]
+
+
 @pytest.fixture
 def multi_file(tmp_path: Path, uuid_factory: Callable[[], str]) -> list[FileForTesting]:
     """Creates multiple test files (same content, same bucket/folder path, different file names)"""
@@ -223,6 +284,33 @@ async def test_client_upload(
     assert image_info.get("metadata", {}).get("mimetype") == file.mime_type
 
 
+async def test_client_update(
+    storage_file_client: AsyncBucketProxy,
+    two_files: list[FileForTesting],
+) -> None:
+    """Ensure we can upload files to a bucket"""
+    await storage_file_client.upload(
+        two_files[0].bucket_path,
+        two_files[0].local_path,
+        {"content-type": two_files[0].mime_type},
+    )
+
+    await storage_file_client.update(
+        two_files[0].bucket_path,
+        two_files[1].local_path,
+        {"content-type": two_files[1].mime_type},
+    )
+
+    image = await storage_file_client.download(two_files[0].bucket_path)
+    file_list = await storage_file_client.list(two_files[0].bucket_folder)
+    image_info = next(
+        (f for f in file_list if f.get("name") == two_files[0].name), None
+    )
+
+    assert image == two_files[1].file_content
+    assert image_info.get("metadata", {}).get("mimetype") == two_files[1].mime_type
+
+
 @pytest.mark.parametrize(
     "path", ["foobar.txt", "example/nested.jpg", "/leading/slash.png"]
 )