fix(toolkit): fix repository signatures (#422)

fal-ai · Feb 18, 2025 · 25ea9a9 · 25ea9a9
1 parent 445cc21
commit 25ea9a9
Show file tree

Hide file tree

Showing 7 changed files with 61 additions and 6 deletions.
diff --git a/projects/fal/src/fal/toolkit/file/file.py b/projects/fal/src/fal/toolkit/file/file.py
@@ -166,7 +166,9 @@ def from_bytes(
             fallback_repo = get_builtin_repository(fallback_repository)
 
             url = fallback_repo.save(
-                fdata, object_lifecycle_preference, **fallback_save_kwargs
+                fdata,
+                object_lifecycle_preference=object_lifecycle_preference,
+                **fallback_save_kwargs,
             )
 
         return cls(

diff --git a/projects/fal/src/fal/toolkit/file/providers/fal.py b/projects/fal/src/fal/toolkit/file/providers/fal.py
@@ -185,7 +185,13 @@ def _save(self, file: FileData, storage_type: str) -> str:
 @dataclass
 class FalFileRepository(FalFileRepositoryBase):
     def save(
-        self, file: FileData, object_lifecycle_preference: dict[str, str] | None = None
+        self,
+        file: FileData,
+        multipart: bool | None = None,
+        multipart_threshold: int | None = None,
+        multipart_chunk_size: int | None = None,
+        multipart_max_concurrency: int | None = None,
+        object_lifecycle_preference: dict[str, str] | None = None,
     ) -> str:
         return self._save(file, "gcs")
 
@@ -834,7 +840,10 @@ def save_file(
                     content_type=content_type,
                     file_name=os.path.basename(file_path),
                 )
-            url = self.save(data, object_lifecycle_preference)
+            url = self.save(
+                data,
+                object_lifecycle_preference=object_lifecycle_preference,
+            )
 
         return url, data
 
@@ -844,6 +853,10 @@ class InMemoryRepository(FileRepository):
     def save(
         self,
         file: FileData,
+        multipart: bool | None = None,
+        multipart_threshold: int | None = None,
+        multipart_chunk_size: int | None = None,
+        multipart_max_concurrency: int | None = None,
         object_lifecycle_preference: dict[str, str] | None = None,
     ) -> str:
         return f'data:{file.content_type};base64,{b64encode(file.data).decode("utf-8")}'
@@ -865,6 +878,10 @@ def _object_lifecycle_headers(
     def save(
         self,
         file: FileData,
+        multipart: bool | None = None,
+        multipart_threshold: int | None = None,
+        multipart_chunk_size: int | None = None,
+        multipart_max_concurrency: int | None = None,
         object_lifecycle_preference: dict[str, str] | None = None,
     ) -> str:
         headers = {
@@ -1013,7 +1030,10 @@ def save_file(
                     content_type=content_type,
                     file_name=os.path.basename(file_path),
                 )
-            url = self.save(data, object_lifecycle_preference)
+            url = self.save(
+                data,
+                object_lifecycle_preference=object_lifecycle_preference,
+            )
 
         return url, data
 
@@ -1119,6 +1139,9 @@ def save_file(
                     content_type=content_type,
                     file_name=os.path.basename(file_path),
                 )
-            url = self.save(data, object_lifecycle_preference)
+            url = self.save(
+                data,
+                object_lifecycle_preference=object_lifecycle_preference,
+            )
 
         return url, data
diff --git a/projects/fal/src/fal/toolkit/file/providers/gcp.py b/projects/fal/src/fal/toolkit/file/providers/gcp.py
@@ -56,6 +56,10 @@ def bucket(self):
     def save(
         self,
         data: FileData,
+        multipart: bool | None = None,
+        multipart_threshold: int | None = None,
+        multipart_chunk_size: int | None = None,
+        multipart_max_concurrency: int | None = None,
         object_lifecycle_preference: Optional[dict[str, str]] = None,
     ) -> str:
         destination_path = posixpath.join(

diff --git a/projects/fal/src/fal/toolkit/file/providers/r2.py b/projects/fal/src/fal/toolkit/file/providers/r2.py
@@ -73,6 +73,10 @@ def bucket(self):
     def save(
         self,
         data: FileData,
+        multipart: bool | None = None,
+        multipart_threshold: int | None = None,
+        multipart_chunk_size: int | None = None,
+        multipart_max_concurrency: int | None = None,
         object_lifecycle_preference: Optional[dict[str, str]] = None,
     ) -> str:
         destination_path = posixpath.join(

diff --git a/projects/fal/src/fal/toolkit/file/providers/s3.py b/projects/fal/src/fal/toolkit/file/providers/s3.py
@@ -57,6 +57,10 @@ def storage_client(self):
     def save(
         self,
         data: FileData,
+        multipart: bool | None = None,
+        multipart_threshold: int | None = None,
+        multipart_chunk_size: int | None = None,
+        multipart_max_concurrency: int | None = None,
         object_lifecycle_preference: Optional[dict[str, str]] = None,
         key: Optional[str] = None,
     ) -> str:

diff --git a/projects/fal/src/fal/toolkit/file/types.py b/projects/fal/src/fal/toolkit/file/types.py
@@ -39,6 +39,10 @@ class FileRepository:
     def save(
         self,
         data: FileData,
+        multipart: bool | None = None,
+        multipart_threshold: int | None = None,
+        multipart_chunk_size: int | None = None,
+        multipart_max_concurrency: int | None = None,
         object_lifecycle_preference: Optional[dict[str, str]] = None,
     ) -> str:
         raise NotImplementedError()
@@ -59,4 +63,11 @@ def save_file(
         with open(file_path, "rb") as fobj:
             data = FileData(fobj.read(), content_type, Path(file_path).name)
 
-        return self.save(data, object_lifecycle_preference), data
+        return self.save(
+            data,
+            multipart=multipart,
+            multipart_threshold=multipart_threshold,
+            multipart_chunk_size=multipart_chunk_size,
+            multipart_max_concurrency=multipart_max_concurrency,
+            object_lifecycle_preference=object_lifecycle_preference,
+        ), data
diff --git a/projects/fal/tests/test_stability.py b/projects/fal/tests/test_stability.py
@@ -603,6 +603,13 @@ def hello_file():
         return File.from_bytes(
             b"Hello fal storage from isolated",
             repository=repo_type,
+            save_kwargs={
+                "multipart": False,
+                "multipart_threshold": 1024 * 1024,
+                "multipart_chunk_size": 1024 * 1024,
+                "multipart_max_concurrency": 10,
+                "object_lifecycle_preference": {},
+            },
             fallback_repository=None,
         )